From b283215c28d56d24b3f72a120b8f34222699aeda Mon Sep 17 00:00:00 2001 From: Yueqian Lin <70319226+linyueqian@users.noreply.github.com> Date: Mon, 17 Aug 2026 21:33:13 -0700 Subject: [PATCH] [Model] Add MiniMax Music 3 text-to-music (#6186) Signed-off-by: Yueqian Lin --- docs/models/supported_models.md | 1 + recipes/MiniMaxAI/MiniMax-Music3.md | 173 +++ .../test_invalid_audio_speech.py | 65 ++ .../e2e/online_serving/test_minimax_music3.py | 112 ++ tests/engine/test_async_omni_engine_input.py | 103 ++ .../openai_api/test_tts_detection.py | 1 + .../test_minimax_music3_batch_payload.py | 132 +++ .../models/test_minimax_music3_frame_state.py | 214 ++++ .../models/test_minimax_music3_prompt.py | 203 ++++ .../models/test_minimax_music3_row_binding.py | 217 ++++ tests/model_executor/test_cfg_pairing.py | 389 +++++++ vllm_omni/config/model.py | 5 + vllm_omni/config/pipeline_registry.py | 2 + vllm_omni/data_entry_keys.py | 6 + vllm_omni/deploy/minimax_music3.yaml | 130 +++ vllm_omni/deploy/minimax_music3_2gpu.yaml | 97 ++ vllm_omni/engine/arg_utils.py | 2 + vllm_omni/engine/async_omni_engine.py | 86 +- vllm_omni/engine/stage_engine_core_proc.py | 9 +- vllm_omni/engine/stage_init_utils.py | 57 +- .../openai/tts_adapters/__init__.py | 1 + .../openai/tts_adapters/minimax_music3.py | 163 +++ vllm_omni/model_executor/models/audex/cfg.py | 391 +------ .../models/common/cfg_pairing.py | 460 ++++++++ .../models/minimax_music3/__init__.py | 34 + .../models/minimax_music3/acoustic.py | 459 ++++++++ .../models/minimax_music3/chunking.py | 96 ++ .../models/minimax_music3/constants.py | 111 ++ .../models/minimax_music3/dav.py | 294 ++++++ .../models/minimax_music3/depth_graph.py | 137 +++ .../models/minimax_music3/dit.py | 484 +++++++++ .../models/minimax_music3/frame_state.py | 158 +++ .../models/minimax_music3/pipeline.py | 103 ++ .../models/minimax_music3/prompt.py | 172 +++ .../models/minimax_music3/rvq_decoder.py | 179 ++++ .../models/minimax_music3/sampling.py | 223 ++++ .../models/minimax_music3/staging.py | 66 ++ .../models/minimax_music3/talker.py | 990 ++++++++++++++++++ .../models/minimax_music3/weights.py | 133 +++ vllm_omni/model_executor/models/registry.py | 11 + .../stage_input_processors/minimax_music3.py | 544 ++++++++++ .../transformers_utils/configs/__init__.py | 3 + .../configs/minimax_music3.py | 65 ++ 43 files changed, 6876 insertions(+), 405 deletions(-) create mode 100644 recipes/MiniMaxAI/MiniMax-Music3.md create mode 100644 tests/e2e/online_serving/test_minimax_music3.py create mode 100644 tests/model_executor/models/test_minimax_music3_batch_payload.py create mode 100644 tests/model_executor/models/test_minimax_music3_frame_state.py create mode 100644 tests/model_executor/models/test_minimax_music3_prompt.py create mode 100644 tests/model_executor/models/test_minimax_music3_row_binding.py create mode 100644 tests/model_executor/test_cfg_pairing.py create mode 100644 vllm_omni/deploy/minimax_music3.yaml create mode 100644 vllm_omni/deploy/minimax_music3_2gpu.yaml create mode 100644 vllm_omni/entrypoints/openai/tts_adapters/minimax_music3.py create mode 100644 vllm_omni/model_executor/models/common/cfg_pairing.py create mode 100644 vllm_omni/model_executor/models/minimax_music3/__init__.py create mode 100644 vllm_omni/model_executor/models/minimax_music3/acoustic.py create mode 100644 vllm_omni/model_executor/models/minimax_music3/chunking.py create mode 100644 vllm_omni/model_executor/models/minimax_music3/constants.py create mode 100644 vllm_omni/model_executor/models/minimax_music3/dav.py create mode 100644 vllm_omni/model_executor/models/minimax_music3/depth_graph.py create mode 100644 vllm_omni/model_executor/models/minimax_music3/dit.py create mode 100644 vllm_omni/model_executor/models/minimax_music3/frame_state.py create mode 100644 vllm_omni/model_executor/models/minimax_music3/pipeline.py create mode 100644 vllm_omni/model_executor/models/minimax_music3/prompt.py create mode 100644 vllm_omni/model_executor/models/minimax_music3/rvq_decoder.py create mode 100644 vllm_omni/model_executor/models/minimax_music3/sampling.py create mode 100644 vllm_omni/model_executor/models/minimax_music3/staging.py create mode 100644 vllm_omni/model_executor/models/minimax_music3/talker.py create mode 100644 vllm_omni/model_executor/models/minimax_music3/weights.py create mode 100644 vllm_omni/model_executor/stage_input_processors/minimax_music3.py create mode 100644 vllm_omni/transformers_utils/configs/minimax_music3.py diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index d9fdb09a4..31f3ac35d 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -78,6 +78,7 @@ th { | `SoulXSingerPipeline` | SoulX-Singer (SVS) | `Soul-AILab/SoulX-Singer` | ✅︎ | | | | — | | `SoulXSingerSVCPipeline` | SoulX-Singer-SVC | `Soul-AILab/SoulX-Singer` (`model-svc.pt`) | ✅︎ | | | | — | | `AudioXPipeline` | AudioX | `zhangj1an/AudioX` | ✅︎ | ✅︎ | | | — | +| `MiniMaxMusic3ForConditionalGeneration` | MiniMax Music 3 (text-to-music) | `MiniMaxAI/MiniMax-Music3` | ✅︎ | | | | — | | `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | ✅︎ | | | | [Published](https://recipes.vllm.ai/Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice) | | `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — | | `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-Base | `Qwen/Qwen3-TTS-12Hz-1.7B-Base` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — | diff --git a/recipes/MiniMaxAI/MiniMax-Music3.md b/recipes/MiniMaxAI/MiniMax-Music3.md new file mode 100644 index 000000000..3d078f645 --- /dev/null +++ b/recipes/MiniMaxAI/MiniMax-Music3.md @@ -0,0 +1,173 @@ +# MiniMax Music 3 + +> Text-to-music: lyrics plus a style caption in, a 32 kHz stereo song out + +## Summary + +- Vendor: MiniMax +- Model: `MiniMaxAI/MiniMax-Music3` +- Task: Text-to-music generation (vocals and instrumental together) +- Mode: Online serving through the OpenAI-compatible `/v1/audio/speech` API +- Maintainer: @linyueqian + +## When to use this recipe + +Use it to serve `MiniMaxAI/MiniMax-Music3` for full-song generation up to six +minutes. This is not a TTS model with a musical voice: there is no reference +speaker, no `voice`, and no `temperature`. What you control is the lyrics, the +caption, the seed, and the length. + +## References + +- Model card: [MiniMaxAI/MiniMax-Music3](https://huggingface.co/MiniMaxAI/MiniMax-Music3) +- Demo samples: [MiniMax Music 3 demo](https://minimax-ai.github.io/music3-demo/) +- Deploy config: `vllm_omni/deploy/minimax_music3.yaml` (single GPU), + `vllm_omni/deploy/minimax_music3_2gpu.yaml` (split across two) + +## Architecture + +| Component | Spec | +|---|---| +| Backbone | Qwen3 decoder, 36 layers, hidden 4096, GQA 32/8, vocab 200k | +| Depth decoder | 4 layers, hidden 4096, 16 heads; 7 residual codebooks of 1024 | +| Frame | 8 codebooks; `c0` in the backbone vocabulary, `c1..c7` in the depth decoder | +| AR guidance | Classifier-free, scale 1.5, `c0` masked to the conditioned branch's top 50 | +| Acoustic stage | Flow-matching transformer (36 layers, dim 2048) plus a DAC-style decoder, 512x upsampling | +| Solver | Euler, 30 steps, acoustic guidance scale 1.7 | +| Window | 200 frames per acoustic chunk, 100-frame hop | +| Frame rate | 25 frames per second | +| Context length | 10,240 tokens | +| Output | 32 kHz stereo WAV | + +Generation is two-stage. The backbone predicts one audio frame per decode step, +each frame eight RVQ codebooks deep. Every 200 frames are handed to the +flow-matching transformer, which solves a latent that the decoder turns into +waveform. + +**Guidance doubles the rows.** Every request decodes twice: one row sees the +real prompt, the other sees the same prompt with the caption-and-lyrics span +replaced by `<|audio_cfg|>`. Both rows hold their own KV cache for the whole +song, so a request costs twice the KV its length suggests. Size the pool for +rows, not requests. + +## Hardware Support + +## GPU + +### 1x H200 141GB + +#### Environment + +- OS: Linux +- Python: 3.12 +- CUDA: 12.9 +- vLLM version: 0.27.0 +- vLLM-Omni version or commit: branch `feat/minimax-music3` + +#### Command + +```bash +vllm serve MiniMaxAI/MiniMax-Music3 \ + --host 0.0.0.0 --port 8000 \ + --trust-remote-code --omni +``` + +The default deploy config `vllm_omni/deploy/minimax_music3.yaml` is discovered +automatically from the checkpoint's `model_type=minimax_music3`. It colocates +both stages on device 0. + +#### Verification + +```bash +curl -X POST http://localhost:8000/v1/audio/speech \ + -H "Content-Type: application/json" \ + -d '{ + "model": "MiniMaxAI/MiniMax-Music3", + "input": "[Verse]\nWalking down the empty street at midnight\nStreetlights flicker like a broken dream\n[Chorus]\nAnd I keep on walking\nTill the morning finds me", + "instructions": "A melancholic lo-fi hip-hop track at 85 BPM in F minor: mellow Rhodes piano riff, soft vinyl crackle, dusty boom-bap drums with a laid-back swing, warm upright bass. Intimate bedroom production, gentle tape saturation, no bright cymbals.", + "seed": 1, + "max_new_tokens": 750 + }' \ + --output song.wav +``` + +Expected: a 30 second 32 kHz stereo WAV. `max_new_tokens` counts audio frames +at 25 per second, so 750 frames is at most 30 seconds. It is a cap, not a +target: the model ends the song itself when it emits the audio-end token, and a +response shorter than the cap is the model finishing rather than a truncation. + +#### Notes + +- **Memory:** stage 0 takes `gpu_memory_utilization: 0.6` for the bf16 backbone + and its KV pool; stage 1 takes 0.2 for the transformer and decoder weights, + which sit outside any KV budget. +- **Key flags:** none required. Guidance, the solver step count and the window + geometry are fixed by the checkpoint and are not request parameters. +- **Known limitations:** + - The external API is non-streaming; `stream: true` is rejected. + - `max_model_len` is 10,240, so a maximum-length prompt (5,000 tokens) and a + maximum-length song (9,000 frames) cannot both fit. Long captions shorten + the maximum song. This is a property of the checkpoint. + - The acoustic stage runs in float32. bfloat16 is accepted but measurably + degrades the solver. + +### 2x H200 141GB + +Use `minimax_music3_2gpu.yaml` to put the acoustic stage on the second device: + +```bash +vllm serve MiniMaxAI/MiniMax-Music3 \ + --host 0.0.0.0 --port 8000 \ + --trust-remote-code --omni \ + --deploy-config minimax_music3_2gpu.yaml +``` + +Only the placement differs; both layouts run the acoustic stage in float32. + +## Writing the request + +Two fields carry the request, and both are required. + +- `input` is the **lyrics**. Structure tags (`[Verse]`, `[Chorus]`, `[Bridge]`, + `[Outro]`) steer the arrangement and are part of the prompt contract. + **Put a tag on its own line.** Normalization keeps only the tags on a line + that starts with one and drops the rest of that line, so lyrics written next + to a tag are silently lost: + + ```text + "[Verse]\nWalking down the street" -> [start] [verse] Walking down the street + "[Verse] Walking down the street" -> [start] [verse] + ``` + +- `instructions` is the **caption** describing genre, instrumentation, tempo, + mood and production. It is the strongest control available: vague captions + give generic arrangements. Long structured captions work well, since the + model was trained on descriptions covering global attributes, emotional + progression, vocal detail and arrangement. + +A request is deterministic in its seed: the same lyrics, caption, seed and +length return the same audio. Changing only the seed gives another take of the +same song. + +## Request parameters + +| Parameter | Type | Default | Description | +|---|---|---|---| +| `input` | string | required | Lyrics, non-empty. Tags on their own lines | +| `instructions` | string | required | Caption, non-empty | +| `seed` | int | `0` | Non-negative. Fixes the output for a given request | +| `max_new_tokens` | int | `9000` | Cap on audio frames, 25 per second. Maximum 9,000 | +| `response_format` | string | `"wav"` | Output container | + +## Parameters the model rejects + +The request is refused rather than silently ignored, so a mistake is visible +immediately. + +| Parameter | Reason | +|---|---| +| `temperature`, `top_p`, `top_k`, `repetition_penalty` | Sampling is fixed: guidance at 1.5, then a seeded top-k 50 draw | +| `voice` | There is no speaker to select; the vocal comes from the caption | +| `ref_audio`, `ref_text`, `language`, `task_type` | No reference-audio conditioning and no language tag | +| `speed` | Only `1.0`. Tempo belongs in the caption, e.g. "at 92 BPM" | +| `stream: true` | The external API is non-streaming | diff --git a/tests/dfx/reliability/invalid_param_test/test_invalid_audio_speech.py b/tests/dfx/reliability/invalid_param_test/test_invalid_audio_speech.py index 61cd30ec8..e30591c19 100644 --- a/tests/dfx/reliability/invalid_param_test/test_invalid_audio_speech.py +++ b/tests/dfx/reliability/invalid_param_test/test_invalid_audio_speech.py @@ -976,3 +976,68 @@ def test_voices_delete_invalid_requests( "err_message": err_message, } ) + + +# ─── POST /v1/audio/speech · MiniMax Music 3 ─── + +# Text-to-music on the speech endpoint. Most of the speech contract does not +# apply: there is no speaker, no reference audio and no temperature, and the +# lyrics alone are not a request. Each case below pins a rejection so a caller +# who assumes the usual TTS shape finds out immediately. +_MINIMAX_MUSIC3_SPEECH = [ + pytest.param( + OmniServerParams( + model="MiniMaxAI/MiniMax-Music3", + stage_config_path=get_deploy_config_path("minimax_music3.yaml"), + server_args=_SPEECH_SERVER_ARGS, + ), + id="minimax_music3", + marks=hardware_marks(res={"cuda": "H100"}), + ), +] + +_MINIMAX_MUSIC3_LYRICS = "[Verse]\nWalking down the empty street at midnight" +_MINIMAX_MUSIC3_CAPTION = "A melancholic lo-fi hip-hop track at 85 BPM in F minor." +# 25 frames per second, capped at six minutes of song. +_MINIMAX_MUSIC3_MAX_FRAMES = 9000 + + +@pytest.mark.parametrize( + "overrides, err_message", + [ + pytest.param({"instructions": None}, ("instructions",), id="caption_missing"), + pytest.param({"instructions": " "}, ("instructions",), id="caption_blank"), + pytest.param({"input": " "}, ("input",), id="lyrics_blank"), + pytest.param({"voice": "alloy"}, ("voice",), id="voice_rejected"), + pytest.param({"speed": 1.5}, ("speed",), id="speed_rejected"), + pytest.param( + {"max_new_tokens": _MINIMAX_MUSIC3_MAX_FRAMES + 1}, + ("max_new_tokens", str(_MINIMAX_MUSIC3_MAX_FRAMES)), + id="max_new_tokens_over_cap", + ), + pytest.param({"max_new_tokens": 0}, ("max_new_tokens",), id="max_new_tokens_zero"), + ], +) +@pytest.mark.parametrize("omni_server", _MINIMAX_MUSIC3_SPEECH, indirect=True) +def test_minimax_music3_speech_invalid_field_values( + omni_server: OmniServer, + openai_client: OpenAIClientHandler, + overrides: dict[str, object], + err_message: str | tuple[str, ...], +) -> None: + body: dict[str, Any] = { + "model": omni_server.model, + "input": _MINIMAX_MUSIC3_LYRICS, + "instructions": _MINIMAX_MUSIC3_CAPTION, + "seed": 1, + "max_new_tokens": 250, + "response_format": "wav", + } + for key, value in overrides.items(): + if value is None: + body.pop(key, None) + else: + body[key] = value + openai_client.send_audio_speech_http_request( + {"json": body, "timeout": 120, "err_code": 400, "err_message": err_message} + ) diff --git a/tests/e2e/online_serving/test_minimax_music3.py b/tests/e2e/online_serving/test_minimax_music3.py new file mode 100644 index 000000000..3d8718e70 --- /dev/null +++ b/tests/e2e/online_serving/test_minimax_music3.py @@ -0,0 +1,112 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""E2E online tests for MiniMax Music 3 text-to-music. + +This is a music model on the speech endpoint, so the request shape differs +from every other TTS model here: ``input`` carries the lyrics, ``instructions`` +carries the caption that decides genre and arrangement, and both are required. + +Weekly rather than per-PR: the checkpoint is 27 GB across five components and a +single request occupies two decode rows, because classifier-free guidance is +mandatory for this model. Rejection cases live in +``tests/dfx/reliability/invalid_param_test/test_invalid_audio_speech.py``. +""" + +import os + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + +import pytest + +from tests.helpers.mark import hardware_test +from tests.helpers.runtime import OmniServerParams +from tests.helpers.stage_config import get_deploy_config_path + +MODEL = "MiniMaxAI/MiniMax-Music3" +DEFAULT_AUDIO_SPEECH_TIMEOUT_S = 600.0 + +# 25 frames per second, so 250 frames is ten seconds and 750 is thirty. The +# byte floors sit below the true payload (32 kHz stereo, 16-bit) to leave room +# for the model ending a song before the cap. +_FRAMES_10S = 250 +_FRAMES_30S = 750 +_MIN_BYTES_10S = 500_000 +_MIN_BYTES_30S = 1_500_000 + +LYRICS = ( + "[Verse]\nWalking down the empty street at midnight\n" + "Streetlights flicker like a broken dream\n" + "[Chorus]\nAnd I keep on walking\nTill the morning finds me" +) + +CAPTION = ( + "A melancholic lo-fi hip-hop track at 85 BPM in F minor: mellow Rhodes piano riff, " + "soft vinyl crackle, dusty boom-bap drums with a laid-back swing, warm upright bass. " + "Intimate bedroom production, gentle tape saturation, no bright cymbals." +) + + +tts_server_params = [ + pytest.param( + OmniServerParams( + model=MODEL, + stage_config_path=get_deploy_config_path("minimax_music3.yaml"), + server_args=["--trust-remote-code", "--disable-log-stats"], + ), + id="minimax_music3", + ) +] + + +@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_text_to_music_001(omni_server, openai_client) -> None: + """Baseline: lyrics plus caption in, one ten-second song out. + + Deploy Setting: minimax_music3.yaml + Input Modal: text (lyrics) + text (caption) + Output Modal: audio, 32 kHz stereo + Input Setting: stream=False, seed pinned, 250 frames + """ + request_config = { + "model": omni_server.model, + "input": LYRICS, + "instructions": CAPTION, + "seed": 1, + "max_new_tokens": _FRAMES_10S, + "stream": False, + "response_format": "wav", + "timeout": DEFAULT_AUDIO_SPEECH_TIMEOUT_S, + "min_audio_bytes": _MIN_BYTES_10S, + } + openai_client.send_audio_speech_request(request_config) + + +@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_text_to_music_multi_window_002(omni_server, openai_client) -> None: + """A song long enough to span several acoustic windows. + + The acoustic stage decodes overlapping 200-frame windows on a 100-frame + hop and crops the overlap back off, so a request that fits in one window + never exercises the stitching. 750 frames covers seven of them. + + Deploy Setting: minimax_music3.yaml + Input Setting: stream=False, seed pinned, 750 frames + """ + request_config = { + "model": omni_server.model, + "input": LYRICS, + "instructions": CAPTION, + "seed": 1, + "max_new_tokens": _FRAMES_30S, + "stream": False, + "response_format": "wav", + "timeout": DEFAULT_AUDIO_SPEECH_TIMEOUT_S, + "min_audio_bytes": _MIN_BYTES_30S, + } + openai_client.send_audio_speech_request(request_config) diff --git a/tests/engine/test_async_omni_engine_input.py b/tests/engine/test_async_omni_engine_input.py index 740064173..ee6d322ca 100644 --- a/tests/engine/test_async_omni_engine_input.py +++ b/tests/engine/test_async_omni_engine_input.py @@ -371,3 +371,106 @@ def test_build_add_request_message_releases_preselected_replica_on_preprocess_er ) assert stage_pool.get_bound_replica_id("req-error") is None + + +def test_cfg_companion_build_failure_admits_nothing(mocker: MockerFixture): + """A guided request is all-or-nothing. + + Admitting the parent before its companion is built leaves an orphan: a + model whose guidance is mandatory never completes the pair, so the request + holds scheduler and KV capacity for the scheduler's whole hold budget and + then produces no audio. The build has to raise before anything is enqueued. + """ + engine = object.__new__(AsyncOmniEngine) + params = SamplingParams(max_tokens=8) + engine.prompt_expand_func = mocker.Mock(side_effect=ValueError("cannot build the null twin")) + engine.supported_tasks = ("generate",) + engine.input_processor = mocker.Mock() + engine.request_queue = mocker.Mock() + + with pytest.raises(ValueError, match="null twin"): + engine._build_cfg_companions( + parent_id="req", + original_prompt={"prompt": "positive"}, + stage0_params=params, + sampling_params_list=[params], + ) + + engine.request_queue.sync_q.put.assert_not_called() + + +def test_cfg_companion_processing_failure_admits_nothing(mocker: MockerFixture): + """The same holds when expansion succeeds but input processing does not.""" + engine = object.__new__(AsyncOmniEngine) + params = SamplingParams(max_tokens=8) + engine.prompt_expand_func = lambda *_args: [ + ExpandedPrompt( + prompt={"prompt": "negative"}, + role="cfg_text", + request_id_suffix="__cfg_text", + ) + ] + engine.supported_tasks = ("generate",) + engine.input_processor = mocker.Mock() + engine.input_processor.process_inputs.side_effect = RuntimeError("tokenizer rejected the prompt") + engine.request_queue = mocker.Mock() + + with pytest.raises(RuntimeError, match="tokenizer rejected"): + engine._build_cfg_companions( + parent_id="req", + original_prompt={"prompt": "positive"}, + stage0_params=params, + sampling_params_list=[params], + ) + + engine.request_queue.sync_q.put.assert_not_called() + + +def test_cfg_companion_build_returns_messages_without_enqueueing(mocker: MockerFixture): + """Construction and admission are separate steps.""" + engine = object.__new__(AsyncOmniEngine) + params = SamplingParams(max_tokens=8) + engine.prompt_expand_func = lambda *_args: [ + ExpandedPrompt( + prompt={"prompt": "negative"}, + role="cfg_text", + request_id_suffix="__cfg_text", + ) + ] + engine.supported_tasks = ("generate",) + engine.input_processor = mocker.Mock() + engine.input_processor.process_inputs.return_value = _make_engine_core_request("req__cfg_text") + engine.request_queue = mocker.Mock() + + companions = engine._build_cfg_companions( + parent_id="req", + original_prompt={"prompt": "positive"}, + stage0_params=params, + sampling_params_list=[params], + ) + + assert len(companions) == 1 + assert companions[0].companion_id == "req__cfg_text" + assert companions[0].parent_id == "req" + engine.request_queue.sync_q.put.assert_not_called() + + +def test_cfg_expansion_returning_nothing_is_not_an_error(mocker: MockerFixture): + """An unguided request expands to no companions and admits normally.""" + engine = object.__new__(AsyncOmniEngine) + params = SamplingParams(max_tokens=8) + engine.prompt_expand_func = lambda *_args: [] + engine.supported_tasks = ("generate",) + engine.input_processor = mocker.Mock() + engine.request_queue = mocker.Mock() + + assert ( + engine._build_cfg_companions( + parent_id="req", + original_prompt={"prompt": "positive"}, + stage0_params=params, + sampling_params_list=[params], + ) + == [] + ) + engine.request_queue.sync_q.put.assert_not_called() diff --git a/tests/entrypoints/openai_api/test_tts_detection.py b/tests/entrypoints/openai_api/test_tts_detection.py index 2ac27bb18..6b8b97849 100644 --- a/tests/entrypoints/openai_api/test_tts_detection.py +++ b/tests/entrypoints/openai_api/test_tts_detection.py @@ -153,6 +153,7 @@ _PIPELINE_STAGES = [ "indextts2_talker", "latent_generator", "llm", + "minimax_music3_ar", "ming_tts", "moss_tts", "moss_tts_codec", diff --git a/tests/model_executor/models/test_minimax_music3_batch_payload.py b/tests/model_executor/models/test_minimax_music3_batch_payload.py new file mode 100644 index 000000000..531500cd4 --- /dev/null +++ b/tests/model_executor/models/test_minimax_music3_batch_payload.py @@ -0,0 +1,132 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Regression tests for MiniMax Music 3 batched payload partitioning. + +The stage-0 payload is split per request by indexing list values: a list is +indexed by batch row, anything else is handed to every request unchanged. So +per-request metadata has to be carried as lists. Merging it into one dict made +every request in a batch receive the last request's seed, window offset and +terminal flag, which crosses seeds between songs, misplaces acoustic windows, +and lets one request's terminal flag end another request's stream early. + +These drive ``make_omni_output`` directly with a hand-built batch, so no GPU +and no engine are needed. +""" + +from __future__ import annotations + +import pytest +import torch + +from vllm_omni.model_executor.models.minimax_music3.talker import ( + MiniMaxMusic3TalkerForConditionalGeneration as Talker, +) +from vllm_omni.utils.mm_outputs import to_payload_element + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_WIDTH = 4 + + +def _info(*, seed: int, start: int, frames: int, last: bool) -> dict: + """One batch row's intermediate buffer, as ``postprocess`` leaves it.""" + return { + "latent": torch.full((frames, _WIDTH), float(seed)), + "codes": {"audio": torch.full((frames, 8), seed)}, + "meta": { + "audio_seed": seed, + "num_processed_tokens": start, + "codec_chunk_frames": frames, + "last_chunk": last, + }, + } + + +def _payload_for(output, row: int) -> dict: + """Partition the batch payload down to one request, as the runner does.""" + flat = {} + for key, value in output.multimodal_outputs.items(): + if isinstance(value, dict): + for qual, val in value.items(): + flat[f"{key}.{qual}"] = val + else: + flat[key] = value + return {key: to_payload_element(value, row, 1, 1) for key, value in flat.items()} + + +class TestBatchedMetadataPartitioning: + """Two requests in one batch must keep their own metadata.""" + + @staticmethod + def _two_request_batch(): + infos = [ + _info(seed=111, start=100, frames=3, last=False), + _info(seed=222, start=200, frames=5, last=True), + ] + talker = Talker.__new__(Talker) + return talker.make_omni_output(torch.zeros(2, _WIDTH), model_intermediate_buffer=infos) + + def test_seeds_are_not_crossed(self) -> None: + """A song decoded with another request's seed is a different song.""" + output = self._two_request_batch() + assert _payload_for(output, 0)["meta.audio_seed"] == 111 + assert _payload_for(output, 1)["meta.audio_seed"] == 222 + + def test_window_offsets_stay_with_their_request(self) -> None: + """The start frame decides where a window lands in the song.""" + output = self._two_request_batch() + assert _payload_for(output, 0)["meta.num_processed_tokens"] == 100 + assert _payload_for(output, 1)["meta.num_processed_tokens"] == 200 + + def test_terminal_flag_does_not_leak_across_requests(self) -> None: + """One request finishing must not end another request's stream.""" + output = self._two_request_batch() + assert _payload_for(output, 0)["meta.last_chunk"] is False + assert _payload_for(output, 1)["meta.last_chunk"] is True + + def test_span_length_matches_its_own_latent(self) -> None: + output = self._two_request_batch() + for row, frames in ((0, 3), (1, 5)): + entry = _payload_for(output, row) + assert entry["meta.codec_chunk_frames"] == frames + assert entry["latent"].shape[0] == frames + + def test_latent_and_codes_follow_the_same_row(self) -> None: + """The tensors were already per-request; pin that they stay aligned.""" + output = self._two_request_batch() + for row, seed in ((0, 111), (1, 222)): + entry = _payload_for(output, row) + assert entry["latent"][0, 0].item() == pytest.approx(float(seed)) + assert int(entry["codes.audio"][0, 0]) == seed + + +class TestPartialBatches: + """Only some rows carry a span on any given step.""" + + def test_row_without_a_span_does_not_take_another_row_metadata(self) -> None: + infos = [{}, _info(seed=777, start=300, frames=2, last=True)] + talker = Talker.__new__(Talker) + output = talker.make_omni_output(torch.zeros(2, _WIDTH), model_intermediate_buffer=infos) + + idle = _payload_for(output, 0) + assert idle["latent"].numel() == 0 + assert idle["meta.audio_seed"] is None + assert idle["meta.last_chunk"] is None + + active = _payload_for(output, 1) + assert active["meta.audio_seed"] == 777 + assert active["meta.last_chunk"] is True + + def test_batch_with_no_spans_emits_no_payload(self) -> None: + talker = Talker.__new__(Talker) + output = talker.make_omni_output(torch.zeros(2, _WIDTH), model_intermediate_buffer=[{}, {}]) + assert output.multimodal_outputs is None + + def test_consumed_entries_are_cleared_from_the_buffer(self) -> None: + """The runner never clears this buffer, so a span left behind would be + re-emitted on every later step.""" + infos = [_info(seed=1, start=0, frames=2, last=False)] + talker = Talker.__new__(Talker) + talker.make_omni_output(torch.zeros(1, _WIDTH), model_intermediate_buffer=infos) + assert "latent" not in infos[0] + assert "meta" not in infos[0] diff --git a/tests/model_executor/models/test_minimax_music3_frame_state.py b/tests/model_executor/models/test_minimax_music3_frame_state.py new file mode 100644 index 000000000..73dc9215b --- /dev/null +++ b/tests/model_executor/models/test_minimax_music3_frame_state.py @@ -0,0 +1,214 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for MiniMax Music 3 per-request AR frame bookkeeping. + +The talker emits overlapping windows while dropping frames it can no longer +need, and flushes whatever is left as one span when the song ends. Getting the +eviction boundary wrong by one frame silently corrupts the conditioning handed +to the acoustic stage from the second window onward, which is audible but hard +to attribute, so the lifecycle is simulated here frame by frame. +""" + +from __future__ import annotations + +import pytest +import torch + +from vllm_omni.model_executor.models.minimax_music3.chunking import chunk_windows +from vllm_omni.model_executor.models.minimax_music3.constants import ( + AR_CHUNK_FRAMES, + AR_CHUNK_HOP_FRAMES, + NUM_CODEBOOKS, +) +from vllm_omni.model_executor.models.minimax_music3.frame_state import ( + FrameBuffer, + RequestARState, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +# One scalar per frame instead of 32768 dims: the bookkeeping is what is under +# test, and a frame's value doubles as its identity. +_WIDTH = 4 + + +def _frame(index: int) -> torch.Tensor: + return torch.full((1, _WIDTH), float(index)) + + +def _identities(chunk: torch.Tensor) -> list[int]: + return [int(row[0].item()) for row in chunk] + + +class TestFrameBuffer: + def test_slice_returns_the_requested_frames(self) -> None: + buffer = FrameBuffer() + for index in range(10): + buffer.append(_frame(index)) + assert _identities(buffer.slice(3, 7)) == [3, 4, 5, 6] + + def test_slice_is_relative_to_the_live_base(self) -> None: + buffer = FrameBuffer() + for index in range(10): + buffer.append(_frame(index)) + buffer.discard_before(4) + assert buffer.base_frame == 4 + assert _identities(buffer.slice(4, 8)) == [4, 5, 6, 7] + + def test_slicing_a_discarded_range_raises(self) -> None: + buffer = FrameBuffer() + for index in range(10): + buffer.append(_frame(index)) + buffer.discard_before(5) + with pytest.raises(RuntimeError, match="unavailable"): + buffer.slice(3, 6) + + def test_slicing_past_the_end_raises(self) -> None: + buffer = FrameBuffer() + for index in range(4): + buffer.append(_frame(index)) + with pytest.raises(RuntimeError, match="unavailable"): + buffer.slice(2, 9) + + def test_discard_is_clamped_at_both_ends(self) -> None: + buffer = FrameBuffer() + for index in range(5): + buffer.append(_frame(index)) + buffer.discard_before(-3) + assert buffer.base_frame == 0 + buffer.discard_before(99) + assert buffer.base_frame == 5 + assert buffer.end_frame == 5 + + def test_peak_tracks_the_high_water_mark(self) -> None: + buffer = FrameBuffer() + for index in range(10): + buffer.append(_frame(index)) + buffer.discard_before(8) + assert buffer.peak_frames == 10 + + +class TestWindowLifecycle: + """Drive the state exactly as the talker does, one frame per step.""" + + @staticmethod + def _run(total_frames: int) -> tuple[list[list[int]], int]: + """Generate ``total_frames`` frames, returning the emitted spans. + + Returns: + The frame identities of every emitted span, and the peak buffer + occupancy reached along the way. + """ + state = RequestARState(seed=1, max_frames=total_frames) + spans: list[list[int]] = [] + for index in range(total_frames): + state.frames.append(_frame(index)) + state.codes.append(torch.full((1, NUM_CODEBOOKS), index)) + state.generated_frames += 1 + window = state.ready_window() + if window is not None: + spans.append(_identities(state.take_window(window, is_final=False))) + + # Terminal flush: everything not yet emitted, as one contiguous span. + start = state.emitted_windows * AR_CHUNK_HOP_FRAMES + end = state.frames.end_frame + if end > start: + spans.append(_identities(state.frames.slice(start, end))) + return spans, state.frames.peak_frames + + def test_short_song_is_one_span(self) -> None: + spans, _ = self._run(50) + assert spans == [list(range(50))] + + def test_every_frame_is_delivered_exactly_once_in_order(self) -> None: + """Spans overlap by design, but their union must be the whole song. + + Reconstruct by taking each span's non-overlapping remainder, which is + what the acoustic stage does after cropping. + """ + for total in (250, 437, 600, 901): + spans, _ = self._run(total) + rebuilt: list[int] = [] + for span in spans: + rebuilt.extend(frame for frame in span if frame >= len(rebuilt)) + assert rebuilt == list(range(total)), f"gap or duplicate at {total} frames" + + def test_windows_advance_by_the_hop(self) -> None: + spans, _ = self._run(600) + for index, span in enumerate(spans[:-1]): + assert span[0] == index * AR_CHUNK_HOP_FRAMES + assert len(span) == AR_CHUNK_FRAMES + + def test_mid_stream_windows_are_never_the_final_frame(self) -> None: + """A window is held back until a further frame proves it is not last.""" + total = 600 + spans, _ = self._run(total) + for span in spans[:-1]: + assert span[-1] < total - 1 + + def test_terminal_span_covers_the_tail(self) -> None: + total = 437 + spans, _ = self._run(total) + assert spans[-1][-1] == total - 1 + + def test_terminal_span_is_bounded(self) -> None: + """The tail is at most one full window plus one hop.""" + for total in (250, 301, 437, 600, 901, 1000): + spans, _ = self._run(total) + assert len(spans[-1]) <= AR_CHUNK_FRAMES + AR_CHUNK_HOP_FRAMES + + def test_buffer_stays_bounded_over_a_long_song(self) -> None: + """Peak occupancy must not grow with song length; that is the point + of evicting retired frames.""" + _, peak_short = self._run(600) + _, peak_long = self._run(3000) + assert peak_long == peak_short + assert peak_long <= AR_CHUNK_FRAMES + AR_CHUNK_HOP_FRAMES + + def test_emitted_window_count_matches_the_window_plan(self) -> None: + """What the stream emitted must reconcile with chunk_windows().""" + for total in (250, 437, 600): + spans, _ = self._run(total) + planned = chunk_windows(total) + # Every window is accounted for: the pre-terminal spans are one + # window each, and the terminal span carries whatever remains. + assert len(spans) <= len(planned) + assert spans[0][0] == planned[0].start + + +class TestSamplingPositions: + def test_each_frame_reserves_one_slot_per_codebook(self) -> None: + """Codebook streams must never collide, or codes repeat across frames.""" + state = RequestARState(seed=1, max_frames=10) + positions = [state.advance_position() for _ in range(5)] + assert positions == [0, 8, 16, 24, 32] + + def test_draw_slots_are_globally_unique(self) -> None: + state = RequestARState(seed=1, max_frames=64) + slots: set[int] = set() + for _ in range(64): + base = state.advance_position() + for codebook in range(NUM_CODEBOOKS): + slot = base + codebook + assert slot not in slots + slots.add(slot) + assert len(slots) == 64 * NUM_CODEBOOKS + + +class TestCodeRetention: + def test_codes_span_matches_the_frame_span(self) -> None: + state = RequestARState(seed=1, max_frames=10) + for index in range(10): + state.codes.append(torch.full((1, NUM_CODEBOOKS), index)) + codes = state.take_codes(2, 6) + assert codes is not None + assert codes.shape == (4, NUM_CODEBOOKS) + assert [int(row[0]) for row in codes] == [2, 3, 4, 5] + + def test_out_of_range_span_is_dropped_not_raised(self) -> None: + """Codes are diagnostic; the audio path must not fail without them.""" + state = RequestARState(seed=1, max_frames=10) + for index in range(4): + state.codes.append(torch.full((1, NUM_CODEBOOKS), index)) + assert state.take_codes(2, 99) is None + assert state.take_codes(0, 0) is None diff --git a/tests/model_executor/models/test_minimax_music3_prompt.py b/tests/model_executor/models/test_minimax_music3_prompt.py new file mode 100644 index 000000000..8cd9e41b0 --- /dev/null +++ b/tests/model_executor/models/test_minimax_music3_prompt.py @@ -0,0 +1,203 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for MiniMax Music 3 prompt assembly and window arithmetic. + +The prompt transformations are not cosmetic: they decide token positions, and +token positions decide the audio. These tests pin the documented behaviours +that a well-meaning cleanup would otherwise "fix". +""" + +from __future__ import annotations + +import pytest + +from vllm_omni.model_executor.models.minimax_music3.chunking import ( + chunk_windows, + crop_sample_bounds, + overlap_mel_length, +) +from vllm_omni.model_executor.models.minimax_music3.constants import ( + AR_CHUNK_FRAMES, + AR_CHUNK_HOP_FRAMES, + BLEND_MEL_FRAMES, + DAV_HOP_SAMPLES, + HOP_MEL_FRAMES, +) +from vllm_omni.model_executor.models.minimax_music3.prompt import ( + SPECIAL_TOKEN_IDS, + build_cfg_null_token_ids, + build_prompt, + clean_caption, + normalize_lyrics, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class TestNormalizeLyrics: + def test_prepends_start_tag(self) -> None: + assert normalize_lyrics("[Verse]\nhello").startswith("[start]\n") + + def test_lowercases_structure_tags(self) -> None: + assert "[verse]" in normalize_lyrics("[Verse]\nhello") + assert "[Verse]" not in normalize_lyrics("[Verse]\nhello") + + def test_text_after_a_leading_tag_is_dropped(self) -> None: + """Documented contract: a tag must sit on its own line. + + Lyrics written next to a tag are silently lost, which is why the + serving docs tell callers to put tags on their own line. + """ + out = normalize_lyrics("[Verse] Walking down the street") + assert "Walking" not in out + assert "[verse]" in out + + def test_text_on_its_own_line_survives(self) -> None: + out = normalize_lyrics("[Verse]\nWalking down the street") + assert "Walking down the street" in out + + def test_bracket_space_becomes_newline(self) -> None: + assert normalize_lyrics("a [chorus]") == "[start]\na\n[chorus]" + + +class TestCleanCaption: + def test_special_token_form_is_rewritten_not_passed_through(self) -> None: + """A caption must not be able to inject a real special token.""" + assert clean_caption("<|bpm 92|>") == "bpm is 92" + assert "<|" not in clean_caption("<|bpm 92|>") + + def test_lone_tag_keeps_its_inner_text(self) -> None: + assert clean_caption("<|instrumental|>") == "instrumental" + + def test_markdown_is_stripped(self) -> None: + assert clean_caption("## Heading") == "Heading" + assert clean_caption("**bold**") == "bold" + assert clean_caption("- item") == "item" + + def test_blank_lines_collapse_at_two_not_three(self) -> None: + """The source collapses runs of two or more, not three or more. + + This shifts blank-line token positions and therefore the audio, so it + is pinned deliberately. + """ + assert clean_caption("a\n\nb") == "a\nb" + assert clean_caption("a\n\n\n\nb") == "a\nb" + + +class TestBuildPrompt: + def test_framing_tokens_are_present_and_ordered(self) -> None: + prompt = build_prompt("a bright pop song", "[Verse]\nhello") + assert prompt.startswith("<|im_start|><|caption_start|>") + assert prompt.endswith("<|im_end|><|audio_start|>") + assert prompt.index("<|caption_end|>") < prompt.index("<|lyrics_start|>") + + +class TestBuildCfgNullTokenIds: + @staticmethod + def _prompt_ids(body_len: int) -> list[int]: + return ( + [SPECIAL_TOKEN_IDS["<|im_start|>"]] + + list(range(1000, 1000 + body_len)) + + [SPECIAL_TOKEN_IDS["<|im_end|>"], SPECIAL_TOKEN_IDS["<|audio_start|>"]] + ) + + def test_length_is_preserved(self) -> None: + """Equal length is what lets the two rows decode in lockstep.""" + ids = self._prompt_ids(12) + assert len(build_cfg_null_token_ids(ids)) == len(ids) + + def test_body_is_replaced_and_framing_survives(self) -> None: + ids = self._prompt_ids(12) + null = build_cfg_null_token_ids(ids) + assert null[0] == SPECIAL_TOKEN_IDS["<|im_start|>"] + assert null[-2] == SPECIAL_TOKEN_IDS["<|im_end|>"] + assert null[-1] == SPECIAL_TOKEN_IDS["<|audio_start|>"] + assert set(null[1:-2]) == {SPECIAL_TOKEN_IDS["<|audio_cfg|>"]} + + def test_input_is_not_mutated(self) -> None: + ids = self._prompt_ids(6) + before = list(ids) + build_cfg_null_token_ids(ids) + assert ids == before + + def test_minimal_prompt_with_empty_body(self) -> None: + ids = [ + SPECIAL_TOKEN_IDS["<|im_start|>"], + SPECIAL_TOKEN_IDS["<|audio_cfg|>"], + SPECIAL_TOKEN_IDS["<|im_end|>"], + SPECIAL_TOKEN_IDS["<|audio_start|>"], + ] + assert build_cfg_null_token_ids(ids) == ids + + def test_too_short_is_rejected(self) -> None: + with pytest.raises(ValueError, match="too short"): + build_cfg_null_token_ids([1, 2, 3]) + + def test_wrong_framing_is_rejected(self) -> None: + with pytest.raises(ValueError, match="framing"): + build_cfg_null_token_ids([9, 9, 9, 9, 9]) + + +class TestChunkWindows: + def test_no_frames_yields_nothing(self) -> None: + assert chunk_windows(0) == [] + + def test_short_song_is_one_final_window(self) -> None: + windows = chunk_windows(50) + assert len(windows) == 1 + assert (windows[0].start, windows[0].end) == (0, 50) + assert windows[0].is_first and windows[0].is_last + + def test_exactly_one_window(self) -> None: + windows = chunk_windows(AR_CHUNK_FRAMES) + assert len(windows) == 1 + assert windows[0].is_last + + def test_windows_advance_by_the_hop(self) -> None: + windows = chunk_windows(450) + starts = [w.start for w in windows] + assert starts == list(range(0, starts[-1] + 1, AR_CHUNK_HOP_FRAMES)) + assert windows[-1].is_last + assert not any(w.is_last for w in windows[:-1]) + + def test_windows_cover_every_frame(self) -> None: + frames = 437 + windows = chunk_windows(frames) + covered: set[int] = set() + for window in windows: + covered.update(range(window.start, window.end)) + assert covered == set(range(frames)) + + def test_final_window_is_clamped_to_the_song(self) -> None: + windows = chunk_windows(437) + assert windows[-1].end == 437 + + def test_only_the_first_window_is_first(self) -> None: + windows = chunk_windows(600) + assert windows[0].is_first + assert not any(w.is_first for w in windows[1:]) + + +class TestCropBounds: + def test_single_window_song_is_not_cropped(self) -> None: + window = chunk_windows(50)[0] + assert crop_sample_bounds(window) == (0, 0) + + def test_first_window_keeps_its_head(self) -> None: + left, right = crop_sample_bounds(chunk_windows(600)[0]) + assert left == 0 + assert right == (HOP_MEL_FRAMES - BLEND_MEL_FRAMES) * DAV_HOP_SAMPLES + + def test_last_window_keeps_its_tail(self) -> None: + left, right = crop_sample_bounds(chunk_windows(600)[-1]) + assert left == BLEND_MEL_FRAMES * DAV_HOP_SAMPLES + assert right == 0 + + def test_middle_window_is_cropped_on_both_sides(self) -> None: + windows = chunk_windows(900) + middle = windows[len(windows) // 2] + left, right = crop_sample_bounds(middle) + assert left > 0 and right > 0 + + def test_overlap_is_half_a_window(self) -> None: + assert overlap_mel_length() == HOP_MEL_FRAMES // 2 diff --git a/tests/model_executor/models/test_minimax_music3_row_binding.py b/tests/model_executor/models/test_minimax_music3_row_binding.py new file mode 100644 index 000000000..94ff30199 --- /dev/null +++ b/tests/model_executor/models/test_minimax_music3_row_binding.py @@ -0,0 +1,217 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Regression tests for MiniMax Music 3 row-to-request binding. + +vLLM compacts and reorders the persistent batch whenever a request finishes or +is preempted, so a batch row number is only meaningful within the step that +produced it. Every failure pinned here was a real defect: state keyed by row +survived into a step where the row belonged to someone else, or state keyed by +one spelling of the request id was released under another and leaked. + +These exercise the pure helpers and the binding arithmetic directly, so they +need no GPU and no engine. +""" + +from __future__ import annotations + +import pytest + +from vllm_omni.model_executor.models.minimax_music3.constants import ( + DEFAULT_MAX_AUDIO_FRAMES, +) +from vllm_omni.model_executor.models.minimax_music3.talker import ( + _CFG_UNCOND_SUFFIX, + _first_int, + _pair_id_from_request, + _request_key, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class TestRequestKey: + def test_engine_list_form_is_unwrapped(self) -> None: + """The engine injects global_request_id as a single-element list. + + Keying state on the list's ``str()`` would produce "['req-1']", which + never matches the bare id that ``on_requests_finished`` delivers, so + every finished song would leak its frame buffer. + """ + assert _request_key(["req-1"]) == "req-1" + + def test_bare_string_is_unchanged(self) -> None: + assert _request_key("req-1") == "req-1" + + def test_list_and_bare_forms_agree(self) -> None: + """Allocation and release must land on the same key.""" + assert _request_key(["req-1"]) == _request_key("req-1") + + def test_empty_and_missing_are_stable(self) -> None: + assert _request_key(None) == "" + assert _request_key([]) == "" + + def test_tuple_form(self) -> None: + assert _request_key(("req-9",)) == "req-9" + + +class TestPairIdDerivation: + def test_conditioned_request_is_its_own_pair(self) -> None: + assert _pair_id_from_request("req-1") == "req-1" + + def test_companion_maps_to_its_partner(self) -> None: + assert _pair_id_from_request(f"req-1{_CFG_UNCOND_SUFFIX}") == "req-1" + + def test_pair_members_agree_on_the_pair_id(self) -> None: + cond = "req-42" + uncond = f"req-42{_CFG_UNCOND_SUFFIX}" + assert _pair_id_from_request(cond) == _pair_id_from_request(uncond) + + def test_a_request_ending_in_a_similar_string_is_not_confused(self) -> None: + assert _pair_id_from_request("req-cfg_uncond") == "req-cfg_uncond" + + +class TestFirstInt: + def test_serving_layer_list_wrapping(self) -> None: + """tts_params values arrive wrapped in a single-element list.""" + assert _first_int([750], DEFAULT_MAX_AUDIO_FRAMES) == 750 + + def test_bare_value(self) -> None: + assert _first_int(750, DEFAULT_MAX_AUDIO_FRAMES) == 750 + + def test_missing_falls_back(self) -> None: + assert _first_int(None, DEFAULT_MAX_AUDIO_FRAMES) == DEFAULT_MAX_AUDIO_FRAMES + assert _first_int([], DEFAULT_MAX_AUDIO_FRAMES) == DEFAULT_MAX_AUDIO_FRAMES + + def test_nonsense_falls_back_rather_than_raising(self) -> None: + assert _first_int("abc", DEFAULT_MAX_AUDIO_FRAMES) == DEFAULT_MAX_AUDIO_FRAMES + assert _first_int(0, DEFAULT_MAX_AUDIO_FRAMES) == DEFAULT_MAX_AUDIO_FRAMES + assert _first_int(-5, DEFAULT_MAX_AUDIO_FRAMES) == DEFAULT_MAX_AUDIO_FRAMES + + +def _feedback_plan( + request_ids: list[str], + computed: list[int], + scheduled: list[int], + prompt_tokens: dict[str, int], + has_codes: set[str], +) -> list[int]: + """Mirror of the talker's feedback row selection. + + Returns the flat token offsets whose embedding gets replaced by an audio + frame. Kept as a standalone function so the arithmetic can be checked + without constructing a model. + """ + offsets: list[int] = [] + running = 0 + for count in scheduled: + offsets.append(running) + running += count + + rows: list[int] = [] + for row, request_id in enumerate(request_ids): + if request_id not in has_codes: + continue + if scheduled[row] != 1 or computed[row] < prompt_tokens.get(request_id, 0): + continue + rows.append(offsets[row]) + return rows + + +class TestFeedbackRowSelection: + """The guard that decides which token embeddings become audio frames.""" + + def test_decode_only_step_replaces_every_row(self) -> None: + ids = ["a", "a__cfg_uncond", "b", "b__cfg_uncond"] + rows = _feedback_plan( + ids, + computed=[120, 120, 90, 90], + scheduled=[1, 1, 1, 1], + prompt_tokens={i: 50 for i in ids}, + has_codes=set(ids), + ) + assert rows == [0, 1, 2, 3] + + def test_prefill_row_is_never_overwritten(self) -> None: + """A prompt token must keep its real embedding. + + The earlier guard compared the flat token count against a row budget, + so a short prompt slipped through and had its first embeddings + replaced with a previous request's audio frame. + """ + rows = _feedback_plan( + ["a", "b"], + computed=[0, 300], + scheduled=[8, 1], + prompt_tokens={"a": 8, "b": 40}, + has_codes={"b"}, + ) + # Row 0 occupies flat tokens 0..7 and must be untouched; row 1's single + # decode token sits at flat offset 8. + assert rows == [8] + + def test_short_prefill_below_the_row_budget_is_still_safe(self) -> None: + """The exact shape that defeated the old token-count guard.""" + rows = _feedback_plan( + ["a"], + computed=[0], + scheduled=[10], + prompt_tokens={"a": 10}, + has_codes={"a"}, + ) + assert rows == [] + + def test_chunked_prefill_midway_is_not_fed_back(self) -> None: + rows = _feedback_plan( + ["a"], + computed=[2048], + scheduled=[2048], + prompt_tokens={"a": 5000}, + has_codes={"a"}, + ) + assert rows == [] + + def test_final_prefill_chunk_is_not_fed_back(self) -> None: + """The step that completes a prompt still carries prompt tokens.""" + rows = _feedback_plan( + ["a"], + computed=[4096], + scheduled=[904], + prompt_tokens={"a": 5000}, + has_codes={"a"}, + ) + assert rows == [] + + def test_mixed_batch_offsets_follow_the_flat_layout(self) -> None: + """Tokens are a flat concatenation, so offsets are a running sum.""" + ids = ["p", "d1", "d2"] + rows = _feedback_plan( + ids, + computed=[0, 500, 700], + scheduled=[64, 1, 1], + prompt_tokens={"p": 64, "d1": 40, "d2": 40}, + has_codes={"d1", "d2"}, + ) + assert rows == [64, 65] + + def test_request_without_codes_yet_is_skipped(self) -> None: + """The first decode step after prefill has no previous frame.""" + rows = _feedback_plan( + ["a"], + computed=[50], + scheduled=[1], + prompt_tokens={"a": 50}, + has_codes=set(), + ) + assert rows == [] + + def test_reordered_batch_binds_by_request_not_position(self) -> None: + """The same two requests in swapped rows get swapped offsets. + + This is the compaction case: holding codes in a row-indexed buffer + would have handed each request the other's frame. + """ + prompt_tokens = {"a": 50, "b": 50} + first = _feedback_plan(["a", "b"], [100, 200], [1, 1], prompt_tokens, {"a"}) + swapped = _feedback_plan(["b", "a"], [200, 100], [1, 1], prompt_tokens, {"a"}) + assert first == [0] + assert swapped == [1] diff --git a/tests/model_executor/test_cfg_pairing.py b/tests/model_executor/test_cfg_pairing.py new file mode 100644 index 000000000..c00313a2a --- /dev/null +++ b/tests/model_executor/test_cfg_pairing.py @@ -0,0 +1,389 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for the model-neutral CFG request-pairing scheduler patches. + +The patches keep a cond/uncond request pair step-locked inside one engine: +admitted together, kept at equal ``num_computed_tokens``, finished together. +They are installed onto a lightweight fake scheduler here; vLLM's real +``Scheduler`` needs a full engine (and a GPU) to construct. +""" + +from __future__ import annotations + +from collections import deque +from dataclasses import dataclass, field +from types import SimpleNamespace +from typing import Any + +import pytest + +from vllm_omni.model_executor.models.common import cfg_pairing +from vllm_omni.model_executor.models.common.cfg_pairing import ( + _equalize_cfg_pair_progress, + _patch_scheduler, + cfg_uncond_suffixes, + register_cfg_uncond_suffix, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +# ───────── fakes ───────── + + +class _FakeRequest: + """Minimal stand-in for ``vllm.v1.request.Request`` (identity equality).""" + + def __init__(self, request_id: str, extra_args: dict[str, Any] | None = None, step_tokens: int = 1) -> None: + self.request_id = request_id + self.sampling_params = SimpleNamespace(extra_args=extra_args) + self.num_computed_tokens = 0 + self.step_tokens = step_tokens + self.skip_reading_prefix_cache = False + + +class _FakeWaitingQueue(deque): + """FCFS waiting queue: vLLM's is also a ``deque`` subclass.""" + + def remove_requests(self, requests: list[_FakeRequest]) -> None: + keep = [req for req in self if not any(req is dropped for dropped in requests)] + self.clear() + self.extend(keep) + + def prepend_request(self, request: _FakeRequest) -> None: + self.appendleft(request) + + +@dataclass +class _FakeSchedulerOutput: + num_scheduled_tokens: dict[str, int] = field(default_factory=dict) + total_num_scheduled_tokens: int = 0 + + +def _make_scheduler_cls() -> type: + """Fresh class per test so the in-place patches never stack.""" + + class FakeScheduler: + def __init__(self) -> None: + self.waiting = _FakeWaitingQueue() + self.running: list[_FakeRequest] = [] + self.requests: dict[str, _FakeRequest] = {} + self.finished: list[set[str]] = [] + self.observed_waiting: list[list[str]] = [] + self.max_num_scheduled_tokens = 512 + self.scheduler_config = SimpleNamespace(long_prefill_token_threshold=4096) + + def add_request(self, request: _FakeRequest) -> None: + self.requests[request.request_id] = request + self.waiting.append(request) + + def finish_requests(self, request_ids: Any, finished_status: Any) -> None: + if request_ids is None: + ids = set(self.requests) + elif isinstance(request_ids, str): + ids = {request_ids} + else: + ids = set(request_ids) + self.finished.append(ids) + self.running = [req for req in self.running if req.request_id not in ids] + for req_id in ids: + self.requests.pop(req_id, None) + + def schedule(self) -> _FakeSchedulerOutput: + self.observed_waiting.append([req.request_id for req in self.waiting]) + while self.waiting: + self.running.append(self.waiting.popleft()) + output = _FakeSchedulerOutput() + for req in self.running: + req.num_computed_tokens += req.step_tokens + output.num_scheduled_tokens[req.request_id] = req.step_tokens + output.total_num_scheduled_tokens += req.step_tokens + return output + + return FakeScheduler + + +@pytest.fixture +def scheduler() -> Any: + cls = _make_scheduler_cls() + _patch_scheduler(cls) + return cls() + + +@pytest.fixture(autouse=True) +def _restore_suffix_registry(): + original = set(cfg_pairing._UNCOND_REQUEST_ID_SUFFIXES) + yield + cfg_pairing._UNCOND_REQUEST_ID_SUFFIXES.clear() + cfg_pairing._UNCOND_REQUEST_ID_SUFFIXES.update(original) + + +def _add_pair( + scheduler: Any, + pair_id: str | None = None, + base_id: str = "req-1", + suffix: str = "__cfg_uncond", +) -> tuple[_FakeRequest, _FakeRequest]: + cond_extra: dict[str, Any] = {"cfg_role": "cond"} + uncond_extra: dict[str, Any] = {"cfg_role": "uncond"} + if pair_id is not None: + cond_extra["cfg_pair_id"] = pair_id + uncond_extra["cfg_pair_id"] = pair_id + cond = _FakeRequest(base_id, cond_extra) + uncond = _FakeRequest(f"{base_id}{suffix}", uncond_extra) + scheduler.add_request(cond) + scheduler.add_request(uncond) + return cond, uncond + + +# ───────── pair-id derivation ───────── + + +def test_pair_id_is_derived_from_request_id_when_absent(scheduler: Any) -> None: + cond, uncond = _add_pair(scheduler) + + assert scheduler._cfg_pairs == {"req-1": {"cond": cond.request_id, "uncond": uncond.request_id}} + assert scheduler._cfg_req_to_pair == {cond.request_id: "req-1", uncond.request_id: "req-1"} + + +def test_registered_companion_suffix_is_stripped(scheduler: Any) -> None: + register_cfg_uncond_suffix("::negative") + assert "::negative" in cfg_uncond_suffixes() + + cond, uncond = _add_pair(scheduler, base_id="music-7", suffix="::negative") + + assert scheduler._cfg_pairs == {"music-7": {"cond": cond.request_id, "uncond": uncond.request_id}} + + +def test_register_rejects_empty_suffix() -> None: + with pytest.raises(ValueError): + register_cfg_uncond_suffix("") + + +def test_uncond_without_known_suffix_keeps_its_own_request_id(scheduler: Any) -> None: + # Degenerate input: the companion never joins the cond pair, so the pair + # stays incomplete and the hold budget eventually releases it. + scheduler.add_request(_FakeRequest("solo-uncond", {"cfg_role": "uncond"})) + + assert scheduler._cfg_pairs == {"solo-uncond": {"uncond": "solo-uncond"}} + + +def test_explicit_pair_id_takes_precedence(scheduler: Any) -> None: + cond, uncond = _add_pair(scheduler, pair_id="pair-A") + + assert scheduler._cfg_pairs == {"pair-A": {"cond": cond.request_id, "uncond": uncond.request_id}} + assert cond.request_id not in scheduler._cfg_pairs + assert scheduler._cfg_req_to_pair[uncond.request_id] == "pair-A" + + +def test_cfg_requests_skip_prefix_cache_reads(scheduler: Any) -> None: + cond, uncond = _add_pair(scheduler) + + assert cond.skip_reading_prefix_cache is True + assert uncond.skip_reading_prefix_cache is True + + +# ───────── admission holding ───────── + + +def test_lone_member_is_held_then_released(scheduler: Any) -> None: + scheduler.add_request(_FakeRequest("req-1", {"cfg_role": "cond"})) + + scheduler.schedule() + + assert scheduler.observed_waiting == [[]] # the base scheduler saw nothing + assert [req.request_id for req in scheduler.waiting] == ["req-1"] + assert scheduler.running == [] + + +def test_complete_pair_is_admitted_together(scheduler: Any) -> None: + scheduler.add_request(_FakeRequest("req-1", {"cfg_role": "cond"})) + scheduler.schedule() + scheduler.add_request(_FakeRequest("req-1__cfg_uncond", {"cfg_role": "uncond"})) + + scheduler.schedule() + + assert scheduler.observed_waiting[-1] == ["req-1", "req-1__cfg_uncond"] + assert [req.request_id for req in scheduler.running] == ["req-1", "req-1__cfg_uncond"] + assert list(scheduler.waiting) == [] + + +def test_hold_budget_expiry_drops_the_stale_pair(scheduler: Any, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(cfg_pairing, "_MAX_PAIR_HOLD_STEPS", 2) + scheduler.add_request(_FakeRequest("req-1", {"cfg_role": "cond"})) + + scheduler.schedule() + scheduler.schedule() + assert scheduler.running == [] + assert scheduler._cfg_pairs # still hoping for the partner + + scheduler.schedule() + + assert scheduler._cfg_pairs == {} + assert scheduler._cfg_req_to_pair == {} + assert [req.request_id for req in scheduler.running] == ["req-1"] + + +def test_waiting_queue_reorder_puts_partners_adjacent(scheduler: Any) -> None: + scheduler.add_request(_FakeRequest("req-1", {"cfg_role": "cond"})) + scheduler.add_request(_FakeRequest("other", None)) + scheduler.add_request(_FakeRequest("req-1__cfg_uncond", {"cfg_role": "uncond"})) + + scheduler.schedule() + + assert scheduler.observed_waiting[-1] == ["req-1", "req-1__cfg_uncond", "other"] + + +# ───────── progress equalization ───────── + + +def _equalize_fixture( + cond_computed: int, + uncond_computed: int, + cond_sched: int, + uncond_sched: int, +) -> tuple[Any, _FakeSchedulerOutput, _FakeRequest, _FakeRequest]: + cond = _FakeRequest("req-1", {"cfg_role": "cond"}) + uncond = _FakeRequest("req-1__cfg_uncond", {"cfg_role": "uncond"}) + cond.num_computed_tokens = cond_computed + uncond.num_computed_tokens = uncond_computed + fake = SimpleNamespace( + _cfg_pairs={"req-1": {"cond": cond.request_id, "uncond": uncond.request_id}}, + requests={cond.request_id: cond, uncond.request_id: uncond}, + ) + output = _FakeSchedulerOutput( + num_scheduled_tokens={cond.request_id: cond_sched, uncond.request_id: uncond_sched}, + total_num_scheduled_tokens=cond_sched + uncond_sched, + ) + return fake, output, cond, uncond + + +def test_equalize_pulls_the_ahead_member_back() -> None: + fake, output, cond, uncond = _equalize_fixture(10, 6, 5, 4) + + _equalize_cfg_pair_progress(fake, output) + + assert cond.num_computed_tokens == 6 + assert uncond.num_computed_tokens == 6 + assert output.num_scheduled_tokens == {cond.request_id: 1, uncond.request_id: 4} + assert output.total_num_scheduled_tokens == 5 + + +def test_equalize_is_a_noop_when_the_pair_is_already_aligned() -> None: + fake, output, cond, uncond = _equalize_fixture(8, 8, 4, 4) + + _equalize_cfg_pair_progress(fake, output) + + assert (cond.num_computed_tokens, uncond.num_computed_tokens) == (8, 8) + assert output.total_num_scheduled_tokens == 8 + + +def test_equalize_skips_an_infeasible_rollback() -> None: + # Clawing back 4 tokens would leave the cond row with 0 scheduled tokens, + # which the scheduler output cannot express; leave the pair to _drop_split_pairs. + fake, output, cond, uncond = _equalize_fixture(10, 6, 4, 4) + + _equalize_cfg_pair_progress(fake, output) + + assert cond.num_computed_tokens == 10 + assert output.total_num_scheduled_tokens == 8 + + +# ───────── finish ───────── + + +def test_finish_requests_finishes_the_partner(scheduler: Any) -> None: + cond, uncond = _add_pair(scheduler) + scheduler.schedule() + + scheduler.finish_requests(cond.request_id, "FINISHED_STOPPED") + + assert scheduler.finished == [{cond.request_id, uncond.request_id}] + assert scheduler._cfg_pairs == {} + assert scheduler._cfg_req_to_pair == {} + assert scheduler.running == [] + + +def test_finish_all_requests_clears_the_registry(scheduler: Any) -> None: + _add_pair(scheduler) + scheduler.schedule() + + scheduler.finish_requests(None, "FINISHED_ABORTED") + + assert scheduler._cfg_pairs == {} + assert scheduler._cfg_req_to_pair == {} + + +# ───────── non-CFG traffic ───────── + + +def test_plain_requests_are_untouched(scheduler: Any) -> None: + scheduler.add_request(_FakeRequest("plain-1", None)) + scheduler.add_request(_FakeRequest("plain-2", {"seed": 3})) + + output = scheduler.schedule() + + assert scheduler._cfg_pairs == {} + assert scheduler.observed_waiting == [["plain-1", "plain-2"]] + assert output.num_scheduled_tokens == {"plain-1": 1, "plain-2": 1} + assert output.total_num_scheduled_tokens == 2 + # No pair ever existed, so the hold bookkeeping is never allocated. + assert not hasattr(scheduler, "_cfg_hold_counts") + + +def test_default_suffix_matches_the_audex_stage_processor() -> None: + from vllm_omni.model_executor.stage_input_processors.audex import CFG_UNCOND_SUFFIX + + assert CFG_UNCOND_SUFFIX in cfg_uncond_suffixes() + + +def test_audex_cfg_module_still_exports_the_patch_entry_points() -> None: + from vllm_omni.model_executor.models.audex.cfg import apply_cfg_patches, cfg_patches_applied + + assert apply_cfg_patches is cfg_pairing.apply_cfg_patches + assert cfg_patches_applied is cfg_pairing.cfg_patches_applied + + +# ───────── install gate ───────── + + +def test_gate_triggers_on_a_cfg_logits_processor() -> None: + from vllm_omni.engine.stage_init_utils import _stage_declares_cfg_pairs + + class AudexCFGLogitsProcessor: + pass + + assert _stage_declares_cfg_pairs(SimpleNamespace(logits_processors=[AudexCFGLogitsProcessor])) + # vLLM keeps unresolved entries as dotted strings. + assert _stage_declares_cfg_pairs( + SimpleNamespace(logits_processors=["vllm_omni.model_executor.models.audex.cfg:AudexCFGLogitsProcessor"]) + ) + + +def test_gate_triggers_on_a_declared_cfg_role_default() -> None: + from vllm_omni.engine.stage_init_utils import _stage_declares_cfg_pairs + + assert _stage_declares_cfg_pairs( + SimpleNamespace(logits_processors=None, sampling_extra_args_keys=("cfg_role", "cfg_scale")) + ) + + +def test_gate_stays_off_for_non_cfg_stages() -> None: + from vllm_omni.engine.stage_init_utils import _stage_declares_cfg_pairs + + assert not _stage_declares_cfg_pairs(SimpleNamespace(logits_processors=[], sampling_extra_args_keys=())) + assert not _stage_declares_cfg_pairs( + SimpleNamespace(logits_processors=["some.module:MinPLogitsProcessor"], sampling_extra_args_keys=("tta_rvq",)) + ) + assert not _stage_declares_cfg_pairs(SimpleNamespace()) + + +def test_sampling_extra_args_keys_are_derived_from_stage_defaults() -> None: + from vllm_omni.engine.stage_init_utils import _sampling_extra_args_keys + + assert _sampling_extra_args_keys({"extra_args": {"cfg_scale": 1.5, "cfg_role": "cond"}}) == ( + "cfg_role", + "cfg_scale", + ) + assert _sampling_extra_args_keys({"temperature": 0.1}) == () + assert _sampling_extra_args_keys(None) == () diff --git a/vllm_omni/config/model.py b/vllm_omni/config/model.py index 9164e1352..ef08dde8d 100644 --- a/vllm_omni/config/model.py +++ b/vllm_omni/config/model.py @@ -146,6 +146,11 @@ class OmniModelConfig(ModelConfig): task_type: str | None = None enable_sleep_mode: bool = False has_sampling_extra_args: bool = False + # Key names (not values) of the stage's default sampling ``extra_args``. + # Engine-core code runs before any request arrives, so this is the only + # place it can learn which request-shaping conventions a stage uses (e.g. + # ``cfg_role`` for classifier-free-guidance request pairs). + sampling_extra_args_keys: tuple[str, ...] = () @property def registry(self): diff --git a/vllm_omni/config/pipeline_registry.py b/vllm_omni/config/pipeline_registry.py index e8e3448b2..f498b7bde 100644 --- a/vllm_omni/config/pipeline_registry.py +++ b/vllm_omni/config/pipeline_registry.py @@ -85,6 +85,7 @@ from vllm_omni.model_executor.models.ming_tts.pipeline import ( MING_TTS_PIPELINE, ) from vllm_omni.model_executor.models.minicpmo_4_5.pipeline import MINICPMO_4_5_PIPELINE +from vllm_omni.model_executor.models.minimax_music3.pipeline import MINIMAX_MUSIC3_PIPELINE from vllm_omni.model_executor.models.moss_tts.pipeline import ( MOSS_TTS_LOCAL_PIPELINE, MOSS_TTS_PIPELINE, @@ -177,6 +178,7 @@ OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = { "moss_tts_realtime": MOSS_TTS_REALTIME_PIPELINE, "moss_tts_local": MOSS_TTS_LOCAL_PIPELINE, "minicpmo_4_5": MINICPMO_4_5_PIPELINE, + "minimax_music3": MINIMAX_MUSIC3_PIPELINE, "higgs_audio_v2": HIGGS_AUDIO_V2_PIPELINE, "higgs_multimodal_qwen3": HIGGS_AUDIO_V3_PIPELINE, "dynin_omni": DYNIN_OMNI_PIPELINE, diff --git a/vllm_omni/data_entry_keys.py b/vllm_omni/data_entry_keys.py index db2cf15ea..98575b689 100644 --- a/vllm_omni/data_entry_keys.py +++ b/vllm_omni/data_entry_keys.py @@ -91,6 +91,11 @@ class OmniPayloadMeta(TypedDict, total=False): ref_context_included: bool talker_prefill_offset: int omni_final_stage_id: int + # Per-request audio seed. Stages that draw their own noise (flow-matching + # / diffusion decoders) need the producing stage's seed to stay + # reproducible, and the producing stage's SamplingParams do not travel + # with the payload. + audio_seed: int class OmniPayload(TypedDict, total=False): @@ -189,6 +194,7 @@ class MetaStruct(_StructBase): codec_left_context_frames: int | None = None code_flat_numel: int | None = None omni_final_stage_id: int | None = None + audio_seed: int | None = None class OmniPayloadStruct(_StructBase): diff --git a/vllm_omni/deploy/minimax_music3.yaml b/vllm_omni/deploy/minimax_music3.yaml new file mode 100644 index 000000000..6e8a13d9d --- /dev/null +++ b/vllm_omni/deploy/minimax_music3.yaml @@ -0,0 +1,130 @@ +# MiniMax-Music3 deploy: Stage 0 (AR talker) -> Stage 1 (acoustic decoder), +# both on one 140 GB card, streaming 200-frame windows over shared memory. +# +# Auto-discovered default for model_type=minimax_music3. +# +# Classifier-free guidance is mandatory for this checkpoint: every request +# decodes as a cond/uncond pair in stage 0, so stage-0 KV, max_num_seqs and +# prefill budget are all sized for TWO rows per user request. The stage-0 +# extra_args block below is what turns on has_sampling_extra_args (so per-row +# sampling extra args reach the model's forward) and is also how the CFG +# scheduler patches learn this stage is paired. +pipeline: minimax_music3 +async_chunk: true + +# The repo-root config.json reports model_type=minimax_music3, which upstream +# transformers does not ship a config class for. +trust_remote_code: true + +connectors: + connector_of_shared_memory: + name: SharedMemoryConnector + extra: + # Stream stereo PCM back to the client as stage 1 finishes each window. + codec_streaming: true + # Window geometry is fixed by the checkpoint (200-frame windows on a + # 100-frame hop, models/minimax_music3/constants.py), so there are no + # codec_chunk_frames knobs here: changing the window changes the audio. + shm_threshold_bytes: 65536 + connector_get_sleep_s: 0.01 + # One window is 200 AR frames = 8 s of music, so stage 1 waits far longer + # for a chunk than a per-step codec stage would. First chunk additionally + # covers stage-0 prefill of both CFG rows. + connector_get_max_wait_first_chunk: 30000 + connector_get_max_wait: 10000 + +stages: + - stage_id: 0 + # 16 rows = 8 concurrent guided requests, since CFG doubles every request. + max_num_seqs: 16 + # ~16 GB of bf16 Qwen3-8B weights plus KV for 16 full-length rows + # (36 layers x 8 KV heads x 128 head_dim ~= 144 KB/token, so ~1.4 GB per + # 10240-token row) fits comfortably inside 0.6 of a 140 GB card. + gpu_memory_utilization: 0.6 + dtype: bfloat16 + # The checkpoint's max_position_embeddings is 10240 and the whole song is + # decoded in-context. Note this cannot hold a max-length prompt AND a + # max-length song at the same time: the prompt cap is 5000 tokens and the + # frame budget is 9000, which overruns 10240. That is a property of the + # checkpoint, not of this file; long captions shorten the maximum song. + max_model_len: 10240 + # One scheduler step must be able to prefill both rows of a pair + # (2 x 5000 prompt tokens) so the cond/uncond pair stays step-locked. + max_num_batched_tokens: 10240 + # Two rows at 8B is a launch-bound decode step, so the backbone is + # captured. FULL_DECODE_ONLY only ever dispatches a graph for a uniform + # decode batch; prefill and mixed steps stay eager, which is what the + # talker's frame-feedback fallback is for. The talker's own sampler and + # depth decoder run after the graph, not inside it. + # + # What makes the forward capturable: the frame feedback writes a + # fixed-size slab in prepare_runner_inputs (outside the captured region) + # and the forward selects over it with torch.where, instead of an + # index_copy against an index tensor whose address and length changed + # every step. The row-to-request binding the sampler reads is likewise + # rebuilt in prepare_runner_inputs, because a replayed step never runs + # the forward's Python body. + enforce_eager: false + compilation_config: + # No compilation. The talker walks the backbone's decoder layers itself, + # so the compiled Qwen3Model.forward is never entered and inductor has + # nothing to do; what the mode still decides is which kernels vLLM + # defaults to. Left at its non-eager default it picks the elementwise + # PyTorch norms over the fused vllm_c ones, which is both slower and + # numerically different: measured here, that alone moved the sampled + # codes off the eager trajectory by frame 4. Graph capture is + # independent of this and still happens. + mode: 0 + # 16 rows is the stage ceiling, i.e. 8 guided requests. + cudagraph_capture_sizes: [1, 2, 4, 8, 16] + cudagraph_mode: FULL_DECODE_ONLY + cudagraph_num_of_warmups: 1 + # Belt and braces on the same point: the non-eager default is "none", + # which hands the norm and activation ops to elementwise PyTorch. + custom_ops: ["all"] + # A CFG pair must be sampled inside one engine step for the logits blend + # and the cond->uncond token copy to line up. + async_scheduling: false + # Off during bring-up: the uncond row is a long run of identical + # <|audio_cfg|> tokens and would take a much larger prefix-cache hit than + # its cond partner, which is exactly the num_computed_tokens skew the CFG + # pair patches then have to undo. + enable_prefix_caching: false + devices: "0" + output_connectors: + to_stage_1: connector_of_shared_memory + default_sampling_params: + temperature: 1.0 + # The frame budget is enforced by the model. This ceiling only has to + # leave room for the extra step that delivers the final window. + max_tokens: 9002 + extra_args: + cfg_role: cond + cfg_scale: 1.5 + + - stage_id: 1 + # One row per user request: the uncond companion is stage-0-final and + # never reaches the acoustic stage. + max_num_seqs: 8 + # No KV cache on this stage; the budget only covers the DiT and DAV + # weights plus per-window latents and activations. + gpu_memory_utilization: 0.2 + # The flow-matching solver degrades measurably in bfloat16 + # (SUPPORTED_ACOUSTIC_DTYPES in models/minimax_music3/constants.py). + dtype: float32 + enforce_eager: true + # Async scheduling builds per-request payloads against the NEXT step's + # input batch, which breaks when chunk arrival staggers requests. + async_scheduling: false + enable_prefix_caching: false + # One placeholder prompt token per arriving window, so a six-minute song + # (90 windows) needs a tiny context. + max_model_len: 4096 + max_num_batched_tokens: 4096 + devices: "0" + input_connectors: + from_stage_0: connector_of_shared_memory + default_sampling_params: + temperature: 0.0 + max_tokens: 4096 + seed: 0 diff --git a/vllm_omni/deploy/minimax_music3_2gpu.yaml b/vllm_omni/deploy/minimax_music3_2gpu.yaml new file mode 100644 index 000000000..8c463e0e7 --- /dev/null +++ b/vllm_omni/deploy/minimax_music3_2gpu.yaml @@ -0,0 +1,97 @@ +# MiniMax-Music3 two-GPU layout: AR talker on device 0, acoustic decoder on +# device 1. +# +# Splitting the stages removes the memory-budget compromise of the single-card +# file: the talker's CFG-doubled KV no longer competes with the DiT and DAV +# weights, and the flow-matching solver no longer competes with AR decode for +# SMs. Everything else (window geometry, CFG contract, sampling defaults) +# matches minimax_music3.yaml, so the two files stay behaviorally identical. +pipeline: minimax_music3 +async_chunk: true + +# The repo-root config.json reports model_type=minimax_music3, which upstream +# transformers does not ship a config class for. +trust_remote_code: true + +connectors: + connector_of_shared_memory: + name: SharedMemoryConnector + extra: + # Stream stereo PCM back to the client as stage 1 finishes each window. + codec_streaming: true + # Window geometry is fixed by the checkpoint (200-frame windows on a + # 100-frame hop, models/minimax_music3/constants.py), so there are no + # codec_chunk_frames knobs here: changing the window changes the audio. + shm_threshold_bytes: 65536 + connector_get_sleep_s: 0.01 + # One window is 200 AR frames = 8 s of music, so stage 1 waits far longer + # for a chunk than a per-step codec stage would. First chunk additionally + # covers stage-0 prefill of both CFG rows. + connector_get_max_wait_first_chunk: 30000 + connector_get_max_wait: 10000 + +stages: + - stage_id: 0 + # 32 rows = 16 concurrent guided requests, since CFG doubles every request. + # Twice the single-card figure: this stage now owns its whole device. + max_num_seqs: 32 + gpu_memory_utilization: 0.85 + dtype: bfloat16 + # The checkpoint's max_position_embeddings is 10240 and the whole song is + # decoded in-context. Note this cannot hold a max-length prompt AND a + # max-length song at the same time: the prompt cap is 5000 tokens and the + # frame budget is 9000, which overruns 10240. That is a property of the + # checkpoint, not of this file; long captions shorten the maximum song. + max_model_len: 10240 + # One scheduler step must be able to prefill both rows of a pair + # (2 x 5000 prompt tokens) so the cond/uncond pair stays step-locked. + max_num_batched_tokens: 10240 + # The talker owns its sampler and drives a depth decoder per AR step; + # leave graph capture off until that path is verified graph-safe. + enforce_eager: true + # A CFG pair must be sampled inside one engine step for the logits blend + # and the cond->uncond token copy to line up. + async_scheduling: false + # Off during bring-up: the uncond row is a long run of identical + # <|audio_cfg|> tokens and would take a much larger prefix-cache hit than + # its cond partner, which is exactly the num_computed_tokens skew the CFG + # pair patches then have to undo. + enable_prefix_caching: false + devices: "0" + output_connectors: + to_stage_1: connector_of_shared_memory + default_sampling_params: + temperature: 1.0 + # The frame budget is enforced by the model. This ceiling only has to + # leave room for the extra step that delivers the final window. + max_tokens: 9002 + extra_args: + cfg_role: cond + cfg_scale: 1.5 + + - stage_id: 1 + # One row per user request: the uncond companion is stage-0-final and + # never reaches the acoustic stage. + max_num_seqs: 16 + # No KV cache on this stage, but it now has a device to itself, so the + # solver can hold more windows in flight than the shared-card profile. + gpu_memory_utilization: 0.6 + # The flow-matching solver degrades measurably in bfloat16 + # (SUPPORTED_ACOUSTIC_DTYPES in models/minimax_music3/constants.py). + dtype: float32 + enforce_eager: true + # Async scheduling builds per-request payloads against the NEXT step's + # input batch, which breaks when chunk arrival staggers requests. + async_scheduling: false + enable_prefix_caching: false + # One placeholder prompt token per arriving window, so a six-minute song + # (90 windows) needs a tiny context. + max_model_len: 4096 + max_num_batched_tokens: 4096 + devices: "1" + input_connectors: + from_stage_0: connector_of_shared_memory + default_sampling_params: + temperature: 0.0 + max_tokens: 4096 + seed: 0 diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index 29c84a1c9..fcd11b173 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -232,6 +232,7 @@ class OmniEngineArgs(EngineArgs): log_stats: bool = False custom_pipeline_args: dict[str, Any] | None = None has_sampling_extra_args: bool = False + sampling_extra_args_keys: tuple[str, ...] = () def __post_init__(self) -> None: if self.worker_cls is None: @@ -409,6 +410,7 @@ class OmniEngineArgs(EngineArgs): omni_kv_config=self.omni_kv_config, task_type=self.task_type, has_sampling_extra_args=self.has_sampling_extra_args, + sampling_extra_args_keys=tuple(self.sampling_extra_args_keys or ()), ) return omni_config diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index 38eea768e..e35315f7e 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -743,23 +743,31 @@ class AsyncOmniEngine: enqueue_ts=time.perf_counter(), ) - def _enqueue_cfg_companions( + def _build_cfg_companions( self, parent_id: str, original_prompt: Any, stage0_params: Any, sampling_params_list: list[Any], - ) -> None: - """Expand prompt into CFG companions, process through InputProcessor, and enqueue.""" - try: - expanded = self.prompt_expand_func(original_prompt, stage0_params) - except Exception: - logger.exception("[AsyncOmniEngine] prompt_expand_func failed for req %s", parent_id) - return + ) -> list[AddCompanionRequestMessage]: + """Expand a prompt into its CFG companions, without enqueueing any. + Construction is separated from admission so a guided request is + all-or-nothing. A model whose guidance is mandatory cannot decode a + request whose companion never arrived: the pair never completes, the + request occupies scheduler and KV capacity for the scheduler's whole + hold budget, and then produces no audio. Raising here instead means the + caller learns immediately and nothing was admitted. + + Raises: + Exception: Whatever prompt expansion or input processing raised. + The caller is expected to let it reach the client. + """ + expanded = self.prompt_expand_func(original_prompt, stage0_params) if not expanded: - return + return [] + companions: list[AddCompanionRequestMessage] = [] for ep in expanded: cid = f"{parent_id}{ep.request_id_suffix}" companion_prompt = ep.prompt @@ -775,6 +783,14 @@ class AsyncOmniEngine: params=companion_params, supported_tasks=self.supported_tasks, ) + # Same restore the parent request gets: the upstream input + # processor drops omni-only prompt fields, so without this the + # companion reaches the worker with no additional_information at + # all. That is where ``global_request_id`` lives, and it is what + # was just injected above, so skipping it silently undoes the + # injection: the model sees the companion row with no id and + # cannot match it to its conditioned partner. + request = upgrade_to_omni_request(request, companion_prompt) request.external_req_id = cid # Companions are stage-0-final for ordinary downstream payloads, # but diffusion still needs their CFG KV caches. @@ -783,7 +799,7 @@ class AsyncOmniEngine: # Registration of this companion on stage-0's output processor is # deferred to Orchestrator._handle_add_companion, which routes # admission through StagePool.submit_initial(..., affinity_request_id=...). - self.request_queue.sync_q.put( + companions.append( AddCompanionRequestMessage( companion_id=cid, parent_id=parent_id, @@ -793,11 +809,35 @@ class AsyncOmniEngine: sampling_params_list=companion_spl, ) ) + return companions + + def _enqueue_cfg_companions( + self, + parent_id: str, + original_prompt: Any, + stage0_params: Any, + sampling_params_list: list[Any], + ) -> None: + """Build and enqueue CFG companions, tolerating a build failure. + + Kept for callers that admit the parent first and cannot roll it back. + Prefer building with :meth:`_build_cfg_companions` before the parent is + admitted, so the pair is atomic. + """ + try: + companions = self._build_cfg_companions(parent_id, original_prompt, stage0_params, sampling_params_list) + except Exception: + logger.exception("[AsyncOmniEngine] CFG companion build failed for req %s", parent_id) + return + for companion in companions: + self.request_queue.sync_q.put(companion) + if not companions: + return logger.info( "[AsyncOmniEngine] CFG expansion for req %s: %d companions", parent_id, - len(expanded), + len(companions), ) @staticmethod @@ -1279,16 +1319,28 @@ class AsyncOmniEngine: reasoning_ended=reasoning_ended, resumable=resumable, ) - self.request_queue.sync_q.put(msg) - - # CFG companion expansion: create and enqueue companion requests - # so the AR stage also generates their KV caches. + # CFG companions are built before the parent is admitted, so the group + # is all-or-nothing: a build failure raises here, nothing is enqueued, + # and the caller sees the error. Admitting the parent first would leave + # an orphan holding scheduler and KV capacity that can never complete, + # because a model whose guidance is mandatory cannot decode a request + # whose companion never arrived. + companions: list[AddCompanionRequestMessage] = [] if self.prompt_expand_func is not None and final_stage_id > 0: - original_prompt = msg.original_prompt effective_spl = msg.sampling_params_list stage0_params = effective_spl[0] if effective_spl else None if stage0_params is not None: - self._enqueue_cfg_companions(request_id, original_prompt, stage0_params, effective_spl) + companions = self._build_cfg_companions(request_id, msg.original_prompt, stage0_params, effective_spl) + + self.request_queue.sync_q.put(msg) + for companion in companions: + self.request_queue.sync_q.put(companion) + if companions: + logger.info( + "[AsyncOmniEngine] CFG expansion for req %s: %d companions", + request_id, + len(companions), + ) async def add_request_async( self, diff --git a/vllm_omni/engine/stage_engine_core_proc.py b/vllm_omni/engine/stage_engine_core_proc.py index 3ac719a6b..4bcf47d3b 100644 --- a/vllm_omni/engine/stage_engine_core_proc.py +++ b/vllm_omni/engine/stage_engine_core_proc.py @@ -31,7 +31,7 @@ from vllm.v1.engine.utils import ( from vllm_omni.distributed.omni_coordinator import create_stage_coord_client from vllm_omni.engine import OmniEngineCoreRequest from vllm_omni.engine.stage_init_utils import ( - maybe_apply_audex_cfg_patches, + maybe_apply_cfg_scheduler_patches, set_death_signal, ) @@ -134,9 +134,10 @@ class StageEngineCoreProc(EngineCoreProc): _vllm_engine_core_module.EngineCoreRequest, ) - # Audex CFG scheduler patches must land before EngineCore builds - # its Scheduler; gated on the stage's logits_processors config. - maybe_apply_audex_cfg_patches(kwargs.get("vllm_config")) + # CFG pairing scheduler patches must land before EngineCore builds + # its Scheduler; gated on the stage's logits_processors and its + # default sampling extra_args. + maybe_apply_cfg_scheduler_patches(kwargs.get("vllm_config")) engine_core = StageEngineCoreProc( *args, diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index 054a2fb22..9da6f2cfa 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/vllm_omni/engine/stage_init_utils.py @@ -871,6 +871,20 @@ def _project_omni_stage_engine_args( return engine_args +def _sampling_extra_args_keys(default_sampling_params: Any) -> tuple[str, ...]: + """Key names of a stage's default sampling ``extra_args``, sorted. + + Only the keys travel into the engine config: engine-core code needs to know + which request-shaping conventions a stage uses (e.g. CFG request pairing) + before any request exists, while the values stay a serving-layer concern. + """ + extra_args = _to_dict(default_sampling_params or {}).get("extra_args") or {} + try: + return tuple(sorted(str(key) for key in extra_args)) + except TypeError: + return () + + def _finalize_engine_args_dict( engine_args_dict: dict[str, Any], *, @@ -880,6 +894,7 @@ def _finalize_engine_args_dict( stage_connector_spec: dict[str, Any] | None, cli_tokenizer: str | None, has_sampling_extra_args: bool, + sampling_extra_args_keys: tuple[str, ...] = (), ) -> dict[str, Any]: """Apply representation-independent engine adapter behavior.""" pipeline_model_root = model @@ -960,6 +975,7 @@ def _finalize_engine_args_dict( engine_args_dict.setdefault("enable_prefix_caching", False) engine_args_dict["has_sampling_extra_args"] = has_sampling_extra_args + engine_args_dict["sampling_extra_args_keys"] = sampling_extra_args_keys # TODO: Remove this after the performance regression is fixed # Set VLLM_USE_FLASHINFER_MOE_FP16=0 for Qwen3-Omni to avoid performance regression @@ -990,6 +1006,7 @@ def build_legacy_engine_args_dict( stage_connector_spec=stage_connector_spec, cli_tokenizer=cli_tokenizer, has_sampling_extra_args=bool(default_sp.get("extra_args")), + sampling_extra_args_keys=_sampling_extra_args_keys(default_sp), ) @@ -1036,6 +1053,7 @@ def build_engine_args_dict_from_omni_stage_config( stage_connector_spec=stage_connector_spec, cli_tokenizer=cli_tokenizer, has_sampling_extra_args=stage_config.model_config.has_sampling_extra_args, + sampling_extra_args_keys=_sampling_extra_args_keys(stage_config.model_config.default_sampling_params), ) @@ -1445,24 +1463,45 @@ def initialize_diffusion_stage( return create_diffusion_client(model, od_config, metadata, stage_init_timeout, batch_size, use_inline) -def maybe_apply_audex_cfg_patches(vllm_config: Any) -> None: - """Install the Audex CFG scheduler patches for CFG-configured engines. +def _stage_declares_cfg_pairs(model_config: Any) -> bool: + """Whether this stage submits classifier-free-guidance request pairs. + + Two independent declarations, because a model may own either side of the + mechanism without the other: + + * a CFG logits processor is configured (blending implies pairing), or + * the stage's default sampling ``extra_args`` carry ``cfg_role``, which is + how a model declares paired requests without owning a logits processor. + Only the *key set* survives into the engine config (see + ``sampling_extra_args_keys``); the values stay in the serving layer. + """ + processors = getattr(model_config, "logits_processors", None) or [] + if any("CFGLogitsProcessor" in getattr(proc, "__name__", str(proc)) for proc in processors): + return True + return "cfg_role" in (getattr(model_config, "sampling_extra_args_keys", None) or ()) + + +def maybe_apply_cfg_scheduler_patches(vllm_config: Any) -> None: + """Install the CFG pairing scheduler patches for CFG-configured engines. Must run in the engine-core process BEFORE ``Scheduler`` is constructed: the patch wraps ``Scheduler.__init__`` to add the pair registry, so a - scheduler built earlier would never become pair-aware. Gated on the - engine's ``logits_processors`` so non-CFG stages stay untouched. + scheduler built earlier would never become pair-aware. Gated so non-CFG + stages stay untouched. """ model_config = getattr(vllm_config, "model_config", None) - processors = getattr(model_config, "logits_processors", None) or [] - if not any("AudexCFGLogitsProcessor" in getattr(proc, "__name__", str(proc)) for proc in processors): + if model_config is None or not _stage_declares_cfg_pairs(model_config): return - from vllm_omni.model_executor.models.audex.cfg import apply_cfg_patches + from vllm_omni.model_executor.models.common.cfg_pairing import apply_cfg_patches apply_cfg_patches() from vllm.v1.core.sched.scheduler import Scheduler - if not getattr(Scheduler.schedule, "_audex_cfg_patched", False): - raise RuntimeError("Audex CFG scheduler patches failed to install before Scheduler construction") + if not getattr(Scheduler.schedule, "_cfg_pairing_patched", False): + raise RuntimeError("CFG pairing scheduler patches failed to install before Scheduler construction") + + +# Name kept for callers written against the Audex-only gate. +maybe_apply_audex_cfg_patches = maybe_apply_cfg_scheduler_patches diff --git a/vllm_omni/entrypoints/openai/tts_adapters/__init__.py b/vllm_omni/entrypoints/openai/tts_adapters/__init__.py index 74597a2d1..e725e65cd 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/__init__.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/__init__.py @@ -111,6 +111,7 @@ from vllm_omni.entrypoints.openai.tts_adapters import ( # noqa: E402,F401 indextts2, ming_flash_omni_tts, ming_tts, + minimax_music3, moss_tts, omnivoice, qwen3_tts, diff --git a/vllm_omni/entrypoints/openai/tts_adapters/minimax_music3.py b/vllm_omni/entrypoints/openai/tts_adapters/minimax_music3.py new file mode 100644 index 000000000..50b9c2d71 --- /dev/null +++ b/vllm_omni/entrypoints/openai/tts_adapters/minimax_music3.py @@ -0,0 +1,163 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""MiniMax Music 3 serving adapter for ``/v1/audio/speech``. + +This is a text-to-music model on a speech endpoint, so most of the speech +contract does not apply. ``input`` carries the lyrics and ``instructions`` the +caption that describes genre, instrumentation, tempo and production. There is +no speaker to select, no reference audio, and no temperature: sampling is +fixed by the checkpoint. Unsupported parameters are rejected rather than +silently ignored, so a mistaken request fails immediately instead of quietly +returning something the caller did not ask for. +""" + +from typing import TYPE_CHECKING, Any + +from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter +from vllm_omni.entrypoints.openai.tts_adapters.base import ARTTSAdapter, PreparedRequest +from vllm_omni.model_executor.models.minimax_music3.constants import ( + AUDIO_FRAME_RATE, + MAX_AUDIO_FRAMES, + MAX_PROMPT_TOKENS, +) +from vllm_omni.model_executor.models.minimax_music3.prompt import ( + build_prompt, + validate_tokenizer_ids, +) + +if TYPE_CHECKING: + from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest + +# Captions are long by design: the checkpoint was trained on structured +# descriptions covering global attributes, emotional progression, vocal detail +# and arrangement. The reference captions run to roughly 5,000 characters, so +# the generic 500-character instructions cap does not apply here. The real +# limit is the tokenized prompt cap, enforced downstream. +_MAX_CAPTION_CHARS = 20_000 + +# Sampling is fixed by the checkpoint: guidance at 1.5, then a seeded top-k 50 +# draw. Accepting these would imply a control the model does not have. +_UNSUPPORTED_SAMPLING_PARAMS = ("temperature", "top_p", "top_k", "repetition_penalty") + +# CFG pair plumbing is derived from the request id by the scheduler. A caller +# supplying these would desync the guided pair. +_INTERNAL_CFG_KEYS = ("cfg_role", "cfg_pair_id", "cfg_null_prompt") + + +@register_tts_adapter +class MiniMaxMusic3Adapter(ARTTSAdapter): + """Lyrics plus caption in, one 32 kHz stereo song out.""" + + name = "minimax_music3" + stage_keys = frozenset({"minimax_music3_ar"}) + # 9,000 frames at 25 fps is six minutes, the checkpoint's own ceiling. + max_new_tokens_min = 1 + max_new_tokens_max = MAX_AUDIO_FRAMES + + def __init__(self, ctx: Any) -> None: + super().__init__(ctx) + self._cached_tokenizer: Any = None + + def validate(self, request: "OpenAICreateSpeechRequest") -> str | None: + if not request.input or not request.input.strip(): + return "MiniMax Music 3 requires non-empty 'input' carrying the lyrics" + + instructions = request.instructions + if not isinstance(instructions, str) or not instructions.strip(): + return ( + "MiniMax Music 3 requires 'instructions' describing the music " + "(genre, instrumentation, tempo, mood); it is what decides the arrangement" + ) + if len(instructions) > _MAX_CAPTION_CHARS: + return f"'instructions' exceeds the maximum length of {_MAX_CAPTION_CHARS} characters" + + if request.max_new_tokens is not None: + if request.max_new_tokens < self.max_new_tokens_min: + return f"'max_new_tokens' must be at least {self.max_new_tokens_min}" + if request.max_new_tokens > self.max_new_tokens_max: + return ( + f"'max_new_tokens' counts audio frames at {AUDIO_FRAME_RATE} per second " + f"and cannot exceed {self.max_new_tokens_max}" + ) + + voice = (request.voice or "").strip().lower() + if voice not in ("", "default"): + return ( + "MiniMax Music 3 has no speaker to select; the vocal comes from " + f"'instructions'. Omit 'voice' or pass 'default'; got {request.voice!r}" + ) + if request.ref_audio is not None or request.ref_text is not None: + return "MiniMax Music 3 does not support reference-audio conditioning" + if request.language is not None: + return "MiniMax Music 3 has no language tag; the language follows the lyrics" + if request.task_type is not None: + return "MiniMax Music 3 does not support 'task_type'" + if request.speed is not None and float(request.speed) != 1.0: + return "MiniMax Music 3 only supports speed=1.0; put the tempo in 'instructions', e.g. 'at 92 BPM'" + + extra: dict[str, Any] = request.extra_params or {} + for key in _INTERNAL_CFG_KEYS: + if key in extra: + return f"extra_params.{key} is managed internally by the server and cannot be set" + rejected = sorted(key for key in _UNSUPPORTED_SAMPLING_PARAMS if extra.get(key) is not None) + if rejected: + return "MiniMax Music 3 has fixed sampling and does not accept: " + ", ".join(rejected) + return None + + async def build( + self, + request: "OpenAICreateSpeechRequest", + sampling_params_list: list, + has_inline_ref_audio: bool, + ) -> PreparedRequest: + del sampling_params_list, has_inline_ref_audio + prompt = build_prompt(request.instructions, request.input) + # Submit token ids, not text. Guidance needs a companion whose prompt is + # the same length with the caption-and-lyrics span replaced, and that + # length match can only be guaranteed on ids. Handing over text would + # leave the expansion with nothing to build the twin from, and a request + # that decodes unguided produces a different song. + token_ids = self._tokenizer().encode(prompt) + if len(token_ids) > MAX_PROMPT_TOKENS: + raise ValueError( + f"MiniMax Music 3 prompt is {len(token_ids)} tokens; " + f"the maximum is {MAX_PROMPT_TOKENS}. Shorten the caption or the lyrics." + ) + # The frame budget is enforced by the model, not by the engine's + # max_tokens. The model needs one decode step beyond the last frame to + # deliver the song's final window, so capping the engine at exactly + # this many tokens would cut the tail off. + max_frames = int(request.max_new_tokens or MAX_AUDIO_FRAMES) + # Carried on the prompt rather than through tts_params: the adapter + # path hands PreparedRequest.prompt to the engine untouched, so only + # what is attached here reaches the model's per-request info. + return PreparedRequest( + prompt={ + "prompt_token_ids": token_ids, + "additional_information": {"max_audio_frames": [max_frames]}, + }, + tts_params={"max_audio_frames": [max_frames]}, + model_type=self.name, + ) + + def _tokenizer(self) -> Any: + """The music tokenizer, loaded once per process. + + Its special-token ids are pinned: a checkpoint whose tokenizer + disagrees is not this model, and quietly producing noise is worse than + failing at the first request. + """ + if self._cached_tokenizer is None: + from transformers import AutoTokenizer + + stage = getattr(self.ctx.server, "_tts_stage", None) + model_path = getattr(getattr(stage, "engine_args", None), "tokenizer", None) + if not model_path: + model_path = self.ctx.server.engine_client.model_config.tokenizer + tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True) + validate_tokenizer_ids(tokenizer) + self._cached_tokenizer = tokenizer + return self._cached_tokenizer + + +__all__ = ["MiniMaxMusic3Adapter"] diff --git a/vllm_omni/model_executor/models/audex/cfg.py b/vllm_omni/model_executor/models/audex/cfg.py index 488097216..3c1837140 100644 --- a/vllm_omni/model_executor/models/audex/cfg.py +++ b/vllm_omni/model_executor/models/audex/cfg.py @@ -3,8 +3,7 @@ # # Adapted from nvidia/Nemotron-Labs-Audex-2B (Apache-2.0): # inference_scripts_vllm/audiogen_scripts/cfg_logits_processor.py -# inference_scripts_vllm/audiogen_scripts/vllm_cfg_patch.py -# Divergences from the official files are commented at the divergence site. +# Divergences from the official file are commented at the divergence site. """Classifier-free guidance for the Audex thinker stage. CFG pairs a conditional request with an unconditional (null-prompt) @@ -20,16 +19,13 @@ Requests opt in via ``SamplingParams.extra_args``: uncond: {"cfg_scale": 1.5, "cfg_role": "uncond", "cfg_pair_id": } Blending is only correct when both pair members decode the same position in -the same engine step, so ``apply_cfg_patches()`` makes the vLLM scheduler -pair-aware: a lone member waits for its partner, partners stay adjacent, -their ``num_computed_tokens`` are equalized after every schedule, and they -finish together. The patch is a no-op for engines that never see a CFG -request (the pair registry stays empty). +the same engine step. Keeping the pair step-locked is model-neutral and lives +in :mod:`vllm_omni.model_executor.models.common.cfg_pairing`; the scheduler +patches are re-exported here so existing imports keep working. """ from __future__ import annotations -from collections import deque from typing import Any import torch @@ -42,18 +38,20 @@ from vllm.v1.sample.logits_processor import ( MoveDirectionality, ) +from vllm_omni.model_executor.models.common.cfg_pairing import ( + _COND, + _UNCOND, + apply_cfg_patches, + cfg_patches_applied, +) + logger = init_logger(__name__) -_COND = "cond" -_UNCOND = "uncond" - - -def _cfg_extra(request: Any, key: str) -> Any: - sampling_params = getattr(request, "sampling_params", None) - extra_args = getattr(sampling_params, "extra_args", None) - if extra_args: - return extra_args.get(key) - return None +__all__ = [ + "AudexCFGLogitsProcessor", + "apply_cfg_patches", + "cfg_patches_applied", +] class AudexCFGLogitsProcessor(LogitsProcessor): @@ -208,360 +206,3 @@ class AudexCFGLogitsProcessor(LogitsProcessor): logits[uncond_idx] = blended return logits - - -# --------------------------------------------------------------------------- -# Scheduler pair synchronization -# --------------------------------------------------------------------------- - - -def _wait_queues(scheduler: Any) -> list[Any]: - # Divergence from the official (vLLM 0.20-era) file: vLLM 0.24 parks - # structurally not-ready requests in a second ``skipped_waiting`` queue; - # pair-hold must cover both queues or a lone member parked there would - # be admitted without its partner. - queues = [scheduler.waiting] - skipped = getattr(scheduler, "skipped_waiting", None) - if skipped is not None: - queues.append(skipped) - return queues - - -def _pair_complete(scheduler: Any, request_id: str, blocked_ids: set[str] | None = None) -> bool: - pair_id = scheduler._cfg_req_to_pair.get(request_id) - if pair_id is None: - return True - roles = scheduler._cfg_pairs.get(pair_id, {}) - if len(roles) != 2 or not all(rid in scheduler.requests for rid in roles.values()): - return False - if blocked_ids: - # A partner parked in skipped_waiting (e.g. waiting for remote KVs) - # may fail promotion this step while this member gets scheduled from - # waiting — hold this member until the partner is schedulable. - for rid in roles.values(): - if rid != request_id and rid in blocked_ids: - return False - return True - - -# Schedule steps a lone pair member may wait for its partner. Partners arrive -# within a handful of steps in practice; exhausting this budget means the -# companion was never created (e.g. prompt expansion failed), and the request -# is released to decode unguided instead of hanging forever. -_MAX_PAIR_HOLD_STEPS = 512 - - -def _drop_stale_pair(scheduler: Any, request_id: str) -> None: - pair_id = scheduler._cfg_req_to_pair.get(request_id) - if pair_id is None: - return - for rid in scheduler._cfg_pairs.pop(pair_id, {}).values(): - scheduler._cfg_req_to_pair.pop(rid, None) - logger.warning( - "CFG pair %s: partner of %s never arrived after %d schedule steps; releasing the request unguided", - pair_id, - request_id, - _MAX_PAIR_HOLD_STEPS, - ) - - -def _hold_incomplete_pairs(scheduler: Any) -> list[tuple[Any, Any]]: - """Pull CFG requests whose partner is not yet admitted out of the wait queues. - - The engine schedules as soon as the first member of a pair arrives over - IPC; without the hold, the lone member prefills one step early and the - pair is permanently offset by one token. Held requests are re-prepended - after the schedule step. - """ - hold_counts: dict[str, int] = getattr(scheduler, "_cfg_hold_counts", None) or {} - scheduler._cfg_hold_counts = hold_counts - - skipped = getattr(scheduler, "skipped_waiting", None) - blocked_ids = {req.request_id for req in list(skipped)} if skipped is not None else set() - - held: list[tuple[Any, Any]] = [] - for queue in _wait_queues(scheduler): - # A member is only "blocked" for its partner's sake when it sits in - # the OTHER queue; members of the queue being scanned move together. - partner_blockers = blocked_ids if queue is not skipped else set() - held_here = [] - for request in list(queue): - if _pair_complete(scheduler, request.request_id, partner_blockers): - hold_counts.pop(request.request_id, None) - continue - count = hold_counts.get(request.request_id, 0) + 1 - if count > _MAX_PAIR_HOLD_STEPS: - hold_counts.pop(request.request_id, None) - _drop_stale_pair(scheduler, request.request_id) - continue - hold_counts[request.request_id] = count - held_here.append(request) - if held_here: - queue.remove_requests(held_here) - held.extend((queue, req) for req in held_here) - return held - - -def _release_held(held: list[tuple[Any, Any]]) -> None: - for queue, request in reversed(held): - queue.prepend_request(request) - - -def _reorder_waiting_for_cfg(scheduler: Any) -> None: - """Move CFG pair partners adjacent in the FCFS waiting queue.""" - waiting = scheduler.waiting - # Only the FCFS queue (a deque) has caller-controlled ordering; priority - # queues order themselves and pair sync then relies on hold + equalize. - if not isinstance(waiting, deque) or len(waiting) < 2: - return - - requests = list(waiting) - waiting.clear() - - seen: set[str] = set() - result: list[Any] = [] - for request in requests: - rid = request.request_id - if rid in seen: - continue - seen.add(rid) - result.append(request) - - pair_id = scheduler._cfg_req_to_pair.get(rid) - if pair_id is None: - continue - for partner_id in scheduler._cfg_pairs.get(pair_id, {}).values(): - if partner_id != rid and partner_id not in seen: - for candidate in requests: - if candidate.request_id == partner_id: - seen.add(partner_id) - result.append(candidate) - break - - waiting.extend(result) - - -def _equalize_cfg_pair_progress(scheduler: Any, scheduler_output: Any) -> None: - """Bring both members of every CFG pair to the same ``num_computed_tokens``. - - Prefix-cache hits or unequal chunked-prefill budget can leave one member - ahead after a schedule step; reduce the faster member's allocation so - both land on the same position. Over-allocated KV blocks remain and are - consumed next step — nothing leaks. - """ - for roles in scheduler._cfg_pairs.values(): - cond_id = roles.get(_COND) - uncond_id = roles.get(_UNCOND) - if not cond_id or not uncond_id: - continue - - cond_sched = scheduler_output.num_scheduled_tokens.get(cond_id, 0) - uncond_sched = scheduler_output.num_scheduled_tokens.get(uncond_id, 0) - if cond_sched == 0 or uncond_sched == 0: - continue - - cond_req = scheduler.requests.get(cond_id) - uncond_req = scheduler.requests.get(uncond_id) - if not cond_req or not uncond_req: - continue - - if cond_req.num_computed_tokens == uncond_req.num_computed_tokens: - continue - - target = min(cond_req.num_computed_tokens, uncond_req.num_computed_tokens) - - feasible = all( - req.num_computed_tokens - target < sched - for req, sched in ((cond_req, cond_sched), (uncond_req, uncond_sched)) - ) - if not feasible: - continue - - for req_id, req, orig_sched in ( - (cond_id, cond_req, cond_sched), - (uncond_id, uncond_req, uncond_sched), - ): - diff = req.num_computed_tokens - target - if diff > 0: - req.num_computed_tokens = target - scheduler_output.num_scheduled_tokens[req_id] = orig_sched - diff - scheduler_output.total_num_scheduled_tokens -= diff - logger.debug( - "CFG equalize %s: %d -> %d scheduled, computed -> %d", - req_id, - orig_sched, - orig_sched - diff, - target, - ) - - -def _drop_split_pairs(scheduler: Any) -> None: - """Deregister a desynced CFG pair (one member preempted / left behind). - - Manual re-preemption surgery is unsafe against the upstream scheduler. - Instead, a pair that split anyway (KV-pressure preemption of exactly one - member, or unrecoverable progress skew) is dropped from the registry: - the cond request continues UNGUIDED (its logits are never blended again) - rather than being re-paired against a desynced partner, and the - companion's discarded output no longer matters. - """ - if not scheduler._cfg_pairs: - return - - running_ids = {req.request_id for req in scheduler.running} - to_drop: list[tuple[str, str, str, str]] = [] - for pair_id, roles in scheduler._cfg_pairs.items(): - cond_id = roles.get(_COND) - uncond_id = roles.get(_UNCOND) - if cond_id is None or uncond_id is None: - continue - if cond_id not in scheduler.requests or uncond_id not in scheduler.requests: - continue - - cond_running = cond_id in running_ids - uncond_running = uncond_id in running_ids - if cond_running != uncond_running: - # Split only counts once decoding started: a lone member that has - # computed tokens while its partner was pushed back is desynced. - runner = scheduler.requests.get(cond_id if cond_running else uncond_id) - if runner is not None and runner.num_computed_tokens > 0: - to_drop.append((pair_id, cond_id, uncond_id, "one member preempted")) - continue - - if cond_running and uncond_running: - cond_req = scheduler.requests.get(cond_id) - uncond_req = scheduler.requests.get(uncond_id) - if cond_req and uncond_req and cond_req.num_computed_tokens != uncond_req.num_computed_tokens: - logger.warning( - "CFG pair progress mismatch (equalize will retry): %s@%d vs %s@%d", - cond_id, - cond_req.num_computed_tokens, - uncond_id, - uncond_req.num_computed_tokens, - ) - - for pair_id, cond_id, uncond_id, reason in to_drop: - scheduler._cfg_pairs.pop(pair_id, None) - scheduler._cfg_req_to_pair.pop(cond_id, None) - scheduler._cfg_req_to_pair.pop(uncond_id, None) - logger.warning( - "CFG pair %s split (%s); releasing cond=%s unguided, uncond=%s output is discarded", - pair_id, - reason, - cond_id, - uncond_id, - ) - - -_patches_applied = False - - -def cfg_patches_applied() -> bool: - """True once :func:`apply_cfg_patches` has patched the scheduler in this process.""" - return _patches_applied - - -def _patch_scheduler() -> None: - from vllm.v1.core.sched.scheduler import Scheduler - - orig_init = Scheduler.__init__ - orig_add = Scheduler.add_request - orig_finish = Scheduler.finish_requests - orig_schedule = Scheduler.schedule - - def _init(self, *args: Any, **kwargs: Any) -> None: - orig_init(self, *args, **kwargs) - self._cfg_pairs: dict[str, dict[str, str]] = {} - self._cfg_req_to_pair: dict[str, str] = {} - - def _add(self, request: Any) -> None: - orig_add(self, request) - pair_id = _cfg_extra(request, "cfg_pair_id") - role = _cfg_extra(request, "cfg_role") - if pair_id and role: - self._cfg_pairs.setdefault(pair_id, {})[role] = request.request_id - self._cfg_req_to_pair[request.request_id] = pair_id - # An asymmetric prefix-cache hit (the cond prompt shares cached - # branches; the null prompt does not) can advance one - # member past what post-schedule equalization can claw back. - # Skip prefix-cache reads for CFG requests entirely. - if hasattr(request, "skip_reading_prefix_cache"): - request.skip_reading_prefix_cache = True - - def _finish(self, request_ids: Any, finished_status: Any) -> Any: - if request_ids is None: - result = orig_finish(self, request_ids, finished_status) - self._cfg_pairs.clear() - self._cfg_req_to_pair.clear() - return result - - if isinstance(request_ids, str): - request_ids = {request_ids} - else: - request_ids = set(request_ids) - - # Finish both members together: a surviving partner would decode - # unpaired garbage (its logits are never blended again). - partner_ids: set[str] = set() - for req_id in request_ids: - pair_id = self._cfg_req_to_pair.get(req_id) - if pair_id is None: - continue - for rid in self._cfg_pairs.get(pair_id, {}).values(): - if rid != req_id and rid in self.requests: - partner_ids.add(rid) - - all_ids = request_ids | partner_ids - result = orig_finish(self, all_ids, finished_status) - - for req_id in all_ids: - pair_id = self._cfg_req_to_pair.pop(req_id, None) - if pair_id: - self._cfg_pairs.pop(pair_id, None) - return result - - def _schedule(self, *args: Any, **kwargs: Any) -> Any: - if not self._cfg_pairs: - return orig_schedule(self, *args, **kwargs) - - _reorder_waiting_for_cfg(self) - held = _hold_incomplete_pairs(self) - - # Chunked prefill can split one pair member's prompt across steps - # while the other fits in a single step; halving the long-prefill - # threshold makes both chunk, and equalize aligns the remainder. - scheduler_config = getattr(self, "scheduler_config", None) - orig_threshold = getattr(scheduler_config, "long_prefill_token_threshold", None) - if scheduler_config is not None and orig_threshold is not None: - scheduler_config.long_prefill_token_threshold = self.max_num_scheduled_tokens // 2 - try: - scheduler_output = orig_schedule(self, *args, **kwargs) - finally: - if scheduler_config is not None and orig_threshold is not None: - scheduler_config.long_prefill_token_threshold = orig_threshold - _release_held(held) - - _equalize_cfg_pair_progress(self, scheduler_output) - _drop_split_pairs(self) - return scheduler_output - - _schedule._audex_cfg_patched = True # type: ignore[attr-defined] - - Scheduler.__init__ = _init - Scheduler.add_request = _add - Scheduler.finish_requests = _finish - Scheduler.schedule = _schedule - - -def apply_cfg_patches() -> None: - """Make the vLLM v1 scheduler CFG-pair-aware. Idempotent, per process. - - Must run in the engine-core process before ``Scheduler`` is constructed - (the ``__init__`` wrapper installs the pair registry). - """ - global _patches_applied - if _patches_applied: - return - _patches_applied = True - logger.info("Applying Audex CFG scheduler patches to vLLM v1") - _patch_scheduler() diff --git a/vllm_omni/model_executor/models/common/cfg_pairing.py b/vllm_omni/model_executor/models/common/cfg_pairing.py new file mode 100644 index 000000000..ae4b3ac38 --- /dev/null +++ b/vllm_omni/model_executor/models/common/cfg_pairing.py @@ -0,0 +1,460 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# +# Adapted from nvidia/Nemotron-Labs-Audex-2B (Apache-2.0): +# inference_scripts_vllm/audiogen_scripts/vllm_cfg_patch.py +# Divergences from the official file are commented at the divergence site. +"""Model-neutral classifier-free-guidance request pairing for the vLLM v1 scheduler. + +A guided request decodes together with an unconditional companion so their +logits can be blended every step. Blending is only correct when both members +decode the same position in the same engine step, so :func:`apply_cfg_patches` +makes the scheduler pair-aware: a lone member waits for its partner, partners +stay adjacent in the waiting queue, their ``num_computed_tokens`` are equalized +after every schedule, and they finish together. + +Requests opt in through ``SamplingParams.extra_args``:: + + cond: {"cfg_role": "cond", "cfg_pair_id": } + uncond: {"cfg_role": "uncond", "cfg_pair_id": } + +``cfg_pair_id`` is optional. A model whose serving layer cannot mint a pair id +(the companion request id is only known inside the engine) may send ``cfg_role`` +alone; the pair id is then derived from the request id, stripping the +companion's registered suffix (see :func:`register_cfg_uncond_suffix`). + +Logit blending itself is the caller's concern: this module only keeps the two +rows step-locked. The patch is a no-op for engines that never see a CFG request +(the pair registry stays empty). +""" + +from __future__ import annotations + +from collections import deque +from typing import Any + +from vllm.logger import init_logger + +logger = init_logger(__name__) + +__all__ = [ + "apply_cfg_patches", + "cfg_patches_applied", + "cfg_uncond_suffixes", + "register_cfg_uncond_suffix", +] + +_COND = "cond" +_UNCOND = "uncond" + +# Request-id suffix appended to the companion request by the Audex stage input +# processor (``CFG_UNCOND_SUFFIX`` in +# vllm_omni/model_executor/stage_input_processors/audex.py). Models that mint +# the companion request id differently register their own suffix. +_UNCOND_REQUEST_ID_SUFFIXES: set[str] = {"__cfg_uncond"} + + +def register_cfg_uncond_suffix(suffix: str) -> None: + """Register a companion request-id suffix used to derive a missing pair id.""" + if not suffix or not isinstance(suffix, str): + raise ValueError(f"CFG uncond request-id suffix must be a non-empty string, got {suffix!r}") + _UNCOND_REQUEST_ID_SUFFIXES.add(suffix) + + +def cfg_uncond_suffixes() -> frozenset[str]: + """Companion request-id suffixes registered in this process.""" + return frozenset(_UNCOND_REQUEST_ID_SUFFIXES) + + +def _derive_pair_id(request_id: str, role: str) -> str: + """Pair id for a request that carries ``cfg_role`` but no ``cfg_pair_id``. + + The cond request id IS the pair id; its companion carries the same id plus + a registered suffix, so stripping the suffix recovers it. Longest suffix + first, so a registered suffix extending another still resolves. + """ + if role == _UNCOND: + for suffix in sorted(_UNCOND_REQUEST_ID_SUFFIXES, key=len, reverse=True): + if request_id.endswith(suffix): + return request_id[: -len(suffix)] + return request_id + + +def _cfg_extra(request: Any, key: str) -> Any: + sampling_params = getattr(request, "sampling_params", None) + extra_args = getattr(sampling_params, "extra_args", None) + if extra_args: + return extra_args.get(key) + return None + + +def _wait_queues(scheduler: Any) -> list[Any]: + # Divergence from the official (vLLM 0.20-era) file: vLLM 0.24 parks + # structurally not-ready requests in a second ``skipped_waiting`` queue; + # pair-hold must cover both queues or a lone member parked there would + # be admitted without its partner. + queues = [scheduler.waiting] + skipped = getattr(scheduler, "skipped_waiting", None) + if skipped is not None: + queues.append(skipped) + return queues + + +def _pair_complete(scheduler: Any, request_id: str, blocked_ids: set[str] | None = None) -> bool: + pair_id = scheduler._cfg_req_to_pair.get(request_id) + if pair_id is None: + return True + roles = scheduler._cfg_pairs.get(pair_id, {}) + if len(roles) != 2 or not all(rid in scheduler.requests for rid in roles.values()): + return False + if blocked_ids: + # A partner parked in skipped_waiting (e.g. waiting for remote KVs) + # may fail promotion this step while this member gets scheduled from + # waiting; hold this member until the partner is schedulable. + for rid in roles.values(): + if rid != request_id and rid in blocked_ids: + return False + return True + + +# Schedule steps a lone pair member may wait for its partner. Partners arrive +# within a handful of steps in practice; exhausting this budget means the +# companion was never created (e.g. prompt expansion failed), and the request +# is released to decode unguided instead of hanging forever. +_MAX_PAIR_HOLD_STEPS = 512 + + +def _drop_stale_pair(scheduler: Any, request_id: str) -> None: + pair_id = scheduler._cfg_req_to_pair.get(request_id) + if pair_id is None: + return + for rid in scheduler._cfg_pairs.pop(pair_id, {}).values(): + scheduler._cfg_req_to_pair.pop(rid, None) + logger.warning( + "CFG pair %s: partner of %s never arrived after %d schedule steps; releasing the request unguided", + pair_id, + request_id, + _MAX_PAIR_HOLD_STEPS, + ) + + +def _hold_incomplete_pairs(scheduler: Any) -> list[tuple[Any, Any]]: + """Pull CFG requests whose partner is not yet admitted out of the wait queues. + + The engine schedules as soon as the first member of a pair arrives over + IPC; without the hold, the lone member prefills one step early and the + pair is permanently offset by one token. Held requests are re-prepended + after the schedule step. + """ + hold_counts: dict[str, int] = getattr(scheduler, "_cfg_hold_counts", None) or {} + scheduler._cfg_hold_counts = hold_counts + + skipped = getattr(scheduler, "skipped_waiting", None) + blocked_ids = {req.request_id for req in list(skipped)} if skipped is not None else set() + + held: list[tuple[Any, Any]] = [] + for queue in _wait_queues(scheduler): + # A member is only "blocked" for its partner's sake when it sits in + # the OTHER queue; members of the queue being scanned move together. + partner_blockers = blocked_ids if queue is not skipped else set() + held_here = [] + for request in list(queue): + if _pair_complete(scheduler, request.request_id, partner_blockers): + hold_counts.pop(request.request_id, None) + continue + count = hold_counts.get(request.request_id, 0) + 1 + if count > _MAX_PAIR_HOLD_STEPS: + hold_counts.pop(request.request_id, None) + _drop_stale_pair(scheduler, request.request_id) + continue + hold_counts[request.request_id] = count + held_here.append(request) + if held_here: + queue.remove_requests(held_here) + held.extend((queue, req) for req in held_here) + return held + + +def _release_held(held: list[tuple[Any, Any]]) -> None: + for queue, request in reversed(held): + queue.prepend_request(request) + + +def _reorder_waiting_for_cfg(scheduler: Any) -> None: + """Move CFG pair partners adjacent in the FCFS waiting queue.""" + waiting = scheduler.waiting + # Only the FCFS queue (a deque) has caller-controlled ordering; priority + # queues order themselves and pair sync then relies on hold + equalize. + if not isinstance(waiting, deque) or len(waiting) < 2: + return + + requests = list(waiting) + waiting.clear() + + seen: set[str] = set() + result: list[Any] = [] + for request in requests: + rid = request.request_id + if rid in seen: + continue + seen.add(rid) + result.append(request) + + pair_id = scheduler._cfg_req_to_pair.get(rid) + if pair_id is None: + continue + for partner_id in scheduler._cfg_pairs.get(pair_id, {}).values(): + if partner_id != rid and partner_id not in seen: + for candidate in requests: + if candidate.request_id == partner_id: + seen.add(partner_id) + result.append(candidate) + break + + waiting.extend(result) + + +def _equalize_cfg_pair_progress(scheduler: Any, scheduler_output: Any) -> None: + """Bring both members of every CFG pair to the same ``num_computed_tokens``. + + Prefix-cache hits or unequal chunked-prefill budget can leave one member + ahead after a schedule step; reduce the faster member's allocation so + both land on the same position. Over-allocated KV blocks remain and are + consumed next step, so nothing leaks. + """ + for roles in scheduler._cfg_pairs.values(): + cond_id = roles.get(_COND) + uncond_id = roles.get(_UNCOND) + if not cond_id or not uncond_id: + continue + + cond_sched = scheduler_output.num_scheduled_tokens.get(cond_id, 0) + uncond_sched = scheduler_output.num_scheduled_tokens.get(uncond_id, 0) + if cond_sched == 0 or uncond_sched == 0: + continue + + cond_req = scheduler.requests.get(cond_id) + uncond_req = scheduler.requests.get(uncond_id) + if not cond_req or not uncond_req: + continue + + if cond_req.num_computed_tokens == uncond_req.num_computed_tokens: + continue + + target = min(cond_req.num_computed_tokens, uncond_req.num_computed_tokens) + + feasible = all( + req.num_computed_tokens - target < sched + for req, sched in ((cond_req, cond_sched), (uncond_req, uncond_sched)) + ) + if not feasible: + continue + + for req_id, req, orig_sched in ( + (cond_id, cond_req, cond_sched), + (uncond_id, uncond_req, uncond_sched), + ): + diff = req.num_computed_tokens - target + if diff > 0: + req.num_computed_tokens = target + scheduler_output.num_scheduled_tokens[req_id] = orig_sched - diff + scheduler_output.total_num_scheduled_tokens -= diff + logger.debug( + "CFG equalize %s: %d -> %d scheduled, computed -> %d", + req_id, + orig_sched, + orig_sched - diff, + target, + ) + + +def _drop_split_pairs(scheduler: Any) -> None: + """Deregister a desynced CFG pair (one member preempted / left behind). + + Manual re-preemption surgery is unsafe against the upstream scheduler. + Instead, a pair that split anyway (KV-pressure preemption of exactly one + member, or unrecoverable progress skew) is dropped from the registry: + the cond request continues UNGUIDED (its logits are never blended again) + rather than being re-paired against a desynced partner, and the + companion's discarded output no longer matters. + """ + if not scheduler._cfg_pairs: + return + + running_ids = {req.request_id for req in scheduler.running} + to_drop: list[tuple[str, str, str, str]] = [] + for pair_id, roles in scheduler._cfg_pairs.items(): + cond_id = roles.get(_COND) + uncond_id = roles.get(_UNCOND) + if cond_id is None or uncond_id is None: + continue + if cond_id not in scheduler.requests or uncond_id not in scheduler.requests: + continue + + cond_running = cond_id in running_ids + uncond_running = uncond_id in running_ids + if cond_running != uncond_running: + # Split only counts once decoding started: a lone member that has + # computed tokens while its partner was pushed back is desynced. + runner = scheduler.requests.get(cond_id if cond_running else uncond_id) + if runner is not None and runner.num_computed_tokens > 0: + to_drop.append((pair_id, cond_id, uncond_id, "one member preempted")) + continue + + if cond_running and uncond_running: + cond_req = scheduler.requests.get(cond_id) + uncond_req = scheduler.requests.get(uncond_id) + if cond_req and uncond_req and cond_req.num_computed_tokens != uncond_req.num_computed_tokens: + logger.warning( + "CFG pair progress mismatch (equalize will retry): %s@%d vs %s@%d", + cond_id, + cond_req.num_computed_tokens, + uncond_id, + uncond_req.num_computed_tokens, + ) + + for pair_id, cond_id, uncond_id, reason in to_drop: + scheduler._cfg_pairs.pop(pair_id, None) + scheduler._cfg_req_to_pair.pop(cond_id, None) + scheduler._cfg_req_to_pair.pop(uncond_id, None) + logger.warning( + "CFG pair %s split (%s); releasing cond=%s unguided, uncond=%s output is discarded", + pair_id, + reason, + cond_id, + uncond_id, + ) + + +_patches_applied = False + + +def cfg_patches_applied() -> bool: + """True once :func:`apply_cfg_patches` has patched the scheduler in this process.""" + return _patches_applied + + +def _patch_scheduler(scheduler_cls: type | None = None) -> None: + """Wrap the scheduler methods that keep a CFG pair step-locked. + + ``scheduler_cls`` defaults to vLLM's v1 ``Scheduler``; accepting a class + keeps the wrappers exercisable without building a real engine. + """ + if scheduler_cls is None: + from vllm.v1.core.sched.scheduler import Scheduler + + scheduler_cls = Scheduler + + orig_init = scheduler_cls.__init__ + orig_add = scheduler_cls.add_request + orig_finish = scheduler_cls.finish_requests + orig_schedule = scheduler_cls.schedule + + def _init(self, *args: Any, **kwargs: Any) -> None: + orig_init(self, *args, **kwargs) + self._cfg_pairs: dict[str, dict[str, str]] = {} + self._cfg_req_to_pair: dict[str, str] = {} + + def _add(self, request: Any) -> None: + orig_add(self, request) + pair_id = _cfg_extra(request, "cfg_pair_id") + role = _cfg_extra(request, "cfg_role") + if role and not pair_id: + # Models that cannot inject a pair id from the serving layer send + # the role alone and let the request id carry the pairing. It has + # to be the user-facing id: the engine renames every request to an + # internal UUID on admission, and the companion suffix survives + # only on ``external_req_id``. Deriving from the UUID would give + # the two members unrelated pair ids, so they would never be held + # for each other and could be scheduled apart. + external_id = getattr(request, "external_req_id", None) or request.request_id + pair_id = _derive_pair_id(str(external_id), role) + if pair_id and role: + self._cfg_pairs.setdefault(pair_id, {})[role] = request.request_id + self._cfg_req_to_pair[request.request_id] = pair_id + # An asymmetric prefix-cache hit (the cond prompt shares cached + # branches; the null prompt does not) can advance one member past + # what post-schedule equalization can claw back. Skip prefix-cache + # reads for CFG requests entirely. + if hasattr(request, "skip_reading_prefix_cache"): + request.skip_reading_prefix_cache = True + + def _finish(self, request_ids: Any, finished_status: Any) -> Any: + if request_ids is None: + result = orig_finish(self, request_ids, finished_status) + self._cfg_pairs.clear() + self._cfg_req_to_pair.clear() + return result + + if isinstance(request_ids, str): + request_ids = {request_ids} + else: + request_ids = set(request_ids) + + # Finish both members together: a surviving partner would decode + # unpaired garbage (its logits are never blended again). + partner_ids: set[str] = set() + for req_id in request_ids: + pair_id = self._cfg_req_to_pair.get(req_id) + if pair_id is None: + continue + for rid in self._cfg_pairs.get(pair_id, {}).values(): + if rid != req_id and rid in self.requests: + partner_ids.add(rid) + + all_ids = request_ids | partner_ids + result = orig_finish(self, all_ids, finished_status) + + for req_id in all_ids: + pair_id = self._cfg_req_to_pair.pop(req_id, None) + if pair_id: + self._cfg_pairs.pop(pair_id, None) + return result + + def _schedule(self, *args: Any, **kwargs: Any) -> Any: + if not self._cfg_pairs: + return orig_schedule(self, *args, **kwargs) + + _reorder_waiting_for_cfg(self) + held = _hold_incomplete_pairs(self) + + # Chunked prefill can split one pair member's prompt across steps + # while the other fits in a single step; halving the long-prefill + # threshold makes both chunk, and equalize aligns the remainder. + scheduler_config = getattr(self, "scheduler_config", None) + orig_threshold = getattr(scheduler_config, "long_prefill_token_threshold", None) + if scheduler_config is not None and orig_threshold is not None: + scheduler_config.long_prefill_token_threshold = self.max_num_scheduled_tokens // 2 + try: + scheduler_output = orig_schedule(self, *args, **kwargs) + finally: + if scheduler_config is not None and orig_threshold is not None: + scheduler_config.long_prefill_token_threshold = orig_threshold + _release_held(held) + + _equalize_cfg_pair_progress(self, scheduler_output) + _drop_split_pairs(self) + return scheduler_output + + _schedule._cfg_pairing_patched = True # type: ignore[attr-defined] + # Legacy marker name kept for callers written against the Audex-only patch. + _schedule._audex_cfg_patched = True # type: ignore[attr-defined] + + scheduler_cls.__init__ = _init + scheduler_cls.add_request = _add + scheduler_cls.finish_requests = _finish + scheduler_cls.schedule = _schedule + + +def apply_cfg_patches() -> None: + """Make the vLLM v1 scheduler CFG-pair-aware. Idempotent, per process. + + Must run in the engine-core process before ``Scheduler`` is constructed + (the ``__init__`` wrapper installs the pair registry). + """ + global _patches_applied + if _patches_applied: + return + _patches_applied = True + logger.info("Applying CFG pairing scheduler patches to vLLM v1") + _patch_scheduler() diff --git a/vllm_omni/model_executor/models/minimax_music3/__init__.py b/vllm_omni/model_executor/models/minimax_music3/__init__.py new file mode 100644 index 000000000..39e617d04 --- /dev/null +++ b/vllm_omni/model_executor/models/minimax_music3/__init__.py @@ -0,0 +1,34 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""MiniMax Music 3 text-to-music. + +Stage 0 is a Qwen3 backbone with an eight-codebook RVQ frame; stage 1 turns +the resulting conditioning frames into 32 kHz stereo audio through a +flow-matching transformer and a DAC decoder. +""" + +from .constants import ( + AR_CHUNK_FRAMES, + AR_CHUNK_HOP_FRAMES, + AR_HIDDEN_SIZE, + MAX_AUDIO_FRAMES, + OUTPUT_SAMPLE_RATE, +) +from .prompt import ( + AUDIO_CODE_OFFSET, + SPECIAL_TOKEN_IDS, + build_cfg_null_token_ids, + build_prompt, +) + +__all__ = [ + "AR_CHUNK_FRAMES", + "AR_CHUNK_HOP_FRAMES", + "AR_HIDDEN_SIZE", + "AUDIO_CODE_OFFSET", + "MAX_AUDIO_FRAMES", + "OUTPUT_SAMPLE_RATE", + "SPECIAL_TOKEN_IDS", + "build_cfg_null_token_ids", + "build_prompt", +] diff --git a/vllm_omni/model_executor/models/minimax_music3/acoustic.py b/vllm_omni/model_executor/models/minimax_music3/acoustic.py new file mode 100644 index 000000000..879185a57 --- /dev/null +++ b/vllm_omni/model_executor/models/minimax_music3/acoustic.py @@ -0,0 +1,459 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Stage-1 acoustic decoder for MiniMax Music 3. + +Consumes the AR stage's conditioning spans and returns 32 kHz stereo audio. +Each span is cut into 200-frame windows that overlap by half; every window is +solved by the flow-matching DiT with the previous window's tail latent pinned +as a prompt, decoded by the vocoder, and cropped so that concatenating the +survivors reconstructs the song exactly once. + +The window arithmetic lives in :mod:`.chunking` and is driven entirely by the +span's global frame range, so the streaming path (one window per engine step) +and the non-streaming path (the whole song in one payload) produce identical +audio for the same seed. + +The stage runs in float32. TF32 keeps that affordable, and bfloat16 measurably +degrades a 30-step solve. +""" + +from __future__ import annotations + +import hashlib +from collections.abc import Iterable +from dataclasses import dataclass, field +from typing import Any + +import torch +from torch import Tensor, nn +from vllm.config import VllmConfig +from vllm.logger import init_logger + +from vllm_omni.model_executor.models.output_templates import OmniOutput + +from .chunking import ChunkWindow, chunk_windows, crop_sample_bounds, overlap_mel_length +from .constants import ( + AR_CHUNK_HOP_FRAMES, + AR_HIDDEN_SIZE, + DEFAULT_DIT_CFG_SCALE, + DEFAULT_DIT_STEPS, + OUTPUT_CHANNELS, + OUTPUT_SAMPLE_RATE, +) +from .dav import MiniMaxMusic3Vocoder, StreamingResampler, remove_weight_norm +from .dit import MiniMaxMusic3FlowMatchingDiT, remap_transformer_state +from .weights import ( + CONDITION_ENCODER_DIR, + TRANSFORMER_DIR, + VOCODER_DIR, + load_component_state, + resolve_repo_root, +) + +logger = init_logger(__name__) + +__all__ = ["MiniMaxMusic3AcousticForConditionalGeneration"] + +# Namespaced so a request seed reused by another component cannot collide with +# the noise draw, and so the draw is reproducible across processes. Python's +# hash() is randomized per interpreter and must not be used here. +_SEED_PERSON = b"minimax-ttm" +_SEED_BYTES = 8 + + +def _derive_seed(seed: int, *parts: object) -> int: + """Derive a stable 63-bit sub-seed from a request seed and some labels.""" + digest = hashlib.blake2b(digest_size=_SEED_BYTES, person=_SEED_PERSON) + digest.update(int(seed).to_bytes(8, "little", signed=False)) + for part in parts: + raw = str(part).encode("utf-8") + digest.update(len(raw).to_bytes(4, "little")) + digest.update(raw) + return int.from_bytes(digest.digest(), "little") & ((1 << 63) - 1) + + +def _meta_int(value: Any, default: int = 0) -> int: + """Read an int that transport may have wrapped in a tensor or a list.""" + if isinstance(value, torch.Tensor): + return int(value.reshape(-1)[0].item()) if value.numel() else default + if isinstance(value, list | tuple): + return _meta_int(value[0], default) if value else default + try: + return int(value) # type: ignore[arg-type] + except (TypeError, ValueError): + return default + + +def _meta_bool(value: Any) -> bool: + """Read a bool that transport may have wrapped in a tensor or a list.""" + if isinstance(value, torch.Tensor): + return bool(value.reshape(-1)[0].item()) if value.numel() else False + if isinstance(value, list | tuple): + return _meta_bool(value[0]) if value else False + return bool(value) + + +def _meta_str(value: Any) -> str | None: + """Read a string that transport may have wrapped in a list.""" + if isinstance(value, list | tuple): + return _meta_str(value[0]) if value else None + return value if isinstance(value, str) else None + + +@dataclass +class _RequestState: + """Cross-window streaming state, keyed by request id and never by row.""" + + last_latent: Tensor | None = None + last_condition: Tensor | None = None + next_window_index: int = 0 + decoded_windows: int = 0 + resampler: StreamingResampler = field(default_factory=StreamingResampler) + + def release(self) -> None: + self.last_latent = None + self.last_condition = None + self.resampler.reset() + + +class MiniMaxMusic3AcousticForConditionalGeneration(nn.Module): + """Flow-matching DiT plus DAC vocoder, driven by the generation runner.""" + + input_modalities = "audio" + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + del prefix + self.vllm_config = vllm_config + self.model_path = vllm_config.model_config.model + + self.have_multimodal_outputs = True + self.has_preprocess = False + self.has_postprocess = False + self.enable_update_additional_information = True + self.requires_raw_input_tokens = True + + self._configure_backends() + + # Skeleton only. vLLM's memory profiler walks the module tree at + # startup, before any weight loader has run. + self.dit = MiniMaxMusic3FlowMatchingDiT() + self.vocoder = MiniMaxMusic3Vocoder() + + extra = self._connector_extra_config() + self.dit_steps = max(1, _meta_int(extra.get("dit_steps"), DEFAULT_DIT_STEPS)) + cfg_scale = extra.get("dit_cfg_scale") + self.dit_cfg_scale = ( + float(cfg_scale) + if isinstance(cfg_scale, int | float) and not isinstance(cfg_scale, bool) + else DEFAULT_DIT_CFG_SCALE + ) + + self._states: dict[str, _RequestState] = {} + self._loaded = False + + @staticmethod + def _configure_backends() -> None: + """Pin the numeric backends this checkpoint was validated against. + + These are process-global switches, which is safe here only because a + stage owns its own engine-core process. Never call this from a module + shared with the AR stage: cuDNN's fused attention is not + bit-reproducible for grouped-query shapes, so flipping it underneath + the backbone makes the same seed produce a different song depending on + which stage loaded first. + + cuDNN itself is off because its 1-D convolution paths are slower than + the native ones for the vocoder's dilated stacks, and its SDPA backend + is not needed because the DiT runs the flash path. TF32 keeps + float32's range with a 10-bit mantissa, which is what makes a float32 + solve affordable. + """ + torch.backends.cuda.matmul.allow_tf32 = True + torch.set_float32_matmul_precision("high") + if torch.cuda.is_available(): + torch.backends.cudnn.enabled = False + torch.backends.cuda.enable_cudnn_sdp(False) + + def _connector_extra_config(self) -> dict[str, Any]: + """Return the stage connector's ``extra`` block, or an empty mapping.""" + connector = getattr(self.vllm_config.model_config, "stage_connector_config", None) + if isinstance(connector, dict): + extra = connector.get("extra", connector) + return extra if isinstance(extra, dict) else {} + extra = getattr(connector, "extra", None) + return extra if isinstance(extra, dict) else {} + + # ------------------------------------------------------------- vLLM hooks + + def embed_input_ids(self, input_ids: Tensor, **_: Any) -> Tensor: + """Stable dummy embedding: this stage never reads token embeddings.""" + if input_ids.numel() == 0: + return torch.empty((0, 1), device=input_ids.device, dtype=torch.float32) + return torch.zeros((input_ids.shape[0], 1), device=input_ids.device, dtype=torch.float32) + + def compute_logits(self, hidden_states: Tensor | OmniOutput, sampling_metadata: Any = None) -> None: + del hidden_states, sampling_metadata + return None + + def load_weights(self, weights: Iterable[tuple[str, Tensor]]) -> set[str]: + """Load the three acoustic components from their own folders. + + The iterator vLLM hands over covers the stage's own checkpoint folder, + which holds nothing this stage needs. It is drained anyway: the loader + streams from a file handle the caller closes only once the iterator is + exhausted, so abandoning it wedges startup. + """ + try: + for _ in weights: + pass + except RuntimeError as exc: + # The iterator is lazy, so draining it is what actually asks the + # default loader for root-level weight files. This repository keeps + # every component in its own subfolder and has none at the root, so + # that lookup raises. Nothing here needs those weights; any other + # RuntimeError is a real problem and is re-raised. + if "Cannot find any model weights" not in str(exc): + raise + logger.debug("MiniMax Music 3 acoustic stage has no root-level weights, as expected") + return self._load_components() + + def _ensure_loaded(self) -> None: + """Load on first use when the weight loader skipped ``load_weights``. + + ``load_format=dummy`` (used by warmup and memory profiling) never calls + ``load_weights``, which would leave the modules holding their random + skeleton and make the dummy run produce noise or crash. + """ + if not self._loaded: + self._load_components() + + def _load_components(self) -> set[str]: + """Read ``condition_encoder/``, ``transformer/`` and ``vocoder/``. + + Every component loads with ``strict=True``. A silently partial load + would produce audio that sounds plausible but is wrong, which is far + worse than a startup failure. + """ + root = resolve_repo_root(self.model_path) + encoder = self.dit.condition_encoder + encoder.load_state_dict(load_component_state(root, CONDITION_ENCODER_DIR), strict=True) + transformer_state = remap_transformer_state(load_component_state(root, TRANSFORMER_DIR)) + self.dit.transformer.load_state_dict(transformer_state, strict=True) + self.vocoder.load_state_dict(load_component_state(root, VOCODER_DIR), strict=True) + + device = self.vllm_config.device_config.device + self.dit.to(device=device, dtype=torch.float32).eval() + self.vocoder.to(device=device, dtype=torch.float32).eval() + folded = remove_weight_norm(self.vocoder) + self._loaded = True + + loaded = {f"dit.{name}" for name, _ in self.dit.named_parameters()} + loaded |= {f"vocoder.{name}" for name, _ in self.vocoder.named_parameters()} + logger.info( + "MiniMax Music 3 acoustic loaded from %s device=%s dit_steps=%d " + "dit_cfg_scale=%.3f attention=%s folded_weight_norms=%d parameters=%d", + root, + device, + self.dit_steps, + self.dit_cfg_scale, + self.dit.transformer.attention_backend_name, + folded, + sum(p.numel() for p in self.parameters()), + ) + return loaded + + def on_requests_finished(self, finished_req_ids: Iterable[str]) -> None: + """Drop the streaming state of finished or aborted requests.""" + for req_id in finished_req_ids: + state = self._states.pop(str(req_id), None) + if state is not None: + logger.info( + "MiniMax Music 3 acoustic done request=%s windows=%d", + req_id, + state.decoded_windows, + ) + state.release() + + def make_omni_output(self, model_outputs: Tensor | OmniOutput | tuple, **kwargs: Any) -> OmniOutput: + del kwargs + if isinstance(model_outputs, OmniOutput): + return model_outputs + # The runner may unpack the NamedTuple in transit; rebuild it. + if isinstance(model_outputs, tuple) and len(model_outputs) == len(OmniOutput._fields): + return OmniOutput(*model_outputs) + raise TypeError(f"MiniMaxMusic3AcousticForConditionalGeneration expected OmniOutput, got {type(model_outputs)}") + + # --------------------------------------------------------------- decoding + + @staticmethod + def _silence() -> Tensor: + """Placeholder for a step that produced no audio. + + Deliberately rank 2. A rank-1 empty tensor mixed with rank-2 chunks + makes the multimodal collector's ``torch.cat(..., dim=-1)`` fall back + to a flattening path that would planar-ize the stereo. + """ + return torch.zeros((OUTPUT_CHANNELS, 0), dtype=torch.float32) + + def _decode_window(self, state: _RequestState, hidden: Tensor, window: ChunkWindow, seed: int) -> Tensor: + """Solve and vocode one window, cropped to its unique contribution. + + Everything stays on the device: the crop is a slice and the result is + handed straight to the resampler, so a step costs one device-to-host + copy however many windows it decoded. + """ + device = self.vllm_config.device_config.device + hidden = hidden.unsqueeze(0).to(device=device, dtype=torch.float32) + generator = torch.Generator(device=device).manual_seed(_derive_seed(seed, "dit", window.index)) + align = self.dit.aligned_condition(hidden) + latent = self.dit( + align, + generator=generator, + initial_latent=state.last_latent, + initial_condition=state.last_condition, + num_steps=self.dit_steps, + cfg_scale=self.dit_cfg_scale, + ) + waveform = self.vocoder(latent) + + # Carry the middle of the overlap forward. The solver re-imposes it on + # the next window, so the two windows agree where they meet. + overlap = overlap_mel_length() + start = max(0, latent.shape[-1] - 2 * overlap) + end = max(start, latent.shape[-1] - overlap) + state.last_latent = latent[..., start:end].clone() + state.last_condition = align[..., start:end].clone() + + wave = waveform[0].clamp(-1.0, 1.0).float() + left, right = crop_sample_bounds(window) + stop = wave.shape[-1] - right if right else wave.shape[-1] + return wave[:, left:stop] + + def _decode_span(self, req_id: str, span: Tensor | None, meta: dict[str, Any]) -> Tensor: + """Decode every window a span covers and return this step's audio. + + A span is normally one window, but the terminal emission carries the + held-back complete window plus the truncated final one, and the + non-streaming path carries the whole song. Window boundaries and the + crop are re-derived from the span's global frame range so all three + cases go through one code path. + + Raises: + ValueError: If the span is malformed or lands on a window other + than the one this request is waiting for. + """ + span_start = _meta_int(meta.get("num_processed_tokens"), 0) + is_terminal = _meta_bool(meta.get("last_chunk")) or _meta_bool(meta.get("stream_finished")) + seed = _meta_int(meta.get("audio_seed"), 0) + state = self._states.setdefault(req_id, _RequestState()) + + # A metadata-only terminal marker carries no conditioning; it decodes + # zero windows and exists only to flush the resampler. + windows = chunk_windows(int(span.shape[0])) if span is not None else [] + pieces: list[Tensor] = [] + with torch.inference_mode(): + for local in windows: + start = span_start + local.start + index = start // AR_CHUNK_HOP_FRAMES + if index != state.next_window_index: + # Each window is prompted by its predecessor's tail, so a + # gap or a replay corrupts the join. Fail the request here + # rather than return audio with a seam in it. + raise ValueError( + f"MiniMax Music 3 acoustic request {req_id} expected window " + f"{state.next_window_index} but received window {index} " + f"(span starts at AR frame {span_start}); the overlap prompt " + f"for this window is stale and the audio would not join" + ) + window = ChunkWindow( + index=index, + start=start, + end=span_start + local.end, + is_first=start == 0, + is_last=local.is_last and is_terminal, + ) + pieces.append(self._decode_window(state, span[local.start : local.end], window, seed)) + state.next_window_index = index + 1 + state.decoded_windows += 1 + + emitted = state.resampler.push(pieces, flush=is_terminal) + return emitted.cpu() + + def _span_from_info(self, info: dict[str, Any], req_id: str) -> Tensor | None: + """Return the conditioning span as ``[T, 32768]``, or ``None``. + + ``None`` means this step carried no conditioning, which is normal: a + request can be scheduled before its span arrives, and the terminal + marker is metadata only. A span that is present but the wrong shape is + not normal and stops the request. + + Raises: + ValueError: If a span is present but is not ``[T, 32768]``. + """ + latent = info.get("latent") + if not isinstance(latent, Tensor) or latent.numel() == 0: + return None + if latent.ndim == 3 and latent.shape[0] == 1: + latent = latent[0] + if latent.ndim != 2 or latent.shape[1] != AR_HIDDEN_SIZE: + raise ValueError( + f"MiniMax Music 3 acoustic request {req_id} carries a conditioning " + f"span of shape {tuple(latent.shape)}; expected [T,{AR_HIDDEN_SIZE}]" + ) + return latent + + def forward( + self, + input_ids: Tensor | None = None, + positions: Tensor | None = None, + intermediate_tensors: Any = None, + inputs_embeds: Tensor | None = None, + runtime_additional_information: list[dict[str, Any]] | None = None, + **kwargs: Any, + ) -> OmniOutput: + """Decode this step's spans into one stereo waveform per request. + + Raises: + ValueError: If a payload is malformed. Failing the step is + deliberate: a request whose conditioning cannot be decoded has + no audio to return, and reporting an error is better than + leaving the caller waiting on a stream that will never end. + """ + del input_ids, positions, intermediate_tensors, inputs_embeds + self._ensure_loaded() + + infos = runtime_additional_information + if infos is None: + infos = kwargs.get("model_intermediate_buffer") + infos = infos or [] + + # One entry per scheduled request, including steps where a request's + # span has not arrived yet. + seq_token_counts = kwargs.get("seq_token_counts") + num_reqs = len(seq_token_counts) if seq_token_counts is not None else len(infos) + + audios: list[Tensor] = [] + sr_tensor = torch.tensor(OUTPUT_SAMPLE_RATE, dtype=torch.int32) + for index in range(num_reqs): + info = infos[index] if index < len(infos) else {} + info = info if isinstance(info, dict) else {} + meta = info.get("meta") + meta = meta if isinstance(meta, dict) else {} + req_id = _meta_str(meta.get("req_id")) or _meta_str(info.get("req_id")) + has_span = isinstance(info.get("latent"), Tensor) and info["latent"].numel() > 0 + is_terminal = _meta_bool(meta.get("last_chunk")) or _meta_bool(meta.get("stream_finished")) + if not has_span and not is_terminal: + audios.append(self._silence()) + continue + if req_id is None: + raise ValueError( + "MiniMax Music 3 acoustic payload carries no meta.req_id, so its " + "window cannot be joined to the rest of its request" + ) + audios.append(self._decode_span(req_id, self._span_from_info(info, req_id), meta)) + + return OmniOutput( + text_hidden_states=None, + multimodal_outputs={"model_outputs": audios, "sr": [sr_tensor] * num_reqs}, + ) diff --git a/vllm_omni/model_executor/models/minimax_music3/chunking.py b/vllm_omni/model_executor/models/minimax_music3/chunking.py new file mode 100644 index 000000000..ba917f39d --- /dev/null +++ b/vllm_omni/model_executor/models/minimax_music3/chunking.py @@ -0,0 +1,96 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Chunk boundary and overlap arithmetic for MiniMax Music 3. + +The AR stage emits fixed 200-frame windows that overlap by half. The acoustic +stage decodes each window with the previous window's latent as a prompt, then +crops the overlap back off so the concatenated waveform is seamless. All of +that is pure integer arithmetic and lives here so it can be unit tested +without a GPU. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +from .constants import ( + AR_CHUNK_FRAMES, + AR_CHUNK_HOP_FRAMES, + BLEND_MEL_FRAMES, + DAV_HOP_SAMPLES, + HOP_MEL_FRAMES, +) + + +@dataclass(frozen=True) +class ChunkWindow: + """One acoustic window, in AR frame coordinates.""" + + index: int + start: int + end: int + is_first: bool + is_last: bool + + @property + def length(self) -> int: + return self.end - self.start + + +def chunk_windows(frames: int) -> list[ChunkWindow]: + """Return the fixed-size windows covering ``frames``, overlapping by half. + + A song shorter than one window produces a single short window; otherwise + windows advance by ``AR_CHUNK_HOP_FRAMES`` and the final one is truncated + at the end of the song. + """ + frames = int(frames) + if frames <= 0: + return [] + if frames <= AR_CHUNK_FRAMES: + return [ChunkWindow(0, 0, frames, True, True)] + + windows: list[ChunkWindow] = [] + index = 0 + start = 0 + while start < frames: + end = min(start + AR_CHUNK_FRAMES, frames) + windows.append( + ChunkWindow( + index=index, + start=start, + end=end, + is_first=index == 0, + is_last=end >= frames, + ) + ) + if end >= frames: + break + index += 1 + start += AR_CHUNK_HOP_FRAMES + return windows + + +def overlap_mel_length() -> int: + """Latent frames carried forward as the next window's flow-matching prompt.""" + return HOP_MEL_FRAMES // 2 + + +def crop_sample_bounds(window: ChunkWindow) -> tuple[int, int]: + """Return ``(left, right)`` waveform samples to trim from a decoded window. + + The first window keeps its head and the last keeps its tail; every other + window is cropped on both sides so that concatenating the survivors + reconstructs the song exactly once. + """ + left = 0 if window.is_first else BLEND_MEL_FRAMES * DAV_HOP_SAMPLES + right = 0 if window.is_last else (HOP_MEL_FRAMES - BLEND_MEL_FRAMES) * DAV_HOP_SAMPLES + return left, right + + +__all__ = [ + "ChunkWindow", + "chunk_windows", + "crop_sample_bounds", + "overlap_mel_length", +] diff --git a/vllm_omni/model_executor/models/minimax_music3/constants.py b/vllm_omni/model_executor/models/minimax_music3/constants.py new file mode 100644 index 000000000..4e93139c2 --- /dev/null +++ b/vllm_omni/model_executor/models/minimax_music3/constants.py @@ -0,0 +1,111 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inference-contract constants for MiniMax Music 3. + +These values are fixed by the checkpoint, not tuning knobs. Changing any of +them changes the audio, so they are declared once here and imported +everywhere rather than re-derived per call site. +""" + +# --- Generation limits ----------------------------------------------------- + +# 25 audio frames per second, capped at six minutes of song. +AUDIO_FRAME_RATE = 25 +MAX_AUDIO_FRAMES = 9_000 +DEFAULT_MAX_AUDIO_FRAMES = 9_000 + +# The backbone was trained with a 10,240-token context; the prompt is capped +# well below it so there is room for the audio frames that follow. +MAX_PROMPT_TOKENS = 5_000 + +# --- Frame layout ---------------------------------------------------------- + +# Eight RVQ codebooks per frame: c0 lives in the backbone vocabulary, c1..c7 +# in the depth decoder's own 1024-entry books. +NUM_CODEBOOKS = 8 +C0_VOCAB_SIZE = 16_384 +AUDIO_VOCAB_SIZE = 1_024 + +# The acoustic stage is conditioned on the backbone hidden state concatenated +# with the seven depth-decoder hidden states. +BACKBONE_HIDDEN_SIZE = 4_096 +AR_HIDDEN_SIZE = BACKBONE_HIDDEN_SIZE * NUM_CODEBOOKS # 32768 + +# --- Sampling -------------------------------------------------------------- + +# Classifier-free guidance is not optional and has no request field: this is +# how the reference implementation samples this checkpoint. +AR_CFG_SCALE = 1.5 +AR_CFG_TOP_K = 50 +AR_TOP_K = 50 + +# --- Acoustic stage -------------------------------------------------------- + +DEFAULT_DIT_STEPS = 30 +DEFAULT_DIT_CFG_SCALE = 1.7 +DIT_LATENT_CHANNELS = 128 + +# Windowing: the AR stage hands the acoustic stage 200-frame windows that +# overlap by half, and the acoustic stage crops the overlap back off. +AR_CHUNK_FRAMES = 200 +AR_CHUNK_HOP_FRAMES = 100 + +# DAV runs at 44.1 kHz; the response is resampled to 32 kHz stereo. +DAV_SAMPLE_RATE = 44_100 +OUTPUT_SAMPLE_RATE = 32_000 +OUTPUT_CHANNELS = 2 + +# Condition resampling: AR frames land on a 24 kHz / 960-hop grid and are +# stretched onto the DAV latent's 44.1 kHz / 512-hop grid. +CONDITION_INPUT_SAMPLE_RATE = 24_000 +CONDITION_INPUT_HOP = 960 +CONDITION_OUTPUT_SAMPLE_RATE = 44_100 +CONDITION_OUTPUT_HOP = 512 + +# Mel frames per AR frame: 44100/24000 * 960/512 = 3.4453125. So a 100-frame +# hop is 344 mel frames and a 200-frame window is 689. +# +# HOP_MEL_FRAMES is the hop, not the window: the crop and carry-forward +# arithmetic is expressed in units of the hop, and a window is two of them +# plus one frame of rounding. Reading this as the window length is the easy +# mistake here, and it is off by a factor of two. +HOP_MEL_FRAMES = 344 +BLEND_MEL_FRAMES = HOP_MEL_FRAMES // 4 +DAV_HOP_SAMPLES = 512 + +# Only these two dtypes are supported for the acoustic stage. float32 is the +# default: TF32 keeps it affordable, and bfloat16 measurably degrades the +# solver. +SUPPORTED_ACOUSTIC_DTYPES = frozenset({"float32", "bfloat16"}) + + +__all__ = [ + "AR_CFG_SCALE", + "AR_CFG_TOP_K", + "AR_CHUNK_FRAMES", + "AR_CHUNK_HOP_FRAMES", + "AR_HIDDEN_SIZE", + "AR_TOP_K", + "AUDIO_FRAME_RATE", + "AUDIO_VOCAB_SIZE", + "BACKBONE_HIDDEN_SIZE", + "BLEND_MEL_FRAMES", + "C0_VOCAB_SIZE", + "CONDITION_INPUT_HOP", + "CONDITION_INPUT_SAMPLE_RATE", + "CONDITION_OUTPUT_HOP", + "CONDITION_OUTPUT_SAMPLE_RATE", + "DAV_HOP_SAMPLES", + "DAV_SAMPLE_RATE", + "DEFAULT_DIT_CFG_SCALE", + "DEFAULT_DIT_STEPS", + "DEFAULT_MAX_AUDIO_FRAMES", + "DIT_LATENT_CHANNELS", + "HOP_MEL_FRAMES", + "MAX_AUDIO_FRAMES", + "MAX_PROMPT_TOKENS", + "NUM_CODEBOOKS", + "OUTPUT_CHANNELS", + "OUTPUT_SAMPLE_RATE", + "SUPPORTED_ACOUSTIC_DTYPES", +] diff --git a/vllm_omni/model_executor/models/minimax_music3/dav.py b/vllm_omni/model_executor/models/minimax_music3/dav.py new file mode 100644 index 000000000..bc0a478fd --- /dev/null +++ b/vllm_omni/model_executor/models/minimax_music3/dav.py @@ -0,0 +1,294 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""DAC-style vocoder for MiniMax Music 3. + +Only the decoder half of the autoencoder ships in the checkpoint, which is all +inference needs: the flow-matching stage produces the latent directly, so the +encoder and the quantizer never run. + +Stereo is folded into the batch. A ``[B, 128, T]`` latent is read as two +64-channel mono latents per item, decoded as ``2B`` independent mono streams, +and unfolded again at the end. The four transposed-convolution blocks upsample +by 8, 8, 4 and 2, so one latent frame becomes 512 samples at 44.1 kHz. + +Module names match the ``vocoder/`` component folder key for key, so its state +dict loads with ``strict=True`` and no remapping. The convolutions keep the +checkpoint's weight-norm parameterization (``weight_g``/``weight_v``) through +loading; :func:`remove_weight_norm` folds it away afterwards, since the +factorization is fixed once the weights stop training. + +:class:`StreamingResampler` carries the 44.1 kHz output down to the 32 kHz +response rate across chunk boundaries, which is not something a plain +per-chunk resample can do. +""" + +from __future__ import annotations + +import math +from collections.abc import Sequence + +import torch +from torch import Tensor, nn +from vllm.logger import init_logger + +from .constants import DAV_SAMPLE_RATE, DIT_LATENT_CHANNELS, OUTPUT_CHANNELS, OUTPUT_SAMPLE_RATE + +logger = init_logger(__name__) + +__all__ = ["MiniMaxMusic3Vocoder", "StreamingResampler", "remove_weight_norm", "snake"] + +# Fixed by vocoder/config.json. +_DECODER_INPUT_DIM = 1_024 +_DECODER_HIDDEN_DIM = 1_536 +_UPSAMPLING_RATIOS = (8, 8, 4, 2) +_RESIDUAL_DILATIONS = (1, 3, 9) +_KERNEL_SIZE = 7 +# Each vocoder stream sees half the latent channels; the other half is the +# opposite stereo channel. +_STREAM_CHANNELS = DIT_LATENT_CHANNELS // OUTPUT_CHANNELS + + +def snake(x: Tensor, alpha: Tensor) -> Tensor: + """Periodic ``x + sin^2(alpha x) / alpha`` activation. + + The reshape to three dimensions is what the reference does and is a no-op + for the rank-3 activations used here; it is kept so the numerics are + identical for any rank. + """ + shape = x.shape + flat = x.reshape(shape[0], shape[1], -1) + flat = flat + (alpha + 1e-9).reciprocal() * torch.sin(alpha * flat).pow(2) + return flat.reshape(shape) + + +class Snake1d(nn.Module): + """Snake activation with one learned frequency per channel.""" + + def __init__(self, channels: int) -> None: + super().__init__() + self.alpha = nn.Parameter(torch.ones(1, channels, 1)) + + def forward(self, x: Tensor) -> Tensor: + return snake(x, self.alpha) + + +def _wn_conv(*args: object, **kwargs: object) -> nn.Module: + """A weight-normalized ``Conv1d``, named ``weight_g``/``weight_v``.""" + return nn.utils.weight_norm(nn.Conv1d(*args, **kwargs)) # type: ignore[arg-type] + + +def _wn_conv_transpose(*args: object, **kwargs: object) -> nn.Module: + """A weight-normalized ``ConvTranspose1d``, named ``weight_g``/``weight_v``.""" + return nn.utils.weight_norm(nn.ConvTranspose1d(*args, **kwargs)) # type: ignore[arg-type] + + +def remove_weight_norm(module: nn.Module) -> int: + """Fold weight normalization into every convolution weight in place. + + Returns: + How many submodules were folded. + """ + removed = 0 + for child in module.modules(): + try: + nn.utils.remove_weight_norm(child) + except (ValueError, RuntimeError): + continue + removed += 1 + return removed + + +class ResidualUnit(nn.Module): + """Dilated snake-gated residual block; preserves length and width.""" + + def __init__(self, dim: int, dilation: int) -> None: + super().__init__() + pad = (_KERNEL_SIZE - 1) * dilation // 2 + self.snake1 = Snake1d(dim) + self.conv1 = _wn_conv(dim, dim, kernel_size=_KERNEL_SIZE, dilation=dilation, padding=pad) + self.snake2 = Snake1d(dim) + self.conv2 = _wn_conv(dim, dim, kernel_size=1) + + def forward(self, x: Tensor) -> Tensor: + y = self.conv2(self.snake2(self.conv1(self.snake1(x)))) + if y.shape[-1] != x.shape[-1]: + # Defensive centre crop. The padding above makes this a no-op for + # every dilation the checkpoint ships. + pad = (x.shape[-1] - y.shape[-1]) // 2 + x = x[..., pad : x.shape[-1] - pad] + return x + y + + +class DecoderBlock(nn.Module): + """One upsampling stage: activate, stretch by ``stride``, then refine.""" + + def __init__(self, input_dim: int, output_dim: int, stride: int) -> None: + super().__init__() + self.snake1 = Snake1d(input_dim) + self.conv_t1 = _wn_conv_transpose( + input_dim, + output_dim, + kernel_size=2 * stride, + stride=stride, + padding=math.ceil(stride / 2), + ) + self.res_unit1 = ResidualUnit(output_dim, _RESIDUAL_DILATIONS[0]) + self.res_unit2 = ResidualUnit(output_dim, _RESIDUAL_DILATIONS[1]) + self.res_unit3 = ResidualUnit(output_dim, _RESIDUAL_DILATIONS[2]) + + def forward(self, x: Tensor) -> Tensor: + x = self.conv_t1(self.snake1(x)) + return self.res_unit3(self.res_unit2(self.res_unit1(x))) + + +class MiniMaxMusic3Vocoder(nn.Module): + """Latent to waveform: ``[B, 128, T] -> [B, 2, T * 512]`` at 44.1 kHz.""" + + def __init__(self) -> None: + super().__init__() + self.dec_in_proj = nn.Conv1d(_STREAM_CHANNELS, _DECODER_INPUT_DIM, kernel_size=1) + self.conv_in = _wn_conv( + _DECODER_INPUT_DIM, + _DECODER_HIDDEN_DIM, + kernel_size=_KERNEL_SIZE, + padding=_KERNEL_SIZE // 2, + ) + blocks: list[nn.Module] = [] + output_dim = _DECODER_HIDDEN_DIM + for index, stride in enumerate(_UPSAMPLING_RATIOS): + input_dim = _DECODER_HIDDEN_DIM // (2**index) + output_dim = _DECODER_HIDDEN_DIM // (2 ** (index + 1)) + blocks.append(DecoderBlock(input_dim, output_dim, stride)) + self.blocks = nn.ModuleList(blocks) + self.snake_out = Snake1d(output_dim) + self.conv_out = _wn_conv(output_dim, 1, kernel_size=_KERNEL_SIZE, padding=_KERNEL_SIZE // 2) + + @property + def upsampling_factor(self) -> int: + """Waveform samples produced per latent frame.""" + return math.prod(_UPSAMPLING_RATIOS) + + @torch.inference_mode() + def forward(self, latent: Tensor) -> Tensor: + """Decode a latent to interleaved stereo. + + Args: + latent: ``[B, 128, T]`` vocoder latent. + + Returns: + ``[B, 2, T * 512]`` waveform at 44.1 kHz, unclamped. + + Raises: + ValueError: If the latent is not ``[B, 128, T]``. + """ + if latent.ndim != 3 or latent.shape[1] != DIT_LATENT_CHANNELS: + raise ValueError( + f"MiniMax Music 3 vocoder latent must be [B,{DIT_LATENT_CHANNELS},T], got {tuple(latent.shape)}" + ) + bsz, _, frames = latent.shape + folded = latent.reshape(bsz * OUTPUT_CHANNELS, _STREAM_CHANNELS, frames) + h = self.conv_in(self.dec_in_proj(folded)) + for block in self.blocks: + h = block(h) + wave = torch.tanh(self.conv_out(self.snake_out(h))) + return wave.reshape(bsz, OUTPUT_CHANNELS, -1) + + +class StreamingResampler: + """Resample a chunked stream to the response rate without a seam. + + A polyphase resampler is time-invariant but not offset-invariant: it only + reproduces the same output grid when the input offset is a whole multiple + of ``orig // gcd(orig, new)`` samples, and it zero-pads both ends of every + call. Resampling each decoded chunk on its own therefore shifts the grid + by a fraction of a sample at every join and tapers the signal across it. + Measured against a single whole-stream resample on 20 s of broadband + noise, that costs a peak error of 2.2 (full scale) and drifts the length. + + This buffers instead. Only whole blocks are emitted, so every boundary + lands on an exact output sample; the previously emitted tail is prepended + as real left context and the not-yet-emitted head appended as real right + context, and both are dropped from the result. On the same measurement + the emitted stream is bit-identical to the whole-stream resample. + + One instance per request. It holds at most a couple of blocks of audio. + """ + + def __init__( + self, + *, + orig_freq: int = DAV_SAMPLE_RATE, + new_freq: int = OUTPUT_SAMPLE_RATE, + channels: int = OUTPUT_CHANNELS, + context_blocks: int = 1, + margin_blocks: int = 1, + ) -> None: + self.orig_freq = int(orig_freq) + self.new_freq = int(new_freq) + self.channels = int(channels) + # The output grid repeats every this many input samples. + self._block = self.orig_freq // math.gcd(self.orig_freq, self.new_freq) + self._context_len = self._block * context_blocks + self._margin = self._block * margin_blocks + self._pending: Tensor | None = None + self._context: Tensor | None = None + + def _empty(self, like: Tensor | None) -> Tensor: + if like is None: + return torch.zeros((self.channels, 0), dtype=torch.float32) + return like.new_zeros((self.channels, 0)) + + def reset(self) -> None: + """Drop all buffered audio.""" + self._pending = None + self._context = None + + @torch.inference_mode() + def push(self, chunks: Sequence[Tensor], *, flush: bool = False) -> Tensor: + """Append source-rate audio and return whatever is now emittable. + + Args: + chunks: New ``[channels, T]`` segments at ``orig_freq``, in order. + flush: Emit the remainder too, because the stream has ended. + + Returns: + ``[channels, T]`` at ``new_freq``; empty when nothing is ready. + """ + import torchaudio.functional as audio_functional + + parts = [chunk for chunk in chunks if chunk.numel()] + if self._pending is not None: + parts.insert(0, self._pending) + self._pending = None + if not parts: + return self._empty(self._context) + + buffered = torch.cat(parts, dim=-1) if len(parts) > 1 else parts[0] + total = int(buffered.shape[-1]) + if flush: + emit_len, lookahead = total, 0 + else: + emit_len = max(0, ((total - self._margin) // self._block) * self._block) + lookahead = min(self._margin, total - emit_len) + if emit_len == 0: + self._pending = buffered + return self._empty(buffered) + + padded = buffered[..., : emit_len + lookahead] + start = 0 + if self._context is not None: + padded = torch.cat((self._context, padded), dim=-1) + # Exact because the context is a whole number of blocks. + start = int(self._context.shape[-1]) * self.new_freq // self.orig_freq + resampled = audio_functional.resample(padded.float(), self.orig_freq, self.new_freq) + # Round up, matching how a whole-stream resample sizes its output. This + # only bites on the flush, where the remainder is not a whole block; + # every other emission divides exactly. + stop = start + -(-emit_len * self.new_freq // self.orig_freq) + + if flush: + self.reset() + else: + self._context = buffered[..., max(0, emit_len - self._context_len) : emit_len].clone() + self._pending = buffered[..., emit_len:].clone() if emit_len < total else None + return resampled[..., start:stop] diff --git a/vllm_omni/model_executor/models/minimax_music3/depth_graph.py b/vllm_omni/model_executor/models/minimax_music3/depth_graph.py new file mode 100644 index 000000000..9c9a4f8e5 --- /dev/null +++ b/vllm_omni/model_executor/models/minimax_music3/depth_graph.py @@ -0,0 +1,137 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""CUDA graph capture for the RVQ depth loop. + +The depth loop is seven sequential forwards of a four-layer block over a +sequence that grows from two to eight positions, plus a guided draw after each +one. At the batch sizes this stage runs, none of those kernels is anywhere near +large enough to hide its own launch: measured in isolation on one L20X, the +loop costs 8.4 ms per frame when driven from Python and 4.0 ms when replayed +from a graph, and that figure barely moves between one and eight guided pairs. +It is launch-bound, so it is worth recording. + +Capture is lazy and keyed on the exact pair count, never padded. That keeps the +recorded kernels the same ones the eager loop selected for that shape, which is +what makes a replay bit-identical to the loop it replaces rather than merely +close to it. +""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass + +import torch +from vllm.logger import init_logger +from vllm.platforms import current_platform + +logger = init_logger(__name__) + +__all__ = ["DepthGraphCache"] + +# Inputs, in the order the eager loop takes them. +_DepthInputs = tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor] +_DepthOutputs = tuple[torch.Tensor, torch.Tensor] + + +@dataclass +class _CapturedDepth: + """One recorded depth loop and the buffers it reads and writes.""" + + graph: torch.cuda.CUDAGraph + inputs: _DepthInputs + codes: torch.Tensor + depth_hidden: torch.Tensor + + +class DepthGraphCache: + """Replay the depth loop from a graph once its shape has settled. + + Args: + run_eager: The loop itself, taking ``(cond_hidden, uncond_hidden, c0, + seeds, positions)`` and returning ``(codes, depth_hidden)``. + warmup_steps: How many times a pair count runs eagerly before it is + recorded. + """ + + def __init__( + self, + run_eager: Callable[..., _DepthOutputs], + *, + warmup_steps: int = 2, + ) -> None: + self._run_eager = run_eager + self._warmup_steps = warmup_steps + self._captured: dict[int, _CapturedDepth] = {} + self._eager_steps: dict[int, int] = {} + self._disabled: set[int] = set() + + def run(self, *inputs: torch.Tensor) -> _DepthOutputs: + """Run one depth loop, recording or replaying it where possible.""" + pairs = int(inputs[2].shape[0]) + entry = self._captured.get(pairs) + if entry is None: + if not self._should_capture(pairs, inputs[0]): + return self._run_eager(*inputs) + entry = self._capture(pairs, inputs) + if entry is None: + return self._run_eager(*inputs) + + for static, live in zip(entry.inputs, inputs, strict=True): + static.copy_(live) + entry.graph.replay() + # The codes are fed back as the next frame's input and appended to the + # request's trajectory, so they outlive the step; the graph overwrites + # its output buffer on the next replay. The depth hidden states are + # consumed inside this step by the concatenation that builds the + # conditioning frame, which allocates, so that one can be handed back + # as is. + return entry.codes.clone(), entry.depth_hidden + + def _should_capture(self, pairs: int, sample: torch.Tensor) -> bool: + if pairs in self._disabled or not sample.is_cuda: + return False + if torch.cuda.is_current_stream_capturing(): + # Already inside somebody else's capture; nested graphs are not a + # thing and the outer recording covers this work anyway. + return False + count = self._eager_steps.get(pairs, 0) + 1 + self._eager_steps[pairs] = count + # Let a new pair count run eagerly a few times first. Capture records + # whatever the kernels do at that moment, and the first call into a + # cuBLAS or SDPA path for an unseen shape can allocate a workspace or + # pick an algorithm, neither of which belongs inside the recording. + return count > self._warmup_steps + + def _capture(self, pairs: int, inputs: _DepthInputs) -> _CapturedDepth | None: + try: + static = tuple(torch.empty_like(tensor) for tensor in inputs) + for target, live in zip(static, inputs, strict=True): + target.copy_(live) + graph = torch.cuda.CUDAGraph() + with ( + torch.inference_mode(), + torch.cuda.graph( + graph, + pool=current_platform.get_global_graph_pool(), + # The engine's connector threads keep their own streams + # busy while this runs. They never touch the depth loop's + # buffers, so restrict the legality check to this thread + # rather than failing the capture on their activity. + capture_error_mode="thread_local", + ), + ): + codes, depth_hidden = self._run_eager(*static) + entry = _CapturedDepth(graph=graph, inputs=static, codes=codes, depth_hidden=depth_hidden) + self._captured[pairs] = entry + logger.info("MiniMax Music 3 captured the RVQ depth CUDA graph for %d guided pair(s)", pairs) + return entry + except Exception as exc: # noqa: BLE001 - any capture failure must fall back, not crash + self._disabled.add(pairs) + logger.warning( + "MiniMax Music 3 RVQ depth CUDA graph disabled for %d guided pair(s), " + "falling back to the eager loop: %s", + pairs, + exc, + ) + return None diff --git a/vllm_omni/model_executor/models/minimax_music3/dit.py b/vllm_omni/model_executor/models/minimax_music3/dit.py new file mode 100644 index 000000000..cf180ac52 --- /dev/null +++ b/vllm_omni/model_executor/models/minimax_music3/dit.py @@ -0,0 +1,484 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Condition encoder and flow-matching transformer for MiniMax Music 3. + +Stage 1 turns the AR stage's conditioning frames into a DAC latent in two +steps. The condition encoder collapses the eight 4096-wide layer slices of a +frame into one 2048-wide vector and stretches the result from the AR frame +grid onto the vocoder's latent grid. The transformer is then solved as a +flow-matching velocity field: a fixed-step Euler integrator walks Gaussian +noise to the latent the vocoder expects, under classifier-free guidance. + +Module names mirror the checkpoint's ``condition_encoder/`` and +``transformer/`` component folders key for key, with one deliberate +exception: the per-block ``to_q``/``to_k``/``to_v`` projections are fused into +a single ``to_qkv`` so each block runs one GEMM instead of three. +:func:`remap_transformer_state` performs that fusion, and is the only place +the checkpoint's layout is reinterpreted. + +Attention runs through the project's own diffusion attention layer, so this +model picks up the same backend selection, sequence-parallel plumbing and +KV-quant policy as every other transformer in the tree. That layer resolves +its backend from the diffusion config, which an ``LLM_GENERATION`` stage such +as this one does not set; it tolerates that and falls through to the platform +default. If it cannot be built at all the block runs plain SDPA instead, a +choice made once when the blocks are built and never inside ``forward``. +""" + +from __future__ import annotations + +import math + +import torch +from torch import Tensor, nn +from torch.nn import functional as F +from vllm.logger import init_logger + +from .constants import ( + AR_HIDDEN_SIZE, + BACKBONE_HIDDEN_SIZE, + CONDITION_INPUT_HOP, + CONDITION_INPUT_SAMPLE_RATE, + CONDITION_OUTPUT_HOP, + CONDITION_OUTPUT_SAMPLE_RATE, + DEFAULT_DIT_CFG_SCALE, + DEFAULT_DIT_STEPS, + DIT_LATENT_CHANNELS, + NUM_CODEBOOKS, +) + +logger = init_logger(__name__) + +__all__ = [ + "MiniMaxMusic3ConditionEncoder", + "MiniMaxMusic3FlowMatchingDiT", + "MiniMaxMusic3Transformer1DModel", + "remap_transformer_state", +] + +# Fixed by transformer/config.json. Declared rather than read from the file so +# the skeleton can be built before any checkpoint is on disk, which is what +# vLLM's memory profiler needs at startup. +_DIM = 2_048 +_NUM_LAYERS = 36 +_HEAD_DIM = 64 +_FF_INNER_DIM = 8_192 +_FOURIER_DIM = 256 +_ROTARY_DIM = 32 +_ROTARY_BASE = 10_000.0 +_CONDITION_DIM = 2_048 +# The solver feeds the transformer the current latent, a zero block of the +# same width, and the condition: 128 + 128 + 2048 = 2304. +_TRANSFORMER_IN_DIM = DIT_LATENT_CHANNELS * 2 + _CONDITION_DIM + +# The prompt interpolation floor from the reference sampler. Sigma never +# reaches exactly zero, so the noise term never fully vanishes. +_SIGMA_FLOOR = 1e-6 + + +class FourierFeatures(nn.Module): + """Random-Fourier timestep features: ``[B, 1] -> [B, out_features]``.""" + + def __init__(self, in_features: int, out_features: int) -> None: + super().__init__() + self.weight = nn.Parameter(torch.empty(out_features // 2, in_features)) + + def forward(self, value: Tensor) -> Tensor: + f = 2.0 * math.pi * value @ self.weight.T + return torch.cat((f.cos(), f.sin()), dim=-1) + + +class TimestepEmbedding(nn.Module): + """The checkpoint's ``time_embed``: Linear, SiLU, Linear.""" + + def __init__(self, in_dim: int, dim: int) -> None: + super().__init__() + self.linear_1 = nn.Linear(in_dim, dim) + self.act = nn.SiLU() + self.linear_2 = nn.Linear(dim, dim) + + def forward(self, x: Tensor) -> Tensor: + return self.linear_2(self.act(self.linear_1(x))) + + +def _rotate_half(x: Tensor) -> Tensor: + x1, x2 = x.chunk(2, dim=-1) + return torch.cat((-x2, x1), dim=-1) + + +def _apply_rope(x: Tensor, rope_cos: Tensor, rope_sin: Tensor) -> Tensor: + """Rotate the leading ``rope_cos.shape[-1]`` channels of ``[B, S, H, D]``. + + Only the first 32 of each head's 64 channels are rotated; the rest pass + through untouched, which is what this checkpoint was trained with. The + tables are cached at the exact sequence length, so there is nothing to + slice off the front. + """ + rot_dim = rope_cos.shape[-1] + rotated = x[..., :rot_dim] + rotated = rotated * rope_cos + _rotate_half(rotated) * rope_sin + return torch.cat((rotated, x[..., rot_dim:]), dim=-1) + + +# Set to the failure reason the first time the native attention layer cannot be +# built, so the 36 blocks agree on one path and the reason is logged once. +_NATIVE_ATTENTION_UNAVAILABLE: str | None = None + + +def _build_native_attention(*, num_heads: int, head_size: int, softmax_scale: float, prefix: str) -> nn.Module | None: + """Build the project's diffusion attention layer, or ``None`` to use SDPA. + + Anything that goes wrong here is a legitimate reason to run plain SDPA, + and is reported once rather than once per block. + """ + global _NATIVE_ATTENTION_UNAVAILABLE + if _NATIVE_ATTENTION_UNAVAILABLE is not None: + return None + try: + from vllm_omni.diffusion.attention.layer import Attention as DiffusionAttention + + return DiffusionAttention( + num_heads=num_heads, + head_size=head_size, + causal=False, + softmax_scale=softmax_scale, + prefix=prefix, + role="self", + role_category="self", + ) + except Exception as exc: # noqa: BLE001 - any failure means: use SDPA + _NATIVE_ATTENTION_UNAVAILABLE = repr(exc) + logger.warning( + "MiniMax Music 3 DiT could not build the diffusion attention layer " + "(%s); running torch SDPA instead. The stage decodes in float32, " + "for which the diffusion layer dispatches to SDPA anyway, so this " + "changes performance bookkeeping and not the audio.", + exc, + ) + return None + + +class Attention(nn.Module): + """Bidirectional self-attention with partial rotary position embeddings. + + Q, K and V are kept in ``[B, S, H, D]``, which is the layout the diffusion + attention backends consume and return. + """ + + def __init__(self, dim: int, *, head_dim: int, prefix: str = "") -> None: + super().__init__() + self.num_heads = dim // head_dim + self.head_dim = head_dim + self.softmax_scale = float(head_dim**-0.5) + # Fused from the checkpoint's separate to_q/to_k/to_v; see + # remap_transformer_state for the concatenation order. + self.to_qkv = nn.Linear(dim, dim * 3, bias=False) + self.to_out = nn.Linear(dim, dim, bias=False) + self.backend = _build_native_attention( + num_heads=self.num_heads, + head_size=head_dim, + softmax_scale=self.softmax_scale, + prefix=prefix, + ) + self.backend_name = self.backend.attn_backend.get_name() if self.backend is not None else "TORCH_SDPA" + + def _attend(self, q: Tensor, k: Tensor, v: Tensor) -> Tensor: + if self.backend is not None: + return self.backend(q, k, v) + q, k, v = (t.transpose(1, 2) for t in (q, k, v)) + out = F.scaled_dot_product_attention(q, k, v, is_causal=False, scale=self.softmax_scale) + return out.transpose(1, 2) + + def forward(self, x: Tensor, rope_cos: Tensor, rope_sin: Tensor) -> Tensor: + bsz, seq, dim = x.shape + q, k, v = self.to_qkv(x).chunk(3, dim=-1) + q = q.reshape(bsz, seq, self.num_heads, self.head_dim) + k = k.reshape(bsz, seq, self.num_heads, self.head_dim) + v = v.reshape(bsz, seq, self.num_heads, self.head_dim) + q = _apply_rope(q, rope_cos, rope_sin) + k = _apply_rope(k, rope_cos, rope_sin) + out = self._attend(q, k, v) + return self.to_out(out.contiguous().reshape(bsz, seq, dim)) + + +class TransformerBlock(nn.Module): + """Pre-norm attention followed by a pre-norm gated feed-forward.""" + + def __init__(self, dim: int, *, head_dim: int, inner_dim: int, prefix: str = "") -> None: + super().__init__() + self.norm1 = nn.LayerNorm(dim) + self.attn = Attention(dim, head_dim=head_dim, prefix=prefix) + self.norm2 = nn.LayerNorm(dim) + # One fused projection produces the GLU's value and gate halves. + self.ff_in = nn.Linear(dim, inner_dim * 2) + self.ff_out = nn.Linear(inner_dim, dim) + + def forward(self, x: Tensor, rope_cos: Tensor, rope_sin: Tensor) -> Tensor: + x = x + self.attn(self.norm1(x), rope_cos, rope_sin) + value, gate = self.ff_in(self.norm2(x)).chunk(2, dim=-1) + return x + self.ff_out(value * F.silu(gate)) + + +class MiniMaxMusic3Transformer1DModel(nn.Module): + """The 36-block velocity field: ``([B,128,T], [B], [B,2048,T]) -> [B,128,T]``.""" + + def __init__(self) -> None: + super().__init__() + self.preprocess_conv = nn.Conv1d(_TRANSFORMER_IN_DIM, _TRANSFORMER_IN_DIM, 1, bias=False) + self.postprocess_conv = nn.Conv1d(DIT_LATENT_CHANNELS, DIT_LATENT_CHANNELS, 1, bias=False) + self.proj_in = nn.Linear(_TRANSFORMER_IN_DIM, _DIM, bias=False) + self.proj_out = nn.Linear(_DIM, DIT_LATENT_CHANNELS, bias=False) + self.time_proj = FourierFeatures(1, _FOURIER_DIM) + self.time_embed = TimestepEmbedding(_FOURIER_DIM, _DIM) + self.transformer_blocks = nn.ModuleList( + TransformerBlock( + _DIM, + head_dim=_HEAD_DIM, + inner_dim=_FF_INNER_DIM, + prefix=f"transformer_blocks.{index}.attn", + ) + for index in range(_NUM_LAYERS) + ) + inv_freq = 1.0 / (_ROTARY_BASE ** (torch.arange(0, _ROTARY_DIM, 2).float() / _ROTARY_DIM)) + self.register_buffer("inv_freq", inv_freq, persistent=False) + self._rope_cache: dict[tuple[int, torch.dtype, torch.device], tuple[Tensor, Tensor]] = {} + + @property + def attention_backend_name(self) -> str: + """The attention path the blocks actually run, for startup logging.""" + return self.transformer_blocks[0].attn.backend_name + + def _rope(self, seq_len: int, *, dtype: torch.dtype, device: torch.device) -> tuple[Tensor, Tensor]: + """Return cached ``(cos, sin)`` shaped ``[seq_len, 1, 32]``. + + The head axis is left as a singleton so the tables broadcast over + ``[B, S, H, D]`` without transposing Q and K. + """ + key = (seq_len, dtype, device) + cached = self._rope_cache.get(key) + if cached is None: + t = torch.arange(seq_len, device=self.inv_freq.device, dtype=torch.float32) + freqs = torch.einsum("i,j->ij", t, self.inv_freq) + freqs = torch.cat((freqs, freqs), dim=-1).to(dtype=dtype, device=device).unsqueeze(1) + cached = (freqs.cos(), freqs.sin()) + self._rope_cache[key] = cached + return cached + + def forward(self, x: Tensor, t: Tensor, condition: Tensor) -> Tensor: + """Predict the flow velocity at latent ``x`` and time ``t``. + + Args: + x: Current latent, ``[B, 128, T_mel]``. + t: Flow time in ``[0, 1)``, ``[B]``. + condition: Aligned condition, ``[B, 2048, T_mel]``. + + Returns: + The velocity, ``[B, 128, T_mel]``. + """ + full = torch.cat((x, torch.zeros_like(x), condition), dim=1) + full = self.preprocess_conv(full) + full + timestep_embed = self.time_embed(self.time_proj(t[:, None])) + h = self.proj_in(full.transpose(1, 2)) + # The timestep rides as an extra leading position and is dropped again + # after the stack, so it shifts every rotary phase by one step. + h = torch.cat((timestep_embed.unsqueeze(1), h), dim=1) + rope_cos, rope_sin = self._rope(h.shape[1], dtype=h.dtype, device=h.device) + for block in self.transformer_blocks: + h = block(h, rope_cos, rope_sin) + out = self.proj_out(h[:, 1:]).transpose(1, 2) + return self.postprocess_conv(out) + out + + +def remap_transformer_state(state: dict[str, Tensor]) -> dict[str, Tensor]: + """Rewrite ``transformer/`` checkpoint keys onto this module tree. + + Two rewrites, both confined to the attention block: + + * ``attn.to_q/to_k/to_v.weight`` are concatenated along the output axis + into ``attn.to_qkv.weight``. ``nn.Linear`` output feature ``i`` reads + weight row ``i``, and :meth:`Attention.forward` splits the projection + with ``chunk(3, dim=-1)`` into ``q, k, v``, so query rows come first and + value rows last. + * ``attn.to_out.0.weight`` loses its index: the checkpoint models + ``to_out`` as a two-element list whose second element is a parameterless + dropout, and inference only needs the projection. + + Every other key is copied through untouched, including the fused ``ff_in`` + whose halves are consumed as ``value, gate`` in that order. + + Raises: + ValueError: If a block is missing one of its three projections. + """ + remapped: dict[str, Tensor] = {} + for key, value in state.items(): + if key.endswith(".attn.to_out.0.weight"): + remapped[key.replace(".to_out.0.weight", ".to_out.weight")] = value + continue + if ".attn.to_q.weight" in key: + prefix = key[: -len("to_q.weight")] + try: + fused = [state[f"{prefix}to_{name}.weight"] for name in "qkv"] + except KeyError as exc: + raise ValueError( + f"MiniMax Music 3 transformer block {prefix!r} is missing attention projection {exc}" + ) from exc + remapped[f"{prefix}to_qkv.weight"] = torch.cat(fused, dim=0) + continue + if ".attn.to_k.weight" in key or ".attn.to_v.weight" in key: + continue + remapped[key] = value + return remapped + + +class MiniMaxMusic3ConditionEncoder(nn.Module): + """Collapse an AR frame's eight layer slices and resample onto mel time.""" + + def __init__(self) -> None: + super().__init__() + self.layer_weight_logits = nn.Parameter(torch.zeros(NUM_CODEBOOKS)) + self.layer_scale = nn.Parameter(torch.ones(1)) + self.proj = nn.Conv1d(BACKBONE_HIDDEN_SIZE, _CONDITION_DIM, kernel_size=3, padding=1) + + def aligned_mel_length(self, frames: int) -> int: + """Return how many vocoder latent frames ``frames`` AR frames cover. + + AR frames sit on a 24 kHz / 960-hop grid (25 fps) and the vocoder + latent on a 44.1 kHz / 512-hop grid (about 86 fps), so one AR frame is + roughly 3.445 latent frames. The result is truncated, never rounded. + """ + return max( + 1, + int( + frames + * CONDITION_OUTPUT_SAMPLE_RATE + / CONDITION_INPUT_SAMPLE_RATE + * CONDITION_INPUT_HOP + / CONDITION_OUTPUT_HOP + ), + ) + + def condition(self, hidden: Tensor) -> Tensor: + """Project ``[B, T, 32768]`` AR frames to ``[B, 2048, T]``. + + Raises: + ValueError: If the conditioning is not ``[B, T, 32768]``. + """ + if hidden.ndim != 3 or hidden.shape[-1] != AR_HIDDEN_SIZE: + raise ValueError(f"MiniMax Music 3 condition must be [B,T,{AR_HIDDEN_SIZE}], got {tuple(hidden.shape)}") + hidden_cf = hidden.transpose(1, 2) + bsz, _, frames = hidden_cf.shape + hidden_cf = hidden_cf.reshape(bsz, NUM_CODEBOOKS, BACKBONE_HIDDEN_SIZE, frames) + weights = torch.softmax(self.layer_weight_logits, dim=0).to(hidden_cf.dtype) + hidden_cf = torch.einsum("blht,l->bht", hidden_cf, weights) + hidden_cf = self.layer_scale.to(hidden_cf.dtype) * hidden_cf + return self.proj(hidden_cf) + + def aligned_condition(self, hidden: Tensor) -> Tensor: + """Project AR hidden states and resample them onto the latent grid.""" + hidden = hidden.to(dtype=next(self.parameters()).dtype) + align = self.condition(hidden) + mel_len = self.aligned_mel_length(hidden.shape[1]) + return F.interpolate(align, size=mel_len, mode="nearest") + + +class MiniMaxMusic3FlowMatchingDiT(nn.Module): + """Condition encoder plus velocity field, with the Euler solver on top.""" + + def __init__(self) -> None: + super().__init__() + self.condition_encoder = MiniMaxMusic3ConditionEncoder() + self.transformer = MiniMaxMusic3Transformer1DModel() + + def aligned_mel_length(self, frames: int) -> int: + return self.condition_encoder.aligned_mel_length(frames) + + def condition(self, hidden: Tensor) -> Tensor: + return self.condition_encoder.condition(hidden) + + def aligned_condition(self, hidden: Tensor) -> Tensor: + return self.condition_encoder.aligned_condition(hidden) + + @torch.inference_mode() + def forward( + self, + align: Tensor, + *, + generator: torch.Generator, + initial_latent: Tensor | None = None, + initial_condition: Tensor | None = None, + num_steps: int = DEFAULT_DIT_STEPS, + cfg_scale: float = DEFAULT_DIT_CFG_SCALE, + ) -> Tensor: + """Solve an aligned condition into a vocoder latent ``[1, 128, T_mel]``. + + The previous window's latent is re-imposed at every step rather than + only at the start: the leading ``left`` frames are pinned to the exact + noise-to-latent interpolation the solver would have produced had it + generated them, so the two windows join without a seam. + + Args: + align: Aligned condition, ``[1, 2048, T_mel]``. Overwritten in + place over the prompt span when ``initial_condition`` is + given, which the caller relies on when it saves the condition + for the next window. + generator: Seeded generator for the initial noise draw. + initial_latent: Previous window's tail latent, or ``None``. + initial_condition: Previous window's tail condition, or ``None``. + num_steps: Euler steps. + cfg_scale: Classifier-free guidance weight. + + Returns: + The solved latent, ``[1, 128, T_mel]``. + + Raises: + ValueError: If ``num_steps`` is not positive. + """ + if num_steps < 1: + raise ValueError("MiniMax Music 3 DiT num_steps must be positive") + mel_len = align.shape[-1] + x = torch.randn( + (align.shape[0], DIT_LATENT_CHANNELS, mel_len), + device=align.device, + dtype=align.dtype, + generator=generator, + ) + + left = 0 + latent_prompt: Tensor | None = None + noise_prompt: Tensor | None = None + if initial_latent is not None: + left = min(initial_latent.shape[-1], mel_len) + if left <= 0: + left = 0 + else: + latent_prompt = initial_latent[..., :left].to(x) + noise_prompt = x[..., :left].clone() + if initial_condition is not None: + align[..., :left] = initial_condition[..., :left].to(align) + + dt = 1.0 / num_steps + # Row 0 carries the condition, row 1 stays zero: one batched forward + # evaluates both guidance branches. + cond_cfg = torch.zeros((2, *align.shape[1:]), device=align.device, dtype=align.dtype) + cond_cfg[0].copy_(align[0]) + # Every timestep in one host-to-device copy. Built from the same + # Python divisions the per-step construction would have used, so the + # values are identical, but the solver loop then touches the host only + # to launch kernels. + schedule = torch.tensor( + [step / num_steps for step in range(num_steps)], + device=align.device, + dtype=align.dtype, + ) + + for step in range(num_steps): + t = schedule[step].expand(x.shape[0]) + if left and latent_prompt is not None and noise_prompt is not None: + x[..., :left] = (1.0 - (1.0 - _SIGMA_FLOOR) * t[0]) * noise_prompt + t[0] * latent_prompt + d = self.transformer(x.expand(2, -1, -1), t.expand(2), cond_cfg) + d = cfg_scale * d[:1] + (1.0 - cfg_scale) * d[1:2] + x = x + dt * d + + if left and latent_prompt is not None: + x[..., :left] = latent_prompt + return x diff --git a/vllm_omni/model_executor/models/minimax_music3/frame_state.py b/vllm_omni/model_executor/models/minimax_music3/frame_state.py new file mode 100644 index 000000000..5191cebac --- /dev/null +++ b/vllm_omni/model_executor/models/minimax_music3/frame_state.py @@ -0,0 +1,158 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Per-request AR state for MiniMax Music 3. + +The talker accumulates one 32768-dimensional conditioning frame per decode +step and hands the acoustic stage overlapping 200-frame windows. Windows are +emitted one frame late on purpose: a window can only be declared non-final +once a further frame proves the song did not end at its boundary, and the +tail is flushed when the request finishes. + +State is keyed by request id rather than batch row, because a row's identity +is not stable across engine steps. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field + +import torch + +from .chunking import ChunkWindow +from .constants import AR_CHUNK_FRAMES, AR_CHUNK_HOP_FRAMES, NUM_CODEBOOKS + + +class FrameBuffer: + """Rolling buffer of conditioning frames addressed by absolute frame index. + + Frames before an emitted window's hop point are dropped, so peak memory is + bounded by the window size rather than by the length of the song. + """ + + def __init__(self) -> None: + self.base_frame = 0 + self._frames: list[torch.Tensor] = [] + self.peak_frames = 0 + + @property + def end_frame(self) -> int: + return self.base_frame + len(self._frames) + + def append(self, frame: torch.Tensor) -> None: + """Append one ``[1, hidden]`` frame.""" + self._frames.append(frame) + self.peak_frames = max(self.peak_frames, len(self._frames)) + + def slice(self, start: int, end: int) -> torch.Tensor: + """Return frames ``[start, end)`` stacked as ``[end - start, hidden]``. + + Raises: + RuntimeError: If the range has already been discarded or has not + been produced yet. + """ + rel_start = start - self.base_frame + rel_end = end - self.base_frame + if rel_start < 0 or rel_end > len(self._frames): + raise RuntimeError( + f"MiniMax Music 3 frame range [{start}, {end}) is unavailable; " + f"the buffer holds [{self.base_frame}, {self.end_frame})" + ) + return torch.cat(self._frames[rel_start:rel_end], dim=0) + + def discard_before(self, frame: int) -> None: + drop = min(max(frame - self.base_frame, 0), len(self._frames)) + if drop: + del self._frames[:drop] + self.base_frame += drop + + def clear(self) -> None: + self._frames.clear() + self.base_frame = 0 + + +@dataclass +class RequestARState: + """Everything the talker tracks for one conditioned request.""" + + seed: int + max_frames: int + # Prompt length in tokens, used to tell a decode row apart from a row that + # is still working through a chunked prefill. + prompt_tokens: int = 0 + sampling_position: int = 0 + generated_frames: int = 0 + emitted_windows: int = 0 + finished: bool = False + # The engine is told to stop one step after the song ends. Payloads are + # built from the buffer that ``postprocess`` filled on the previous step, + # so without that extra step the terminal span would be written and then + # discarded with the request. + pending_stop: bool = False + finish_reason: str = "length" + frames: FrameBuffer = field(default_factory=FrameBuffer) + codes: list[torch.Tensor] = field(default_factory=list) + codes_base_frame: int = 0 + # The previous frame's eight codes, fed back as the next step's input + # embedding. Held on the request rather than on a batch row, because the + # row a request occupies is not stable between steps. + last_codes: torch.Tensor | None = None + + def advance_position(self) -> int: + """Return this frame's base draw position and reserve the frame's slot. + + Each frame consumes ``NUM_CODEBOOKS`` positions so that the ``c0`` draw + and the seven depth draws never share an RNG stream. + """ + position = self.sampling_position + self.sampling_position += NUM_CODEBOOKS + return position + + def ready_window(self) -> ChunkWindow | None: + """Return the next window that a lookahead frame has proven non-final.""" + start = self.emitted_windows * AR_CHUNK_HOP_FRAMES + if self.frames.end_frame <= start + AR_CHUNK_FRAMES: + return None + return ChunkWindow( + index=self.emitted_windows, + start=start, + end=start + AR_CHUNK_FRAMES, + is_first=self.emitted_windows == 0, + is_last=False, + ) + + def take_window(self, window: ChunkWindow, *, is_final: bool) -> torch.Tensor: + """Materialize ``window`` and release the frames it retires.""" + chunk = self.frames.slice(window.start, window.end) + self.emitted_windows = max(self.emitted_windows, window.index + 1) + if not is_final: + self.frames.discard_before(window.start + AR_CHUNK_HOP_FRAMES) + return chunk + + def take_codes(self, start: int, end: int) -> torch.Tensor | None: + """Return the sampled codes for frames ``[start, end)``, if retained. + + Codes are carried alongside the conditioning frames so a downstream + stage or a debugging harness can recover the exact RVQ trajectory. + They are dropped rather than raising when a span predates what is + still buffered, since nothing in the audio path depends on them. + """ + if not self.codes or end <= start: + return None + offset = self.codes_base_frame + lo, hi = start - offset, end - offset + if lo < 0 or hi > len(self.codes): + return None + return torch.cat(self.codes[lo:hi], dim=0) + + def discard_codes_before(self, frame: int) -> None: + drop = min(max(frame - self.codes_base_frame, 0), len(self.codes)) + if drop: + del self.codes[:drop] + self.codes_base_frame += drop + + def release(self) -> None: + self.frames.clear() + self.codes.clear() + + +__all__ = ["FrameBuffer", "RequestARState"] diff --git a/vllm_omni/model_executor/models/minimax_music3/pipeline.py b/vllm_omni/model_executor/models/minimax_music3/pipeline.py new file mode 100644 index 000000000..1a8279f05 --- /dev/null +++ b/vllm_omni/model_executor/models/minimax_music3/pipeline.py @@ -0,0 +1,103 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""MiniMax Music 3 pipeline topology: AR talker -> acoustic decoder. + +Stage 0 ``minimax_music3_ar``: a Qwen3 backbone (the repo's ``language_model/`` +subfolder is a plain ``Qwen3ForCausalLM``) plus a depth decoder, wrapped by a +class that owns its sampler. It emits 200-frame windows of frame hidden states +and the matching 8-codebook RVQ codes. + +Stage 1 ``minimax_music3_acoustic``: flow-matching DiT plus DAV vocoder. Pure +compute, no KV cache, weights loaded from the repo's ``condition_encoder/``, +``transformer/`` and ``vocoder/`` subfolders. Output is 32 kHz stereo. + +Classifier-free guidance is mandatory for this checkpoint: every request +decodes as a cond/uncond pair, and the uncond row is produced by +``prompt_expand_func``. The deploy YAMLs therefore carry +``default_sampling_params.extra_args`` with ``cfg_role`` and ``cfg_scale`` so +that ``has_sampling_extra_args`` is on for the stage and per-row extra args +reach the model's forward. +""" + +from vllm_omni.config.stage_config import ( + PipelineConfig, + StageExecutionType, + StagePipelineConfig, +) + +_PROC = "vllm_omni.model_executor.stage_input_processors.minimax_music3" + +# ``<|audio_end|>`` in the checkpoint tokenizer (pinned by +# ``models/minimax_music3/prompt.SPECIAL_TOKEN_IDS``). The talker's own sampler +# forces this at the frame budget; the stop id is the engine-side safety net. +MINIMAX_MUSIC3_AUDIO_END_TOKEN_ID = 151670 + +# MiniMax captions carry the full style description and run to roughly 5000 +# characters, an order of magnitude past the 500-character serving default. +# ``tts_args`` is a pipeline-level stage extra, not a deploy-YAML field: it +# reaches the serving layer through ``StageConfig.yaml_extras``, which is where +# ``ServingSpeech._compute_max_instructions_length`` reads it from. +_MAX_INSTRUCTIONS_LENGTH = 20_000 + +MINIMAX_MUSIC3_PIPELINE = PipelineConfig( + model_type="minimax_music3", + model_arch="MiniMaxMusic3TalkerForConditionalGeneration", + hf_architectures=("MiniMaxMusic3ForConditionalGeneration",), + default_deploy_config_name="minimax_music3.yaml", + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="minimax_music3_ar", + execution_type=StageExecutionType.LLM_AR, + input_sources=(), + owns_tokenizer=True, + engine_output_type="latent", + # The repo root config.json only names the composite architecture; + # the AR weights and the tokenizer live in their own subfolders. + model_subdir="language_model", + tokenizer_subdir="tokenizer", + # detokenize: the sampled stream is audio codes, never text. + # stop_token_ids: <|audio_end|> closes the song. + sampling_constraints={ + "detokenize": False, + "stop_token_ids": [MINIMAX_MUSIC3_AUDIO_END_TOKEN_ID], + }, + # CFG is mandatory here, so this expands on essentially every + # request rather than only on an opt-in flag. + prompt_expand_func=f"{_PROC}.expand_cfg_prompts", + async_chunk_process_next_stage_input_func=f"{_PROC}.ar2acoustic_async_chunk", + # retains_state_across_chunks stays False. It is read only by + # OmniGenerationScheduler, so it is a no-op on an LLM_AR stage; + # more to the point, stage 0 produces chunks rather than parking on + # them, so it never holds a runner slot while waiting for one. + extras={"tts_args": {"max_instructions_length": _MAX_INSTRUCTIONS_LENGTH}}, + ), + StagePipelineConfig( + stage_id=1, + model_stage="minimax_music3_acoustic", + execution_type=StageExecutionType.LLM_GENERATION, + input_sources=(0,), + final_output=True, + final_output_type="audio", + # "audio" (not "latent"): the accumulator concatenates audio over + # time, whereas the latent path would stack the two stereo channels. + engine_output_type="audio", + model_arch="MiniMaxMusic3AcousticForConditionalGeneration", + # This stage reads the repo ROOT so it can reach condition_encoder/, + # transformer/ and vocoder/, and the root ships no tokenizer files. + tokenizer_subdir="tokenizer", + # The acoustic stage carries the previous window's latent forward as + # the next window's flow-matching prompt, so a request parked between + # chunks still owns its model state and must keep counting against + # max_num_seqs. + retains_state_across_chunks=True, + sync_process_input_func=f"{_PROC}.ar2acoustic", + sampling_constraints={"detokenize": True}, + ), + ), +) + +__all__ = [ + "MINIMAX_MUSIC3_AUDIO_END_TOKEN_ID", + "MINIMAX_MUSIC3_PIPELINE", +] diff --git a/vllm_omni/model_executor/models/minimax_music3/prompt.py b/vllm_omni/model_executor/models/minimax_music3/prompt.py new file mode 100644 index 000000000..b9093eae9 --- /dev/null +++ b/vllm_omni/model_executor/models/minimax_music3/prompt.py @@ -0,0 +1,172 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# +# Prompt normalization follows the MiniMaxAI/MiniMax-Music3 input contract. +# The transformations are reproduced byte-for-byte because any divergence +# shifts token positions and therefore changes the generated audio. +"""MiniMax Music 3 prompt construction. + +The prompt is a single flat string, not a chat template:: + + <|im_start|><|caption_start|>{caption}<|caption_end|> + <|lyrics_start|>{lyrics}<|lyrics_end|><|im_end|><|audio_start|> + +The classifier-free-guidance twin reuses the same token ids with the whole +caption-and-lyrics span (delimiters included) overwritten by +``<|audio_cfg|>``, so only the chat markers and ``<|audio_start|>`` survive. +Both rows therefore have identical length and identical audio-start position. +""" + +from __future__ import annotations + +import re +from typing import Any + +# Pinned by ``validate_tokenizer_ids``: a checkpoint whose tokenizer disagrees +# is not this model, and silently producing noise is worse than failing. +SPECIAL_TOKEN_IDS: dict[str, int] = { + "<|im_start|>": 151644, + "<|im_end|>": 151645, + "<|audio_cfg|>": 151654, + "<|audio_start|>": 151669, + "<|audio_end|>": 151670, + "<|caption_start|>": 151671, + "<|caption_end|>": 151672, + "<|lyrics_start|>": 151673, + "<|lyrics_end|>": 151674, +} + +# c0 codes occupy the backbone vocabulary immediately above the text tokens. +AUDIO_CODE_OFFSET = 151675 + +_SPECIAL_TAG_RE = re.compile(r"<\|([^|]*)\|>") +_LEADING_TAGS_RE = re.compile(r"^[ \t]*((?:\[[^\]]+\][ \t]*)+)") + + +def _strip_markdown_residuals(text: str) -> str: + """Strip the Markdown syntax the model input contract does not accept.""" + if not text: + return text + lines_out: list[str] = [] + for raw_line in text.splitlines(): + line = raw_line + line = re.sub(r"^\s{0,3}#{1,6}\s+", "", line) + line = re.sub(r"^\s*[*+-]\s+", "", line) + line = re.sub(r"^\s*\*\s+", "", line) + while "**" in line: + updated = re.sub(r"\*\*([^*]+)\*\*", r"\1", line) + if updated == line: + break + line = updated + line = re.sub(r"(? str: + return _strip_markdown_residuals(text).replace("• ", "").replace(" ", "") + + +def clean_caption(caption: str) -> str: + """Apply the source order: special tokens, then Markdown, then newlines. + + ``<|tag value|>`` forms are rewritten to ``tag is value`` so a caption + cannot inject a real special token into the prompt. + """ + + def _special(match: re.Match[str]) -> str: + inner = match.group(1).strip() + parts = inner.split(None, 1) + return f"{parts[0]} is {parts[1]}" if len(parts) == 2 else inner + + text = _SPECIAL_TAG_RE.sub(_special, caption) + text = _remove_markdown_format(text) + # The source collapses runs of two or more newlines, not three or more; + # this shifts blank-line token positions and therefore the audio. + return re.sub(r"\n{2,}", "\n", text) + + +def _strip_text_after_leading_tags(text: str) -> str: + """Keep only the consecutive structural tags that start a line. + + Lyrics written on the same line as a ``[Verse]`` tag are dropped, which is + the documented contract: a tag must sit on its own line. + """ + if not text: + return text + output: list[str] = [] + for line in text.split("\n"): + match = _LEADING_TAGS_RE.match(line) + output.append(match.group(1).strip() if match else line) + return "\n".join(output) + + +def normalize_lyrics(lyrics: str) -> str: + """Normalize lyrics to the form the checkpoint was trained on.""" + text = _strip_text_after_leading_tags(lyrics) + text = text.replace("] ", "]\n") + text = text.replace(" [", "\n[") + text = text.replace(" ^ ", "\n") + text = re.sub(r"\[([^\]]+)\]", lambda match: f"[{match.group(1).lower()}]", text) + return f"[start]\n{text}" + + +def build_prompt(caption: str, lyrics: str) -> str: + """Build the conditioned prompt string.""" + return ( + "<|im_start|><|caption_start|>" + f"{clean_caption(caption)}" + "<|caption_end|><|lyrics_start|>" + f"{normalize_lyrics(lyrics)}" + "<|lyrics_end|><|im_end|><|audio_start|>" + ) + + +def build_cfg_null_token_ids(prompt_token_ids: list[int]) -> list[int]: + """Return the unconditioned twin's token ids for a conditioned prompt. + + Everything between the leading ``<|im_start|>`` and the trailing + ``<|im_end|><|audio_start|>`` becomes ``<|audio_cfg|>``. Length and the + audio-start position are preserved, which is what lets the two rows decode + in lockstep. + + Raises: + ValueError: If the prompt is too short to carry a caption and lyrics, + or does not have the expected framing tokens. + """ + if len(prompt_token_ids) < 4: + raise ValueError(f"MiniMax Music 3 prompt is too short to build a CFG twin: {len(prompt_token_ids)} tokens") + expected = ( + SPECIAL_TOKEN_IDS["<|im_start|>"], + SPECIAL_TOKEN_IDS["<|im_end|>"], + SPECIAL_TOKEN_IDS["<|audio_start|>"], + ) + actual = (prompt_token_ids[0], prompt_token_ids[-2], prompt_token_ids[-1]) + if actual != expected: + raise ValueError( + "MiniMax Music 3 prompt framing is wrong: expected " + f"(im_start, im_end, audio_start)={expected}, got {actual}" + ) + null_ids = list(prompt_token_ids) + null_ids[1:-2] = [SPECIAL_TOKEN_IDS["<|audio_cfg|>"]] * (len(null_ids) - 3) + return null_ids + + +def validate_tokenizer_ids(tokenizer: Any) -> None: + """Fail fast when a tokenizer is not the supported music tokenizer.""" + for token, expected in SPECIAL_TOKEN_IDS.items(): + token_id = tokenizer.convert_tokens_to_ids(token) + if token_id != expected: + raise ValueError(f"MiniMax Music 3 tokenizer mismatch for {token}: expected {expected}, got {token_id}") + + +__all__ = [ + "AUDIO_CODE_OFFSET", + "SPECIAL_TOKEN_IDS", + "build_cfg_null_token_ids", + "build_prompt", + "clean_caption", + "normalize_lyrics", + "validate_tokenizer_ids", +] diff --git a/vllm_omni/model_executor/models/minimax_music3/rvq_decoder.py b/vllm_omni/model_executor/models/minimax_music3/rvq_decoder.py new file mode 100644 index 000000000..64acc9b62 --- /dev/null +++ b/vllm_omni/model_executor/models/minimax_music3/rvq_decoder.py @@ -0,0 +1,179 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""RVQ depth decoder for MiniMax Music 3 (codebooks c1..c7). + +Each audio frame is eight codebooks deep. The backbone's own ``lm_head`` +emits ``c0``; this four-layer decoder walks ``c1..c7``, conditioning each code +on the previous ones. Its per-codebook hidden states are also the acoustic +stage's conditioning signal, so this module returns them alongside the codes. + +The sequence fed to the decoder is:: + + [proj(backbone_hidden), proj(embed(c0)), proj(embed(c1)), ...] + +and after step ``i`` the head ``audio_heads[i - 1]`` predicts ``c_i``. The +whole thing runs over at most eight positions, so it is written as a dense +causal block rather than anything with a KV cache. +""" + +from __future__ import annotations + +import torch +from torch import Tensor, nn +from torch.nn import functional as F +from torch.nn.attention import SDPBackend, sdpa_kernel + +from .constants import AUDIO_VOCAB_SIZE, BACKBONE_HIDDEN_SIZE, NUM_CODEBOOKS + +# cuDNN's fused attention is not bit-reproducible on this hardware: two +# identical replays of the same input can differ by one bfloat16 ULP, which is +# enough to flip a sampled code and send the rest of the song somewhere else. +# The choice is pinned per call rather than through a process-global switch so +# it cannot be undone by whatever else shares the process. +_SDPA_BACKENDS = [SDPBackend.FLASH_ATTENTION, SDPBackend.MATH] + + +class DepthDecoderBlock(nn.Module): + """Pre-norm transformer block with fused QKV and a SwiGLU MLP.""" + + def __init__(self, hidden_size: int, num_heads: int, intermediate_size: int) -> None: + super().__init__() + self.num_heads = num_heads + self.head_dim = hidden_size // num_heads + self.input_layernorm = nn.RMSNorm(hidden_size, eps=1e-6) + # Stored fused so the three projections are one matmul; the checkpoint + # ships them split and ``load_checkpoint_state`` concatenates. + self.qkv_proj = nn.Linear(hidden_size, hidden_size * 3, bias=False) + self.o_proj = nn.Linear(hidden_size, hidden_size, bias=False) + self.post_attention_layernorm = nn.RMSNorm(hidden_size, eps=1e-6) + self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False) + self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False) + self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False) + + def forward(self, x: Tensor) -> Tensor: + h = self.input_layernorm(x) + bsz, seq, hidden = h.shape + q, k, v = self.qkv_proj(h).chunk(3, dim=-1) + q = q.reshape(bsz, seq, self.num_heads, self.head_dim).transpose(1, 2) + k = k.reshape(bsz, seq, self.num_heads, self.head_dim).transpose(1, 2) + v = v.reshape(bsz, seq, self.num_heads, self.head_dim).transpose(1, 2) + with sdpa_kernel(_SDPA_BACKENDS): + h = F.scaled_dot_product_attention(q, k, v, is_causal=True) + h = h.transpose(1, 2).contiguous().reshape(bsz, seq, hidden) + x = x + self.o_proj(h) + h = self.post_attention_layernorm(x) + h = self.down_proj(F.silu(self.gate_proj(h)) * self.up_proj(h)) + return x + h + + +class RVQDepthDecoder(nn.Module): + """Depth-autoregressive decoder over the seven residual codebooks.""" + + def __init__( + self, + *, + hidden_size: int = BACKBONE_HIDDEN_SIZE, + num_layers: int = 4, + num_heads: int = 16, + intermediate_size: int = 6144, + audio_vocab_size: int = AUDIO_VOCAB_SIZE, + num_codebooks: int = NUM_CODEBOOKS, + max_seq_len: int = 16, + ) -> None: + super().__init__() + self.hidden_size = hidden_size + self.audio_vocab_size = audio_vocab_size + self.num_codebooks = num_codebooks + self.projection = nn.Linear(hidden_size, hidden_size, bias=False) + self.pos_embedding = nn.Embedding(max_seq_len, hidden_size) + self.layers = nn.ModuleList( + DepthDecoderBlock(hidden_size, num_heads, intermediate_size) for _ in range(num_layers) + ) + self.norm = nn.RMSNorm(hidden_size, eps=1e-6) + self.audio_heads = nn.ModuleList( + nn.Linear(hidden_size, audio_vocab_size, bias=False) for _ in range(num_codebooks - 1) + ) + + def forward(self, seq: Tensor) -> Tensor: + positions = torch.arange(seq.shape[1], device=seq.device) + x = seq + self.pos_embedding(positions).unsqueeze(0) + for layer in self.layers: + x = layer(x) + return self.norm(x) + + def load_checkpoint_state(self, state: dict[str, Tensor]) -> None: + """Load the component checkpoint, failing loudly on any mismatch. + + The checkpoint keeps ``q``/``k``/``v`` split and names attention + projections ``attn.to_*``; both are reshaped here. A silent partial + load would produce plausible-sounding but wrong audio, so every key is + accounted for and anything left over is an error. + + Args: + state: Tensors from the ``rvq_depth_decoder`` component, with the + ``audio_embeddings.weight`` entry already removed by the + caller. + + Raises: + ValueError: If any expected key is missing, any shape disagrees, + or the checkpoint carries keys this module does not consume. + """ + loaded: set[str] = set() + missing: list[str] = [] + mismatched: list[str] = [] + + def take(key: str) -> Tensor | None: + value = state.get(key) + if value is None: + missing.append(key) + return value + + def assign(key: str, target: Tensor, source: Tensor, keys: tuple[str, ...]) -> None: + if tuple(source.shape) != tuple(target.shape): + mismatched.append(f"{key}: checkpoint={tuple(source.shape)} model={tuple(target.shape)}") + return + target.data.copy_(source.to(dtype=target.dtype)) + loaded.update(keys) + + def copy_one(key: str, target: Tensor) -> None: + source = take(key) + if source is not None: + assign(key, target, source, (key,)) + + copy_one("projection.weight", self.projection.weight) + copy_one("pos_embedding.weight", self.pos_embedding.weight) + copy_one("norm.weight", self.norm.weight) + for index, head in enumerate(self.audio_heads): + copy_one(f"audio_heads.{index}.weight", head.weight) + + for index, layer in enumerate(self.layers): + prefix = f"layers.{index}." + copy_one(prefix + "input_layernorm.weight", layer.input_layernorm.weight) + copy_one( + prefix + "post_attention_layernorm.weight", + layer.post_attention_layernorm.weight, + ) + qkv_keys = tuple(prefix + f"attn.to_{name}.weight" for name in ("q", "k", "v")) + parts = [take(key) for key in qkv_keys] + if all(part is not None for part in parts): + assign( + prefix + "attn.qkv", + layer.qkv_proj.weight, + torch.cat(parts, dim=0), + qkv_keys, + ) + copy_one(prefix + "attn.to_out.weight", layer.o_proj.weight) + copy_one(prefix + "gate_proj.weight", layer.gate_proj.weight) + copy_one(prefix + "up_proj.weight", layer.up_proj.weight) + copy_one(prefix + "down_proj.weight", layer.down_proj.weight) + + unexpected = sorted(set(state) - loaded) + if missing or mismatched or unexpected: + raise ValueError( + "MiniMax Music 3 RVQ depth decoder failed to load: " + f"missing={missing[:8]}, mismatched={mismatched[:8]}, " + f"unexpected={unexpected[:8]}" + ) + + +__all__ = ["DepthDecoderBlock", "RVQDepthDecoder"] diff --git a/vllm_omni/model_executor/models/minimax_music3/sampling.py b/vllm_omni/model_executor/models/minimax_music3/sampling.py new file mode 100644 index 000000000..226addff5 --- /dev/null +++ b/vllm_omni/model_executor/models/minimax_music3/sampling.py @@ -0,0 +1,223 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Deterministic guided sampling for MiniMax Music 3. + +Two things happen per code, in this order: + +1. **Classifier-free guidance.** Every request decodes twice, once on the real + prompt and once on a prompt whose caption and lyrics have been replaced by + ``<|audio_cfg|>``. Codes are drawn from ``uncond + (cond - uncond) * 1.5``. + For ``c0`` the candidate set is additionally masked to the conditioned + row's own top 50, so the guided draw cannot wander outside what the + conditioned branch already considered plausible. +2. **Seeded top-k.** A batched, order-independent categorical draw: each row + gets its own RNG stream keyed by ``(seed, position)``, so a request's audio + depends only on its own seed and never on who else happens to share the + batch. + +The RNG is a MurmurHash3 x86_32 of ``(seed_lo, seed_hi, position, column)`` +turned into Gumbel noise, and the draw is an argmax over ``logits + gumbel``. +That is mathematically a categorical sample from ``softmax(logits)``, but +unlike ``torch.multinomial`` it needs no per-row generator state, so it is +CUDA-graph safe and reproducible across batch compositions. +""" + +from __future__ import annotations + +import torch +from torch import Tensor + +from .constants import AR_CFG_SCALE, AR_CFG_TOP_K, AR_TOP_K + +_UINT32_MASK = 0xFFFFFFFF +_UINT32_MAX = 4_294_967_295 + +# MurmurHash3 x86_32 constants. +_C1 = 0xCC9E2D51 +_C2 = 0x1B873593 +_R1 = 15 +_R2 = 13 +_M = 5 +_N = 0xE6546B64 + + +def _rotl32(x: Tensor, r: int) -> Tensor: + return ((x << r) | (x >> (32 - r))) & _UINT32_MASK + + +def _murmur3_mix(h: Tensor, k: Tensor) -> Tensor: + """Mix one 32-bit key into the running hash state.""" + k = (k * _C1) & _UINT32_MASK + k = _rotl32(k, _R1) + k = (k * _C2) & _UINT32_MASK + h = h ^ k + h = _rotl32(h, _R2) + return (h * _M + _N) & _UINT32_MASK + + +def _fmix32(h: Tensor) -> Tensor: + """MurmurHash3 finalizer.""" + h = h ^ (h >> 16) + h = (h * 0x85EBCA6B) & _UINT32_MASK + h = h ^ (h >> 13) + h = (h * 0xC2B2AE35) & _UINT32_MASK + return h ^ (h >> 16) + + +# Memoized ``k`` term of the column block's mix. The column block is the same +# ``arange(width)`` on every draw, so its three ops are hoisted out of the per +# draw work and kept per (device, width). Small and bounded: two widths exist, +# the c0 candidate set and the residual codebook. +_COLUMN_MIX_CACHE: dict[tuple[torch.device, int], Tensor] = {} + + +def _column_mix(num_columns: int, device: torch.device) -> Tensor: + """Return the mixed column block, ``[cols]`` int64 holding uint32.""" + key = (device, num_columns) + cached = _COLUMN_MIX_CACHE.get(key) + if cached is None: + columns = torch.arange(num_columns, dtype=torch.int64, device=device) & _UINT32_MASK + k = (columns * _C1) & _UINT32_MASK + k = _rotl32(k, _R1) + cached = (k * _C2) & _UINT32_MASK + _COLUMN_MIX_CACHE[key] = cached + return cached + + +def _absorb_mixed_key(h: Tensor, k: Tensor) -> Tensor: + """Absorb a key whose own three-op transform is already applied.""" + h = h ^ k + h = _rotl32(h, _R2) + return (h * _M + _N) & _UINT32_MASK + + +def row_hash_state(seeds: Tensor, positions: Tensor) -> Tensor: + """Hash state after the seed and position blocks, before the columns. + + Args: + seeds: ``[rows]`` per-request seeds. + positions: ``[rows]`` or ``[rows, draws]`` draw identifiers. The extra + dimension lets every codebook of a frame be hashed in one batch + instead of one dispatch per codebook; the arithmetic is elementwise + integer, so a batched call is bit-identical to the loop it replaces. + + Returns: + A tensor shaped like ``positions``, int64 holding uint32. + """ + seeds = seeds.to(torch.int64) + while seeds.dim() < positions.dim(): + seeds = seeds.unsqueeze(-1) + h = torch.zeros((), dtype=torch.int64, device=seeds.device) + h = _murmur3_mix(h, seeds & _UINT32_MASK) + h = _murmur3_mix(h, (seeds >> 32) & _UINT32_MASK) + return _murmur3_mix(h, positions.to(torch.int64) & _UINT32_MASK) + + +def murmur_hash32(seeds: Tensor, positions: Tensor, columns: Tensor) -> Tensor: + """Hash ``(seed, position, column)`` to a uint32 held in int64. + + The 64-bit seed is consumed as two 32-bit blocks, then the position and + the column, for a hashed length of 16 bytes. + + Args: + seeds: ``[rows]`` per-request seeds. + positions: ``[rows]`` draw identifiers, unique per (frame, codebook). + columns: ``[cols]`` candidate indices, which must be ``arange(cols)``. + + Returns: + ``[rows, cols]`` int64 tensor holding uint32 values. + """ + h = row_hash_state(seeds, positions).unsqueeze(-1) + return _fmix32(_absorb_mixed_key(h, _column_mix(int(columns.shape[0]), h.device)) ^ 16) + + +def gumbel_noise(seeds: Tensor, positions: Tensor, num_columns: int) -> Tensor: + """Return the Gumbel noise for one draw per position, ``[..., cols]``. + + Args: + seeds: ``[rows]`` per-request seeds. + positions: ``[rows]`` or ``[rows, draws]`` draw identifiers. + num_columns: Candidate count, i.e. the width of the logits. + + Returns: + float64 noise shaped ``positions.shape + (num_columns,)``. + """ + h = row_hash_state(seeds, positions).unsqueeze(-1) + hashed = _fmix32(_absorb_mixed_key(h, _column_mix(num_columns, h.device)) ^ 16) + # float64 throughout: the double log is numerically unstable in float32 and + # a single flipped draw changes every frame that follows it. + uniform = hashed.to(torch.float64) / _UINT32_MAX + # Clamp both ends. uniform == 1 gives +inf Gumbel, which turns a -inf + # masked entry into NaN; the upper cap is the hash spacing, so a saturated + # bucket ties with its neighbour instead of always winning. + gumbel = uniform.log().clamp_(min=torch.finfo(torch.float64).min, max=-(2.0**-32)).neg_() + return gumbel.log().neg_() + + +def draw_from_gumbel(logits: Tensor, gumbel: Tensor, *, top_k: int = AR_TOP_K) -> Tensor: + """Argmax of ``top_k``-masked logits plus precomputed Gumbel noise. + + Args: + logits: ``[rows, vocab]`` unnormalized scores. + gumbel: ``[rows, vocab]`` float64 noise from :func:`gumbel_noise`. + top_k: Candidates retained per row before the draw. + + Returns: + ``[rows]`` sampled column indices. + """ + values = torch.nan_to_num(logits.float(), nan=-1e9, posinf=1e9, neginf=-1e9) + k = min(int(top_k), values.shape[-1]) + threshold = torch.topk(values, k, dim=-1).values[..., -1, None] + values = values.masked_fill(values < threshold, -float("inf")) + return (values.to(torch.float64) + gumbel).argmax(dim=-1) + + +def sample_topk_seeded( + logits: Tensor, + seeds: Tensor, + positions: Tensor, + *, + top_k: int = AR_TOP_K, +) -> Tensor: + """Draw one column per row from its own top-k, under its own RNG stream. + + Args: + logits: ``[rows, vocab]`` unnormalized scores. + seeds: ``[rows]`` per-request sampling seeds. + positions: ``[rows]`` draw identifiers, unique per (frame, codebook). + top_k: Candidates retained per row before the draw. + + Returns: + ``[rows]`` sampled column indices. + """ + gumbel = gumbel_noise(seeds, positions, int(logits.shape[-1])) + return draw_from_gumbel(logits, gumbel, top_k=top_k) + + +def apply_c0_cfg(cond: Tensor, uncond: Tensor, *, scale: float = AR_CFG_SCALE) -> Tensor: + """Guide ``cond`` by ``uncond``, restricted to ``cond``'s own top candidates. + + The mask is a second, separate restriction from the top-k the sampler + applies afterwards. Dropping it widens the distribution the model draws + from and audibly loosens how closely the output follows the caption. + """ + guided = uncond + (cond - uncond) * scale + k = min(AR_CFG_TOP_K, cond.shape[-1]) + threshold = torch.topk(cond, k, dim=-1).values[..., -1, None] + return guided.masked_fill(cond < threshold, -float("inf")) + + +def apply_depth_cfg(cond: Tensor, uncond: Tensor, *, scale: float = AR_CFG_SCALE) -> Tensor: + """Guide the depth heads. Unlike ``c0`` these are not additionally masked.""" + return uncond + (cond - uncond) * scale + + +__all__ = [ + "apply_c0_cfg", + "apply_depth_cfg", + "draw_from_gumbel", + "gumbel_noise", + "murmur_hash32", + "row_hash_state", + "sample_topk_seeded", +] diff --git a/vllm_omni/model_executor/models/minimax_music3/staging.py b/vllm_omni/model_executor/models/minimax_music3/staging.py new file mode 100644 index 000000000..ed3eb8a47 --- /dev/null +++ b/vllm_omni/model_executor/models/minimax_music3/staging.py @@ -0,0 +1,66 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Pinned host staging for the small per-step index vectors. + +Building a device tensor from a Python list (``torch.tensor(rows, +device="cuda")``) stages through pageable host memory, which makes the copy +synchronous: the host blocks until the copy retires, and because the copy is +queued behind the step's kernels that means blocking until the whole forward +has finished. On this model that turned four one-line conveniences into the +single largest entry in a CPU profile of the decode loop. + +Writing into a pinned block and pushing it with one non-blocking copy costs +one transfer for the whole step and leaves the host free to keep enqueueing. +""" + +from __future__ import annotations + +from collections.abc import Sequence + +import torch + +__all__ = ["PinnedIndexStage"] + + +class PinnedIndexStage: + """A fixed ``[rows, capacity]`` int64 block staged host-to-device. + + The device view returned by :meth:`push` is only populated once the stream + reaches the copy, so it may be consumed by kernels on that stream and by + nothing else. The host side is double-buffered: the next step writes the + slot the device is not reading, so refilling cannot corrupt a copy that is + still in flight. + """ + + def __init__(self, rows: int, capacity: int, device: torch.device) -> None: + self._capacity = capacity + self._host = [torch.zeros((rows, capacity), dtype=torch.int64, pin_memory=True) for _ in range(2)] + self._device = torch.zeros((rows, capacity), dtype=torch.int64, device=device) + self._slot = 0 + + @property + def capacity(self) -> int: + return self._capacity + + def push(self, values: Sequence[Sequence[int]], count: int) -> torch.Tensor: + """Stage ``count`` entries of each row and return the device view. + + Args: + values: One sequence of Python ints per row of the block. + count: How many entries of each row are meaningful this step. + + Returns: + ``[rows, count]`` int64 view of the device block. + + Raises: + ValueError: If ``count`` exceeds the reserved capacity. + """ + if count > self._capacity: + raise ValueError(f"staged {count} entries but only {self._capacity} were reserved") + self._slot ^= 1 + host = self._host[self._slot] + view = host.numpy() + for index, row in enumerate(values): + view[index, :count] = row + self._device.copy_(host, non_blocking=True) + return self._device[:, :count] diff --git a/vllm_omni/model_executor/models/minimax_music3/talker.py b/vllm_omni/model_executor/models/minimax_music3/talker.py new file mode 100644 index 000000000..c4fb876fa --- /dev/null +++ b/vllm_omni/model_executor/models/minimax_music3/talker.py @@ -0,0 +1,990 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Stage-0 talker for MiniMax Music 3. + +A Qwen3 backbone predicts one audio frame per decode step. Each frame is +eight RVQ codebooks deep: the backbone's own ``lm_head`` emits ``c0`` and the +four-layer depth decoder walks ``c1..c7``. The conditioning signal handed to +the acoustic stage is the backbone hidden state concatenated with the seven +depth hidden states, so 8 x 4096 = 32768 dimensions per frame. + +Two things make this model unusual inside vLLM: + +* **Every request decodes as two rows.** Classifier-free guidance is not + optional here. The conditioned row sees the real prompt and the + unconditioned row sees the same prompt with the caption-and-lyrics span + replaced by ``<|audio_cfg|>``. Codes are drawn from the guided blend and + then fed back to *both* rows, so the pair stays token-identical. Rows are + paired by ``cfg_pair_id`` rather than by adjacency, because batch position + is not stable across engine steps. +* **The decode input is not a token embedding.** After prefill, each step's + input is the composed embedding of the previous frame's eight codes. The + sampled ``c0`` token id still flows through vLLM's normal machinery so that + stop handling and accounting work, but the embedding it would produce is + replaced before the backbone runs. +""" + +from __future__ import annotations + +from collections.abc import Iterable +from dataclasses import dataclass +from typing import Any + +import torch +from torch import nn +from vllm.config import VllmConfig +from vllm.logger import init_logger +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead +from vllm.model_executor.models.qwen3 import Qwen3Model +from vllm.model_executor.models.utils import AutoWeightsLoader +from vllm.v1.outputs import SamplerOutput + +from vllm_omni.model_executor.models.output_templates import OmniOutput + +from .chunking import chunk_windows +from .constants import ( + AR_CHUNK_HOP_FRAMES, + AUDIO_VOCAB_SIZE, + BACKBONE_HIDDEN_SIZE, + C0_VOCAB_SIZE, + DEFAULT_MAX_AUDIO_FRAMES, + NUM_CODEBOOKS, +) +from .depth_graph import DepthGraphCache +from .frame_state import RequestARState +from .prompt import AUDIO_CODE_OFFSET, SPECIAL_TOKEN_IDS +from .rvq_decoder import RVQDepthDecoder +from .sampling import ( + apply_c0_cfg, + apply_depth_cfg, + draw_from_gumbel, + gumbel_noise, + sample_topk_seeded, +) +from .staging import PinnedIndexStage +from .weights import load_depth_decoder_weights + +logger = init_logger(__name__) + +__all__ = ["MiniMaxMusic3TalkerForConditionalGeneration"] + +_COND = "cond" +_UNCOND = "uncond" +_STOP_TOKEN = SPECIAL_TOKEN_IDS["<|audio_end|>"] +# Emitted on the step a song ends, to keep the request alive for one more step +# while its terminal span is delivered. It is never fed back as audio: the +# state is already marked finished, so no frame is composed from it. +_HOLD_TOKEN = AUDIO_CODE_OFFSET + +# The unconditioned twin's request id is the conditioned one plus this suffix, +# which is what lets a pair be recognized without a serving-layer injection. +_CFG_UNCOND_SUFFIX = "__cfg_uncond" + + +@dataclass(frozen=True) +class _RowConstants: + """Per-request values that never change once the request is admitted.""" + + request_id: str + role: str + pair_id: str + seed: int + max_frames: int + + +def _request_key(value: Any) -> str: + """Normalize a request id to the form used to key per-request state. + + The engine injects ``global_request_id`` as a single-element list, while + ``on_requests_finished`` receives the bare string. Keying state on the + unwrapped value keeps allocation and release on the same key; getting this + wrong leaks a song's worth of frames per request. + """ + if isinstance(value, (list, tuple)): + return str(value[0]) if value else "" + return "" if value is None else str(value) + + +def _pair_id_from_request(request_id: str) -> str: + """Derive the CFG pair id when the request does not carry one.""" + if request_id.endswith(_CFG_UNCOND_SUFFIX): + return request_id[: -len(_CFG_UNCOND_SUFFIX)] + return request_id + + +def _role_from_request(request_id: str) -> str: + """Derive the CFG role from the companion suffix.""" + return _UNCOND if request_id.endswith(_CFG_UNCOND_SUFFIX) else _COND + + +def _first_int(value: Any, default: int) -> int: + """Read an int that the serving layer wraps in a single-element list.""" + if isinstance(value, (list, tuple)): + value = value[0] if value else None + if value is None: + return default + try: + parsed = int(value) + except (TypeError, ValueError): + return default + return parsed if parsed > 0 else default + + +class MiniMaxMusic3TalkerForConditionalGeneration(nn.Module): + """Qwen3 backbone plus the audio embedding table and RVQ depth decoder.""" + + # Sampling, feedback and window emission are all model-owned. + prefer_model_sampler: bool = True + have_multimodal_outputs: bool = True + has_postprocess: bool = True + skips_model_sampler_output_token_history: bool = True + postprocess_uses_hidden_states: bool = False + postprocess_uses_multimodal_outputs: bool = False + postprocess_uses_req_infos: bool = True + supports_omni_query_start_loc: bool = True + # A row still completing a chunked prefill reaches the sampler but its + # token is discarded; drawing a frame for it would corrupt the song. + requires_request_sample_eligibility: bool = True + # The acoustic stage consumes the composed 32768-d frame, never the raw + # backbone hidden state, so suppress the runner's automatic "hidden" key. + omni_pooler_payload_include_hidden: bool = False + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + self.vllm_config = vllm_config + config = vllm_config.model_config.hf_config + self.config = config + hidden_size = int(getattr(config, "hidden_size", BACKBONE_HIDDEN_SIZE)) + + self.model = Qwen3Model( + vllm_config=vllm_config, + prefix=f"{prefix}.model" if prefix else "model", + ) + if getattr(config, "tie_word_embeddings", False): + self.lm_head = self.model.embed_tokens + else: + self.lm_head = ParallelLMHead( + int(config.vocab_size), + hidden_size, + prefix=f"{prefix}.lm_head" if prefix else "lm_head", + ) + self.logits_processor = LogitsProcessor(int(config.vocab_size)) + + # c1..c7 share one table; codebook i occupies rows [i*1024, (i+1)*1024). + self.audio_embeddings = nn.Embedding(AUDIO_VOCAB_SIZE * (NUM_CODEBOOKS - 1), hidden_size) + self.rvq_decoder = RVQDepthDecoder(hidden_size=hidden_size) + + self.register_buffer( + "audio_embedding_offsets", + (torch.arange(NUM_CODEBOOKS - 1) * AUDIO_VOCAB_SIZE).unsqueeze(0), + persistent=False, + ) + # The only token ids c0 sampling can return, in vocabulary order: + # <|audio_end|> first, then the c0 code range. Column 0 therefore means + # stop and column code+1 means code. + self.register_buffer( + "c0_logit_ids", + torch.cat( + ( + torch.tensor([_STOP_TOKEN], dtype=torch.long), + torch.arange( + AUDIO_CODE_OFFSET, + AUDIO_CODE_OFFSET + C0_VOCAB_SIZE, + dtype=torch.long, + ), + ) + ), + persistent=False, + ) + self.frame_embedding_scale = float(NUM_CODEBOOKS) ** -0.5 + + scheduler_config = getattr(vllm_config, "scheduler_config", None) + # Guidance doubles the rows, so the batch is twice the admitted + # request count. + self._max_rows = max(64, 2 * int(getattr(scheduler_config, "max_num_seqs", 16))) + + # Per-step scratch, rebuilt every step from request-keyed state. + # Nothing here may survive a step: vLLM compacts and reorders the + # persistent batch, so a row index is only meaningful within the step + # that produced it. + self._last_logits_hidden: torch.Tensor | None = None + self._step_engine_ids: list[str] = [] + self._row_request_ids: list[str] = [] + self._row_roles: list[str] = [] + self._row_pair_ids: list[str] = [] + self._row_seeds: list[int] = [] + self._row_max_frames: list[int] = [] + self._row_prompt_tokens: list[int] = [] + self._row_sample_eligible: list[bool] = [] + self._feedback_positions: torch.Tensor | None = None + self._feedback_codes: torch.Tensor | None = None + + # Per-step index vectors are staged through pinned memory rather than + # built with ``torch.tensor(..., device="cuda")``; see + # ``staging.PinnedIndexStage`` for why that matters here. Allocated on + # first use because the device is not known at construction time. + # Rows are (cond_row, uncond_row, seed, position). + self._sample_stage: PinnedIndexStage | None = None + # One row, the flat token offsets of this step's decoding rows. + self._feedback_stage: PinnedIndexStage | None = None + + # Depth-loop draw offsets, ``[1..NUM_CODEBOOKS)``. Held as a buffer so + # the per-step position arithmetic is one broadcast add rather than one + # scalar add per codebook. + self.register_buffer( + "depth_draw_offsets", + torch.arange(1, NUM_CODEBOOKS, dtype=torch.int64), + persistent=False, + ) + + # Frame feedback slab. A captured decode graph replays fixed addresses + # and fixed shapes, so the rows to overwrite cannot be expressed as a + # variable-length ``index_copy`` against a freshly built index tensor. + # Instead every row of the reserved slab carries a codes entry and a + # flag, both written in place before the graph runs, and the forward + # selects with ``torch.where`` over the whole row set. Sized to the + # widest decode batch; wider prefill batches take the eager path in + # ``_apply_frame_feedback``, and prefill is never graph-replayed. + self.register_buffer( + "frame_feedback_codes", + torch.zeros((self._max_rows, NUM_CODEBOOKS), dtype=torch.long), + persistent=False, + ) + self.register_buffer( + "frame_feedback_active", + torch.zeros(self._max_rows, dtype=torch.bool), + persistent=False, + ) + + # Cross-step state. Keyed by CFG pair id, because a guided request is + # two engine rows sharing one song. The engine renames requests to + # internal UUIDs on admission, so the pair id is derived from the + # user-facing id and the engine ids are mapped onto it separately. + self._states: dict[str, RequestARState] = {} + self._pair_of_engine_id: dict[str, str] = {} + # Per-request constants, keyed by engine id. ``forward`` is the only + # hook the engine hands the sampling extra args to, but its Python body + # does not run on a step replayed from a CUDA graph, so what it learns + # is cached here and re-bound to rows in ``prepare_runner_inputs``. + self._row_constants: dict[str, _RowConstants] = {} + self._row_binding_complete = False + self._incomplete_steps = 0 + + # The depth loop is recorded per pair count and replayed after that. + self._depth_graphs = DepthGraphCache(self._depth_decode_eager) + + # ---------------------------------------------------------------- forward + + def prepare_runner_inputs( + self, + *, + req_ids: list[str], + num_computed_tokens: Any, + num_scheduled_tokens: Any, + input_ids: torch.Tensor, + positions: torch.Tensor, + **_: Any, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Bind this step's batch rows to requests, before the forward runs. + + This is the only hook that sees the authoritative row-to-request + mapping for the step about to execute. vLLM compacts and reorders the + persistent batch whenever a request finishes or is preempted, so any + state indexed by row number from a previous step is stale by the time + it is read. Everything the model carries across steps is therefore + keyed by request id and gathered into row order here. + + It also resolves which rows are decoding. A row is decoding when its + prompt is already computed and it is advancing by a single token; a + row still working through a chunked prefill is not, and must keep its + real prompt embeddings. + + Python runs here on every step, including steps whose forward is + replayed from a CUDA graph rather than executed, so this is also where + the row-to-request binding that ``sample`` reads has to be rebuilt. + """ + self._step_engine_ids = [str(rid) for rid in req_ids] + computed = [int(value) for value in num_computed_tokens] + scheduled = [int(value) for value in num_scheduled_tokens] + # Prompt length for a row that is completing its prefill on this step. + # Only meaningful when the row is about to produce its first logit, + # which is the only moment state is created. + self._row_prompt_tokens = [computed[row] + scheduled[row] for row in range(len(req_ids))] + + # Flat token offset of each row's first scheduled token. + offsets: list[int] = [] + running = 0 + for count in scheduled: + offsets.append(running) + running += count + + feedback_rows: list[int] = [] + feedback_codes: list[torch.Tensor] = [] + for row, engine_id in enumerate(self._step_engine_ids): + # Both members of a guided pair are fed the same frame, so state + # is keyed by pair rather than by request. + pair_id = self._pair_of_engine_id.get(engine_id) + state = self._states.get(pair_id) if pair_id else None + if state is None or state.last_codes is None: + continue + # A decode row advances by exactly one token past a complete + # prompt. A row still working through a chunked prefill keeps its + # real prompt embeddings. + if scheduled[row] != 1 or computed[row] < state.prompt_tokens: + continue + feedback_rows.append(offsets[row]) + feedback_codes.append(state.last_codes) + + self._write_frame_feedback(feedback_rows, feedback_codes, input_ids.device) + self._bind_rows(computed, scheduled) + return input_ids, positions + + def _write_frame_feedback( + self, + feedback_rows: list[int], + feedback_codes: list[torch.Tensor], + device: torch.device, + ) -> None: + """Publish this step's frame feedback for the forward to consume. + + Two forms are written, because two forwards can consume it. The slab + is what a captured decode graph reads: fixed address, fixed shape, + selected row by row with a flag. The loose index and code tensors are + the eager fallback for a batch wider than the slab, which is only ever + a prefill batch and therefore never a replayed one. + """ + if self._feedback_stage is None: + self._feedback_stage = PinnedIndexStage(1, self._max_rows, device) + + active = self.frame_feedback_active + if not feedback_rows: + self._feedback_positions = None + self._feedback_codes = None + active.zero_() + return + + codes = torch.cat(feedback_codes, dim=0).to(device) + if max(feedback_rows) < self._max_rows: + positions = self._feedback_stage.push([feedback_rows], len(feedback_rows))[0] + self.frame_feedback_codes.index_copy_(0, positions, codes) + active.zero_() + active.index_fill_(0, positions, True) + else: + # Too wide for the slab. Nothing reads it on this step, but a stale + # flag would corrupt the next graph replay if the batch shrank. + positions = torch.tensor(feedback_rows, dtype=torch.long, device=device) + active.zero_() + self._feedback_positions = positions + self._feedback_codes = codes + + def _bind_rows(self, computed: list[int], scheduled: list[int]) -> None: + """Rebuild the row-order request metadata from the per-request cache. + + ``forward`` learns these values from the engine's own kwargs, but its + Python body is skipped whenever the step is replayed from a CUDA graph, + so the binding is re-established here. A row whose request has not been + seen by ``forward`` yet leaves the binding incomplete; that only + happens on the step a request is admitted, which is a prefill step, and + prefill is never replayed, so ``forward`` fills it in before the + sampler runs. + """ + constants = [self._row_constants.get(engine_id) for engine_id in self._step_engine_ids] + if not constants or any(entry is None for entry in constants): + self._row_binding_complete = False + return + self._row_request_ids = [entry.request_id for entry in constants] + self._row_roles = [entry.role for entry in constants] + self._row_pair_ids = [entry.pair_id for entry in constants] + self._row_seeds = [entry.seed for entry in constants] + self._row_max_frames = [entry.max_frames for entry in constants] + + # A row is eligible to draw a frame once it has computed its whole + # prompt. Mirrors the engine's own rule; the prompt length is taken + # from the pair's state, which exists for every row whose request has + # already produced a logit. + eligible: list[bool] = [] + for row, entry in enumerate(constants): + state = self._states.get(entry.pair_id) + if state is None: + self._row_binding_complete = False + return + eligible.append(computed[row] + scheduled[row] >= state.prompt_tokens) + self._row_sample_eligible = eligible + self._row_binding_complete = True + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: Any | None = None, + inputs_embeds: torch.Tensor | None = None, + **kwargs: Any, + ) -> torch.Tensor: + del intermediate_tensors + self._capture_row_metadata(kwargs) + + if inputs_embeds is None: + hidden_states = self.model.embed_tokens(input_ids) + hidden_states = self._apply_frame_feedback(hidden_states) + else: + hidden_states = inputs_embeds + + residual: torch.Tensor | None = None + for layer in self.model.layers: + hidden_states, residual = layer(positions, hidden_states, residual) + norm_out = self.model.norm(hidden_states, residual) + if isinstance(norm_out, tuple): + norm_out = norm_out[0] + return norm_out + + def _capture_row_metadata(self, kwargs: dict[str, Any]) -> None: + """Read this step's per-row CFG role, pair id and seed. + + Row order here matches ``prepare_runner_inputs``: both are built from + ``input_batch.req_ids``. The request ids resolved there are preferred, + and the info dicts are only a fallback for paths that do not route + through that hook. + """ + if self._row_binding_complete: + # Every row was resolved from the per-request cache before the + # forward, so there is nothing here the sampler still needs. Bail + # out rather than rebuild it: this runs on every eager decode step. + return + infos = kwargs.get("model_intermediate_buffer") + if infos is None: + infos = kwargs.get("runtime_additional_information") + extra_args = kwargs.get("sampling_extra_args") + eligible = kwargs.get("request_sample_eligible") + self._row_sample_eligible = [bool(flag) for flag in eligible] if eligible else [] + if infos is None and extra_args is None: + # Warm-up dummy run: no omni kwargs are supplied. + self._row_pair_ids = [] + return + + infos = infos or [] + extra_args = extra_args or [] + rows = max(len(infos), len(extra_args), len(self._step_engine_ids)) + request_ids: list[str] = [] + roles: list[str] = [] + pair_ids: list[str] = [] + seeds: list[int] = [] + max_frames: list[int] = [] + for index in range(rows): + info = infos[index] if index < len(infos) else {} + info = info if isinstance(info, dict) else {} + extra = extra_args[index] if index < len(extra_args) else {} + extra = extra if isinstance(extra, dict) else {} + # The user-facing id. The engine renames a request to an internal + # UUID on admission, so this is the only id that both the info + # dicts and postprocess agree on, and the only one that carries + # the companion suffix a pair is recognized by. + request_id = _request_key(info.get("global_request_id")) or f"row-{index}" + pair_id = str(extra.get("cfg_pair_id") or _pair_id_from_request(request_id)) + request_ids.append(request_id) + roles.append(str(extra.get("cfg_role") or _role_from_request(request_id))) + pair_ids.append(pair_id) + seeds.append(int(extra.get("tts_local_seed") or 0)) + max_frames.append(_first_int(info.get("max_audio_frames"), DEFAULT_MAX_AUDIO_FRAMES)) + # Remember which pair each engine-side row belongs to, so the next + # step's feedback can be gathered before the info dicts are seen, + # and cache the whole row constant so a graph-replayed step can be + # bound without the info dicts at all. + if index < len(self._step_engine_ids): + engine_id = self._step_engine_ids[index] + self._pair_of_engine_id[engine_id] = pair_id + self._row_constants[engine_id] = _RowConstants( + request_id=request_id, + role=roles[index], + pair_id=pair_id, + seed=seeds[index], + max_frames=max_frames[index], + ) + self._row_request_ids = request_ids + self._row_roles = roles + self._row_pair_ids = pair_ids + self._row_seeds = seeds + self._row_max_frames = max_frames + + def _apply_frame_feedback(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Replace each decoding row's token embedding with its frame embedding. + + The rows to replace and the codes to write were resolved in + ``prepare_runner_inputs`` against this step's own batch, so a mixed + prefill and decode step touches only the decoding rows and a prompt + token is never overwritten. + + Which of the two forms is read depends only on the row count, which is + a constant inside a captured graph: a decode batch fits the reserved + slab and is selected with ``torch.where`` over fixed buffers, while a + prefill batch wider than the slab falls back to a variable-length + ``index_copy``. Prefill is never replayed, so the fallback never has to + be graph-safe. + """ + rows = hidden_states.shape[0] + if rows <= self._max_rows: + frame_embeds = self.embed_frames(self.frame_feedback_codes[:rows]).to(dtype=hidden_states.dtype) + return torch.where( + self.frame_feedback_active[:rows].unsqueeze(1), + frame_embeds, + hidden_states, + ) + positions = self._feedback_positions + codes = self._feedback_codes + if positions is None or codes is None: + return hidden_states + frame_embeds = self.embed_frames(codes).to(dtype=hidden_states.dtype) + return hidden_states.index_copy(0, positions, frame_embeds) + + def embed_frames(self, codes: torch.Tensor) -> torch.Tensor: + """Compose one conditioning embedding per row from all eight codebooks. + + Args: + codes: ``[rows, 8]`` codebook indices, ``c0`` zero-based. + + Returns: + ``[rows, hidden]`` in the embedding table's dtype. + """ + c0 = self.model.embed_tokens(codes[:, 0] + AUDIO_CODE_OFFSET) + extra = self.audio_embeddings(codes[:, 1:] + self.audio_embedding_offsets) + return (c0 + extra.sum(dim=1).to(c0.dtype)) * self.frame_embedding_scale + + def compute_logits(self, hidden_states: torch.Tensor, sampling_metadata: Any = None) -> torch.Tensor: + # sample() has no other route to the backbone hidden state. + self._last_logits_hidden = hidden_states + return self.logits_processor(self.lm_head, hidden_states) + + def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor: + return self.model.embed_tokens(input_ids) + + # ----------------------------------------------------------------- sample + + def sample(self, logits: torch.Tensor, sampling_metadata: Any) -> SamplerOutput | None: + """Draw one frame per conditioned request and feed it back to both rows.""" + del sampling_metadata + hidden = self._last_logits_hidden + self._last_logits_hidden = None + if hidden is None or not self._row_pair_ids: + return None + + rows = int(logits.shape[0]) + pairs = self._resolve_pairs(rows) + if not pairs: + # Nothing to draw this step. Emit a benign audio code so every row + # stays alive until its partner is admitted; returning None would + # hand the rows to the stock sampler, whose token would be read as + # a real code. + held = torch.full((rows, 1), _HOLD_TOKEN, dtype=torch.long, device=logits.device) + return SamplerOutput(sampled_token_ids=held, logprobs_tensors=None) + + states = [self._state_for(self._row_pair_ids[pair[0]], pair[0]) for pair in pairs] + cond_rows, uncond_rows, seeds, positions = self._stage_step_vectors( + [pair[0] for pair in pairs], + [pair[1] for pair in pairs], + [state.seed for state in states], + [state.advance_position() for state in states], + logits.device, + ) + + narrowed = logits.index_select(1, self.c0_logit_ids).float() + c0_logits = apply_c0_cfg(narrowed.index_select(0, cond_rows), narrowed.index_select(0, uncond_rows)) + drawn = sample_topk_seeded(c0_logits, seeds, positions) + stopped = drawn == 0 + c0 = (drawn - 1).clamp_min(0) + + codes, depth_hidden = self._depth_decode( + hidden.index_select(0, cond_rows), + hidden.index_select(0, uncond_rows), + c0, + seeds, + positions, + ) + frame_hidden = torch.cat((hidden.index_select(0, cond_rows), depth_hidden), dim=-1) + + stopped_flags = stopped.tolist() + drain_rows: list[int] = [] + for index, (cond_row, uncond_row) in enumerate(pairs): + state = states[index] + if state.finished: + # The drain step. The terminal span written last step is on + # its way out in this step's payload, so the pair can stop now. + state.pending_stop = False + drain_rows.extend((index, cond_row, uncond_row)) + continue + if stopped_flags[index]: + state.finished = True + state.pending_stop = True + state.finish_reason = "stop" + state.last_codes = None + continue + # Both rows of the pair are fed the same frame next step; the + # codes live on the request, not on a row, because the row this + # request occupies may change before then. + state.last_codes = codes[index : index + 1] + state.frames.append(frame_hidden[index : index + 1]) + state.codes.append(codes[index : index + 1]) + state.generated_frames += 1 + if state.generated_frames >= state.max_frames: + state.finished = True + state.pending_stop = True + state.finish_reason = "length" + del cond_row, uncond_row + + # Rows that resolved into no pair keep this value. A hold token rather + # than id 0, which is an arbitrary text token: a row waiting for its + # partner should wait, not emit something the vocabulary reads as + # prose. It is never fed back as audio, because feedback is gated on + # the pair's own state. + token_ids = torch.full((rows, 1), _HOLD_TOKEN, dtype=torch.long, device=logits.device) + sampled = self.c0_logit_ids[drawn] + token_ids[cond_rows, 0] = sampled + token_ids[uncond_rows, 0] = sampled + + # The stop token is deferred by exactly one step. A pair that ends this + # step must stay alive so that the terminal span ``postprocess`` is + # about to write can ride out on the next step's payload; a pair that + # already published its tail stops now. + # + # This covers the two ways a song ends on its own: the audio-end token, + # and the request's frame budget, which reaches the model as + # ``max_audio_frames``. It does not cover the engine ending a request + # underneath the model, because the payload for a step is fixed before + # the model is told. Hitting ``max_model_len`` mid-song therefore drops + # the final window. In practice the frame budget is reached first; a + # request can only hit the context limit with a prompt long enough that + # the two limits collide, which the deploy config documents. + for index, (cond_row, uncond_row) in enumerate(pairs): + state = states[index] + if state.pending_stop: + token_ids[cond_row, 0] = _HOLD_TOKEN + token_ids[uncond_row, 0] = _HOLD_TOKEN + for _, cond_row, uncond_row in zip(drain_rows[0::3], drain_rows[1::3], drain_rows[2::3], strict=True): + token_ids[cond_row, 0] = _STOP_TOKEN + token_ids[uncond_row, 0] = _STOP_TOKEN + return SamplerOutput(sampled_token_ids=token_ids, logprobs_tensors=None) + + def _stage_step_vectors( + self, + cond_rows: list[int], + uncond_rows: list[int], + seeds: list[int], + positions: list[int], + device: torch.device, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Move this step's four small int64 vectors to the device in one copy. + + The copy is asynchronous, so the four returned views are only valid + once the stream reaches them. Every consumer is a kernel on that same + stream, which is what makes this safe. + """ + if self._sample_stage is None: + self._sample_stage = PinnedIndexStage(4, self._max_rows, device) + block = self._sample_stage.push([cond_rows, uncond_rows, seeds, positions], len(cond_rows)) + return block[0], block[1], block[2], block[3] + + def _resolve_pairs(self, rows: int) -> list[tuple[int, int]]: + """Map ``cfg_pair_id`` to its ``(cond_row, uncond_row)`` batch indices. + + Guidance is mandatory for this checkpoint, so an incomplete pair is a + hard error rather than something to route around. Skipping the step + would leave the conditioned request producing no frames and the client + waiting forever, which is the failure shape this module has been bitten + by before; raising surfaces it to the caller instead. + + Raises: + RuntimeError: If any pair is missing a member in this step. + """ + by_pair: dict[str, dict[str, int]] = {} + limit = min(rows, len(self._row_pair_ids)) + for row in range(limit): + # A row still finishing a chunked prefill is handed to the sampler + # but its token is discarded, so drawing a frame for it would add + # a bogus frame and desync the seeded RNG. + if row < len(self._row_sample_eligible) and not self._row_sample_eligible[row]: + continue + by_pair.setdefault(self._row_pair_ids[row], {})[self._row_roles[row]] = row + + pairs: list[tuple[int, int]] = [] + incomplete: list[str] = [] + for pair_id, roles in by_pair.items(): + if _COND in roles and _UNCOND in roles: + pairs.append((roles[_COND], roles[_UNCOND])) + else: + incomplete.append(f"{pair_id}({'+'.join(sorted(roles))})") + if incomplete: + # A member can be absent for a step while its partner is still + # being admitted. Holding the present row costs one wasted decode + # step; raising would abort the whole engine step for every other + # request in the batch, so the lone row simply does not sample this + # step and waits. A pair that never converges is bounded by the + # request's own token budget rather than hanging forever. + self._incomplete_steps += 1 + if self._incomplete_steps in (1, 64, 512): + logger.warning( + "MiniMax Music 3 waiting on incomplete CFG pair(s) %s " + "(step %d). If this does not resolve, the companion was " + "never enqueued: check that the stage declares " + "extra_args.cfg_role in its deploy config.", + incomplete[:4], + self._incomplete_steps, + ) + return pairs + + def _depth_decode( + self, + cond_hidden: torch.Tensor, + uncond_hidden: torch.Tensor, + c0: torch.Tensor, + seeds: torch.Tensor, + positions: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Walk ``c1..c7``, replaying the loop from a CUDA graph where possible. + + Seven four-layer forwards over at most eight positions is far too + little work per kernel to cover its own launch, so the loop is recorded + once per pair count and replayed after that. The cache falls back to + the eager loop for any shape it has not recorded or could not record. + """ + return self._depth_graphs.run(cond_hidden, uncond_hidden, c0, seeds, positions) + + def _depth_decode_eager( + self, + cond_hidden: torch.Tensor, + uncond_hidden: torch.Tensor, + c0: torch.Tensor, + seeds: torch.Tensor, + positions: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Walk ``c1..c7``, returning the codes and the stacked depth hiddens. + + Both branches run so the depth heads can be guided the same way the + ``c0`` head is; the drawn code is shared, so the two branches stay + aligned. + """ + decoder = self.rvq_decoder + pair_hidden = torch.cat((cond_hidden, uncond_hidden), dim=0) + c0_paired = c0.repeat(2) + c0_embed = self.model.embed_tokens(c0_paired + AUDIO_CODE_OFFSET) + seq = [ + decoder.projection(pair_hidden).unsqueeze(1), + decoder.projection(c0_embed).unsqueeze(1), + ] + n_pairs = int(c0.shape[0]) + codes = [c0] + depth_parts: list[torch.Tensor] = [] + # Every residual draw's noise comes from the same hash of + # ``(seed, position + codebook, candidate)``, so all seven are produced + # in one batch here instead of re-running the twenty-odd integer ops + # per codebook inside the loop. The arithmetic is elementwise, so the + # values are the same ones the per-codebook calls produced. + depth_gumbel = gumbel_noise( + seeds, + positions.unsqueeze(1) + self.depth_draw_offsets, + decoder.audio_vocab_size, + ) + for index in range(1, NUM_CODEBOOKS): + out = decoder(torch.cat(seq, dim=1))[:, -1] + depth_parts.append(out[:n_pairs].detach()) + head_logits = decoder.audio_heads[index - 1](out).float() + guided = apply_depth_cfg(head_logits[:n_pairs], head_logits[n_pairs:]) + code = draw_from_gumbel(guided, depth_gumbel[:, index - 1]) + codes.append(code) + if index < NUM_CODEBOOKS - 1: + embedding = self.audio_embeddings(code.repeat(2) + (index - 1) * AUDIO_VOCAB_SIZE) + seq.append(decoder.projection(embedding).unsqueeze(1)) + return torch.stack(codes, dim=1), torch.cat(depth_parts, dim=-1) + + def _state_for(self, pair_id: str, row: int) -> RequestARState: + """Fetch or create the state for a guided pair. + + Created on the step the pair first produces a logit, which is the step + its prefill completes, so the recorded prompt length is the real one. + """ + state = self._states.get(pair_id) + if state is None: + seed = self._row_seeds[row] if row < len(self._row_seeds) else 0 + state = RequestARState( + seed=seed, + max_frames=self._row_max_frames[row] if row < len(self._row_max_frames) else DEFAULT_MAX_AUDIO_FRAMES, + prompt_tokens=self._row_prompt_tokens[row] if row < len(self._row_prompt_tokens) else 0, + ) + self._states[pair_id] = state + logger.info( + "MiniMax Music 3 AR start pair=%s seed=%d max_frames=%d prompt_tokens=%d", + pair_id, + state.seed, + state.max_frames, + state.prompt_tokens, + ) + return state + + # ------------------------------------------------------------ emit / free + + def postprocess( + self, + hidden_states_slice: torch.Tensor, + multimodal_outputs: Any = None, + **req_infos: Any, + ) -> dict[str, Any]: + """Publish this request's next span of conditioning frames. + + Mid-stream this is one window, released a frame late so that a window + handed downstream is provably not the last. On the final step it is + instead a single contiguous span covering every window not yet + emitted, because one step produces exactly one payload and the tail + cannot be split across two. + """ + del hidden_states_slice, multimodal_outputs + request_id = _request_key(req_infos.get("global_request_id")) + # The pair's two rows share one song, so only the conditioned row + # publishes it. The companion would otherwise emit a duplicate span + # and the acoustic stage would decode the song twice. + if _role_from_request(request_id) == _UNCOND: + return {} + state = self._states.get(_pair_id_from_request(request_id)) + if state is None: + return {} + if state.finished: + return self._terminal_payload(state) + window = state.ready_window() + if window is None: + return {} + chunk = state.take_window(window, is_final=False) + return self._span_payload(state, chunk, start_frame=window.start, is_final=False) + + def _terminal_payload(self, state: RequestARState) -> dict[str, Any]: + """Emit every remaining frame as one span, once.""" + start = state.emitted_windows * AR_CHUNK_HOP_FRAMES + end = state.frames.end_frame + if end <= start: + return {} + chunk = state.frames.slice(start, end) + state.emitted_windows = len(chunk_windows(state.generated_frames)) + state.frames.clear() + return self._span_payload(state, chunk, start_frame=start, is_final=True) + + def _span_payload( + self, + state: RequestARState, + chunk: torch.Tensor, + *, + start_frame: int, + is_final: bool, + ) -> dict[str, Any]: + """Build the cross-stage payload for one contiguous frame span. + + ``num_processed_tokens`` is the span's first global AR frame index, so + the acoustic stage can re-derive window boundaries and therefore which + window is first and which is last. + """ + frames = int(chunk.shape[0]) + codes = state.take_codes(start_frame, start_frame + frames) + payload: dict[str, Any] = { + "latent": chunk.contiguous(), + "meta": { + "num_processed_tokens": start_frame, + "codec_chunk_frames": frames, + "last_chunk": is_final, + "audio_seed": state.seed, + }, + } + if codes is not None: + payload["codes"] = {"audio": codes} + return payload + + def make_omni_output(self, model_outputs: torch.Tensor | OmniOutput, **kwargs: Any) -> OmniOutput: + if isinstance(model_outputs, OmniOutput): + return model_outputs + infos = kwargs.get("model_intermediate_buffer") + if infos is None: + infos = kwargs.get("runtime_additional_information") + if not infos: + return OmniOutput(text_hidden_states=model_outputs, multimodal_outputs=None) + + latents: list[torch.Tensor] = [] + codes: list[torch.Tensor] = [] + metas: list[dict[str, Any]] = [] + any_payload = False + for info in infos: + info = info if isinstance(info, dict) else {} + latent = info.get("latent") + if isinstance(latent, torch.Tensor) and latent.numel(): + latents.append(latent) + sub = info.get("codes") + codes.append(sub.get("audio") if isinstance(sub, dict) else torch.empty(0)) + meta = info.get("meta") + metas.append(dict(meta) if isinstance(meta, dict) else {}) + any_payload = True + # The runner merges postprocess output into this buffer and + # never clears it, so a window left behind would be re-emitted + # on every later step and the acoustic stage would decode the + # same conditioning again and again. Consume it here. + info.pop("latent", None) + info.pop("codes", None) + info.pop("meta", None) + else: + latents.append(torch.empty(0)) + codes.append(torch.empty(0)) + metas.append({}) + if not any_payload: + return OmniOutput(text_hidden_states=model_outputs, multimodal_outputs=None) + payload: dict[str, Any] = {"latent": latents} + if any(isinstance(entry, torch.Tensor) and entry.numel() for entry in codes): + payload["codes"] = {"audio": codes} + # The span's start frame and terminal flag ride along with it: without + # them the acoustic stage cannot tell which window it is looking at, + # and therefore cannot crop it correctly. + # + # Every field is a list with one entry per batch row, in the same order + # as ``latent``. The payload is partitioned per request by indexing + # list values, so a scalar here would be broadcast to the whole batch + # and every request would decode with the last request's seed, window + # offset and terminal flag. + meta_keys = {key for meta in metas for key in meta} + if meta_keys: + payload["meta"] = {key: [meta.get(key) for meta in metas] for key in sorted(meta_keys)} + return OmniOutput(text_hidden_states=model_outputs, multimodal_outputs=payload) + + def on_requests_finished(self, finished_req_ids: Iterable[str]) -> None: + """Release per-request state once the engine has retired the request. + + The key must be normalized the same way it was when the state was + created, or every finished song leaks its frame buffer. + """ + for engine_id in finished_req_ids: + # These are engine-side ids, so translate through the mapping the + # forward pass built. A pair's state is released when its second + # member retires. + self._row_constants.pop(str(engine_id), None) + pair_id = self._pair_of_engine_id.pop(str(engine_id), None) + if pair_id is None or any(pair_id == value for value in self._pair_of_engine_id.values()): + continue + request_id = pair_id + state = self._states.pop(pair_id, None) + if state is not None: + logger.info( + "MiniMax Music 3 AR done request=%s frames=%d reason=%s peak_buffer=%d", + request_id, + state.generated_frames, + state.finish_reason, + state.frames.peak_frames, + ) + state.release() + + # --------------------------------------------------------------- weights + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + """Load the Qwen3 backbone, then the audio modules from their component. + + The backbone comes from the stage's own checkpoint folder and loads + through vLLM. The audio embedding table and the depth decoder live in + a separate component of the repo and are loaded explicitly. + """ + loader = AutoWeightsLoader(self, skip_prefixes=["audio_embeddings.", "rvq_decoder."]) + loaded = loader.load_weights(weights) + loaded |= load_depth_decoder_weights( + self.vllm_config, + audio_embeddings=self.audio_embeddings, + rvq_decoder=self.rvq_decoder, + ) + return loaded diff --git a/vllm_omni/model_executor/models/minimax_music3/weights.py b/vllm_omni/model_executor/models/minimax_music3/weights.py new file mode 100644 index 000000000..2ab643c65 --- /dev/null +++ b/vllm_omni/model_executor/models/minimax_music3/weights.py @@ -0,0 +1,133 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Component weight loading for MiniMax Music 3. + +The checkpoint is a multi-component repo. The Qwen3 backbone is a normal +``Qwen3ForCausalLM`` under ``language_model/`` and loads through vLLM. The +audio embedding table and the depth decoder live in ``rvq_depth_decoder/``, +and the acoustic stage's three components live in their own folders. None of +those are things vLLM's loader knows how to map, so they are read directly. + +Every loader here is strict. A silently partial load would produce audio that +sounds plausible but is wrong, which is far worse than a startup failure. +""" + +from __future__ import annotations + +import os +from pathlib import Path + +import torch +from torch import nn +from vllm.config import VllmConfig +from vllm.logger import init_logger + +from .constants import AUDIO_VOCAB_SIZE, NUM_CODEBOOKS + +logger = init_logger(__name__) + +_AUDIO_EMBEDDING_KEY = "audio_embeddings.weight" +_COMPONENT_WEIGHT_NAME = "diffusion_pytorch_model.safetensors" + +# Component folders, relative to the repository root. +DEPTH_DECODER_DIR = "rvq_depth_decoder" +CONDITION_ENCODER_DIR = "condition_encoder" +TRANSFORMER_DIR = "transformer" +VOCODER_DIR = "vocoder" + +_ROOT_MARKERS = (DEPTH_DECODER_DIR, TRANSFORMER_DIR, VOCODER_DIR) + + +def resolve_repo_root(model_path: str | os.PathLike[str]) -> Path: + """Find the multi-component repository root from a stage's model path. + + A stage that declares ``model_subdir`` sees the subfolder as its model + path, so the root is usually the parent. Both are checked, and the search + walks up a couple of levels for snapshot layouts. + + Raises: + FileNotFoundError: If no ancestor looks like the repository root. + """ + path = Path(model_path).expanduser().resolve() + candidates = [path, *path.parents[:3]] + for candidate in candidates: + if all((candidate / marker).is_dir() for marker in _ROOT_MARKERS): + return candidate + raise FileNotFoundError( + f"MiniMax Music 3 could not locate the repository root from {path}; expected sibling folders {_ROOT_MARKERS}" + ) + + +def load_component_state(root: Path, component: str, *, device: str = "cpu") -> dict[str, torch.Tensor]: + """Read one component's safetensors, sharded or not. + + Raises: + FileNotFoundError: If the component has no weight files. + """ + from safetensors.torch import load_file + + folder = root / component + single = folder / _COMPONENT_WEIGHT_NAME + if single.exists(): + return load_file(str(single), device=device) + + shards = sorted(folder.glob(f"{Path(_COMPONENT_WEIGHT_NAME).stem}-*.safetensors")) + if not shards: + raise FileNotFoundError(f"MiniMax Music 3 component {component!r} has no weights under {folder}") + state: dict[str, torch.Tensor] = {} + for shard in shards: + state.update(load_file(str(shard), device=device)) + return state + + +def load_depth_decoder_weights( + vllm_config: VllmConfig, + *, + audio_embeddings: nn.Embedding, + rvq_decoder: nn.Module, +) -> set[str]: + """Load the audio embedding table and the RVQ depth decoder in place. + + Args: + vllm_config: The stage's config, used to locate the checkpoint. + audio_embeddings: The ``c1..c7`` embedding table to populate. + rvq_decoder: The depth decoder to populate. + + Returns: + The model-side parameter names that were loaded. + + Raises: + ValueError: If the audio embedding is absent or the wrong shape. + """ + root = resolve_repo_root(vllm_config.model_config.model) + state = load_component_state(root, DEPTH_DECODER_DIR) + + weight = state.pop(_AUDIO_EMBEDDING_KEY, None) + expected = (AUDIO_VOCAB_SIZE * (NUM_CODEBOOKS - 1), audio_embeddings.embedding_dim) + if weight is None or tuple(weight.shape) != expected: + raise ValueError( + f"MiniMax Music 3 audio embedding mismatch: expected {expected}, got " + f"{None if weight is None else tuple(weight.shape)}" + ) + target = audio_embeddings.weight + target.data.copy_(weight.to(device=target.device, dtype=target.dtype)) + + rvq_decoder.load_checkpoint_state(state) + logger.info( + "MiniMax Music 3 loaded audio embedding and depth decoder from %s", + root / DEPTH_DECODER_DIR, + ) + loaded = {"audio_embeddings.weight"} + loaded.update(f"rvq_decoder.{name}" for name, _ in rvq_decoder.named_parameters()) + return loaded + + +__all__ = [ + "CONDITION_ENCODER_DIR", + "DEPTH_DECODER_DIR", + "TRANSFORMER_DIR", + "VOCODER_DIR", + "load_component_state", + "load_depth_decoder_weights", + "resolve_repo_root", +] diff --git a/vllm_omni/model_executor/models/registry.py b/vllm_omni/model_executor/models/registry.py index 3cf888d46..0064f1c4a 100644 --- a/vllm_omni/model_executor/models/registry.py +++ b/vllm_omni/model_executor/models/registry.py @@ -457,6 +457,17 @@ _OMNI_MODELS = { "qwen3_vl", "AuraQwen3VLForConditionalGeneration", ), + ## MiniMax-Music3 (text-to-music; AR talker -> flow-matching acoustic decoder) + "MiniMaxMusic3TalkerForConditionalGeneration": ( + "minimax_music3", + "talker", + "MiniMaxMusic3TalkerForConditionalGeneration", + ), + "MiniMaxMusic3AcousticForConditionalGeneration": ( + "minimax_music3", + "acoustic", + "MiniMaxMusic3AcousticForConditionalGeneration", + ), } diff --git a/vllm_omni/model_executor/stage_input_processors/minimax_music3.py b/vllm_omni/model_executor/stage_input_processors/minimax_music3.py new file mode 100644 index 000000000..44037e4fc --- /dev/null +++ b/vllm_omni/model_executor/stage_input_processors/minimax_music3.py @@ -0,0 +1,544 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Stage-input processors for MiniMax Music 3: AR talker -> acoustic decoder. + +Three entry points, all referenced from +``vllm_omni/model_executor/models/minimax_music3/pipeline.py``: + +* ``expand_cfg_prompts`` (stage 0 ``prompt_expand_func``) builds the + unconditioned twin of a conditioned request. Classifier-free guidance is + mandatory for this checkpoint, so in practice every request expands. +* ``ar2acoustic_async_chunk`` (stage 0 ``async_chunk_process_next_stage_input_func``) + forwards one AR frame span per call to stage 1, and is also the flush point + for the request's tail. +* ``ar2acoustic`` (stage 1 ``sync_process_input_func``) builds the stage-1 + engine input for the non-streaming path, where the whole song arrives at + once after stage 0 finishes. + +Stage-0 emit contract +--------------------- +Per emission the talker publishes a multimodal output shaped like:: + + { + "latent": FloatTensor[1, T, 32768] or FloatTensor[T, 32768], + "codes": {"audio": LongTensor[T, 8]}, + "meta": { + "num_processed_tokens": int, # global AR frame index of frame 0 + "last_chunk": bool, # True only on the terminal emission + "audio_seed": int, # optional + }, + } + +The span is normally one window: ``T == AR_CHUNK_FRAMES``, windows advance by +``AR_CHUNK_HOP_FRAMES``, and window ``k`` starts at AR frame +``k * AR_CHUNK_HOP_FRAMES``. The terminal emission is the exception. A window +cannot be proven non-final until a further frame exists, so the talker holds +the most recently completed window back and, on the step where it knows the +request is ending, publishes the whole remaining frame span at once. That span +covers up to two windows (the held-back complete one plus the truncated final +one), so ``T`` can reach ``AR_CHUNK_FRAMES + AR_CHUNK_HOP_FRAMES``. Stage 1 +re-derives the window boundaries inside a span from +``meta.num_processed_tokens`` and the span length through +``models/minimax_music3/chunking.chunk_windows``. + +``meta.num_processed_tokens`` is required rather than derived: once a terminal +span carries more than one window, a running ``chunk_seq * hop`` count no +longer identifies the span's first frame. The derived value is kept only as a +fallback for pre-terminal emissions. + +``hidden_states.output`` is accepted as an alias for ``latent`` because both +are legal homes for a float payload; ``latent`` is canonical and matches the +stage's ``engine_output_type``. + +Stage-0 -> stage-1 payload +-------------------------- +Every key below is a declared field of ``OmniPayloadStruct`` / +``MetaStruct``, which are ``forbid_unknown_fields=True`` msgspec structs, so +an undeclared key is a construction error rather than silent data loss:: + + latent FloatTensor[T, 32768] span conditioning, bf16 + codes.audio LongTensor[T, 8] RVQ codes for the span + meta.chunk_seq int emission index; is_first == (chunk_seq == 0) + meta.num_processed_tokens int absolute AR frame index of the span start + meta.codec_chunk_frames int T; span end (exclusive) = start + T + meta.last_chunk bool terminal span + meta.stream_finished BoolTensor terminal marker that survives transport + meta.audio_seed int seed for the acoustic stage's noise draw + meta.req_id list[str] single-element external request id + meta.finished BoolTensor overwritten by the chunk adapter + +``codes.audio`` is rank 2 on purpose: the chunk transfer adapter reroutes rank-1 +code tensors onto the token path (``.tolist()``), while rank-2 tensors ride the +payload channel and leave a single ``[0]`` placeholder prompt token per chunk. +``meta.finished`` is stripped before the payload reaches the stage-1 model, so +``meta.stream_finished`` and ``meta.last_chunk`` are what stage 1 reads. + +The overlap crop is deliberately NOT in the payload. Stage 1 recovers it from +the span's global frame range and ``last_chunk`` through +``models/minimax_music3/chunking.crop_sample_bounds``, which is the single +owner of that arithmetic; duplicating it here would let the two copies drift. +""" + +from __future__ import annotations + +from typing import Any + +import torch +from vllm.logger import init_logger + +from vllm_omni.data_entry_keys import CodesStruct, MetaStruct, OmniPayloadStruct +from vllm_omni.inputs.data import OmniTokensPrompt +from vllm_omni.model_executor.models.minimax_music3.constants import ( + AR_CFG_SCALE, + AR_CHUNK_HOP_FRAMES, + AR_HIDDEN_SIZE, + NUM_CODEBOOKS, +) +from vllm_omni.model_executor.models.minimax_music3.prompt import build_cfg_null_token_ids +from vllm_omni.model_executor.stage_input_processors.bagel import ExpandedPrompt + +__all__ = [ + "CFG_UNCOND_SUFFIX", + "ar2acoustic", + "ar2acoustic_async_chunk", + "expand_cfg_prompts", +] + +logger = init_logger(__name__) + +CFG_UNCOND_SUFFIX = "__cfg_uncond" + +_COND = "cond" +_UNCOND = "uncond" + +# Per-request forwarder state, parked on the transfer manager's request map. +_STATE_KEY = "_minimax_music3_async_state" + +# The window conditioning is emitted in bfloat16. Each chunk is its own +# /dev/shm segment and consecutive windows overlap by half, so a float32 +# 200-frame window would move 26 MB of mostly-redundant bytes per hop against +# 13 MB in bfloat16. The backbone runs in bfloat16 anyway, so this discards no +# information the AR stage actually produced. +_LATENT_TRANSPORT_DTYPE = torch.bfloat16 + + +# --------------------------------------------------------------------------- +# Stage 0: classifier-free-guidance prompt expansion +# --------------------------------------------------------------------------- + + +def _prompt_token_ids(prompt: dict[str, Any] | str | Any) -> list[int] | None: + """Return the prompt's token ids, or ``None`` when only text is available.""" + token_ids: Any = None + if isinstance(prompt, dict): + token_ids = prompt.get("prompt_token_ids") + elif isinstance(prompt, (list, tuple)): + token_ids = prompt + if token_ids is None: + return None + if isinstance(token_ids, torch.Tensor): + token_ids = token_ids.reshape(-1).tolist() + if not isinstance(token_ids, (list, tuple)) or not token_ids: + return None + return [int(t) for t in token_ids] + + +def _uncond_extra_args(extra_args: dict[str, Any], cfg_scale: float) -> dict[str, Any]: + """Clone the conditioned row's extra_args for the unconditioned twin. + + Every key is carried through, not only the ``cfg_*`` ones: the stage runs + with ``has_sampling_extra_args`` on, so extra_args reach the model's + ``forward`` per row, and dropping a key here would make the two rows + execute under different settings and decode out of lockstep. + ``cfg_null_prompt`` is the exception: the companion IS the null prompt. + """ + companion = dict(extra_args) + companion.pop("cfg_null_prompt", None) + companion["cfg_role"] = _UNCOND + companion["cfg_scale"] = cfg_scale + return companion + + +def expand_cfg_prompts( + prompt: dict[str, Any] | str, + sampling_params: Any, +) -> list[ExpandedPrompt]: + """Emit the unconditioned companion for a guided MiniMax Music 3 request. + + Classifier-free guidance is not optional for this checkpoint, and the + deploy YAMLs set ``cfg_role: cond`` / ``cfg_scale: 1.5`` as stage-0 + defaults, so an ordinary request expands into exactly one companion and + the pair decodes in lockstep. + + The companion's token ids come from ``build_cfg_null_token_ids``, which + overwrites the caption-and-lyrics span with ``<|audio_cfg|>`` while keeping + the framing tokens. Token ids are preferred over text because the two rows + are only blendable when they have identical length and an identical + ``<|audio_start|>`` position, and only the token-id form guarantees that + without a tokenizer round trip. + + Args: + prompt: The raw user prompt. A dict carrying ``prompt_token_ids`` is + the preferred form; a text-only prompt must supply a pre-built, + length-matched ``cfg_null_prompt`` in ``extra_args``. + sampling_params: Stage-0 sampling params for the conditioned row. + + Returns: + A one-element list with the unconditioned companion, or an empty list + when the request is not CFG-guided. + + Raises: + ValueError: If the request is guided but neither prompt token ids nor + a pre-built ``cfg_null_prompt`` is available. + """ + extra_args = getattr(sampling_params, "extra_args", None) or {} + if extra_args.get("cfg_role") != _COND: + # Not guided, or this IS the companion. Either way, nothing to expand. + return [] + try: + cfg_scale = float(extra_args.get("cfg_scale", AR_CFG_SCALE)) + except (TypeError, ValueError): + logger.warning( + "MiniMax Music 3: ignoring non-numeric cfg_scale %r; request decodes unguided", + extra_args.get("cfg_scale"), + ) + return [] + if cfg_scale <= 1.0: + return [] + + token_ids = _prompt_token_ids(prompt) + companion_prompt: dict[str, Any] | str + if token_ids is not None: + null_ids = build_cfg_null_token_ids(token_ids) + if isinstance(prompt, dict): + companion_prompt = {**prompt, "prompt_token_ids": null_ids} + # Drop the text form so the input processor does not re-tokenize + # and hand back a row of a different length. + companion_prompt.pop("prompt", None) + companion_prompt.pop("prompt_embeds", None) + else: + companion_prompt = {"prompt_token_ids": null_ids} + else: + null_prompt = extra_args.get("cfg_null_prompt") + if not isinstance(null_prompt, str) or not null_prompt: + raise ValueError( + "MiniMax Music 3 CFG request carries neither prompt_token_ids nor a " + "pre-built cfg_null_prompt in extra_args. Submit the prompt as token " + "ids (preferred; the null twin is then built by " + "models.minimax_music3.prompt.build_cfg_null_token_ids) or attach a " + "length-matched cfg_null_prompt string." + ) + if isinstance(prompt, dict): + companion_prompt = {**prompt, "prompt": null_prompt} + companion_prompt.pop("prompt_token_ids", None) + companion_prompt.pop("prompt_embeds", None) + else: + companion_prompt = null_prompt + + return [ + ExpandedPrompt( + prompt=companion_prompt, + role=_UNCOND, + request_id_suffix=CFG_UNCOND_SUFFIX, + sampling_params_override={"extra_args": _uncond_extra_args(extra_args, cfg_scale)}, + ) + ] + + +# --------------------------------------------------------------------------- +# Stage 0 -> stage 1: async-chunk window forwarder +# --------------------------------------------------------------------------- + + +def _payload_section(payload: Any, key: str) -> dict[str, Any]: + """Return a nested payload section as a dict (empty when absent).""" + if not isinstance(payload, dict): + return {} + section = payload.get(key) + return section if isinstance(section, dict) else {} + + +def _window_latent(payload: Any) -> torch.Tensor | None: + """Return this window's conditioning as ``[T, 32768]``, or ``None``. + + ``latent`` is the canonical key; ``hidden_states.output`` is accepted for + talkers that publish through the hidden-state channel instead. A leading + singleton batch dimension is squeezed so stage 1 sees a stable rank. + """ + if not isinstance(payload, dict): + return None + tensor = payload.get("latent") + if not isinstance(tensor, torch.Tensor): + tensor = _payload_section(payload, "hidden_states").get("output") + if not isinstance(tensor, torch.Tensor) or tensor.numel() == 0: + return None + if tensor.ndim == 3: + if tensor.shape[0] != 1: + raise ValueError( + f"MiniMax Music 3 window latent must have a singleton batch dim; got {tuple(tensor.shape)}" + ) + tensor = tensor[0] + if tensor.ndim != 2: + raise ValueError(f"MiniMax Music 3 window latent must be rank 2 or 3; got {tuple(tensor.shape)}") + if tensor.shape[1] != AR_HIDDEN_SIZE: + raise ValueError(f"MiniMax Music 3 window latent has {tensor.shape[1]} features; expected {AR_HIDDEN_SIZE}") + return tensor + + +def _window_codes(payload: Any) -> torch.Tensor | None: + """Return this window's RVQ codes as ``[T, 8]``, or ``None``.""" + codes = _payload_section(payload, "codes").get("audio") + if not isinstance(codes, torch.Tensor) or codes.numel() == 0: + return None + if codes.ndim == 1: + if codes.numel() % NUM_CODEBOOKS != 0: + raise ValueError( + f"MiniMax Music 3 flat code payload of length {codes.numel()} is not a " + f"multiple of {NUM_CODEBOOKS} codebooks" + ) + codes = codes.reshape(-1, NUM_CODEBOOKS) + if codes.ndim != 2 or codes.shape[1] != NUM_CODEBOOKS: + raise ValueError(f"MiniMax Music 3 codes must be [T, {NUM_CODEBOOKS}]; got {tuple(codes.shape)}") + return codes.to(torch.long) + + +def _request_audio_seed(request: Any, payload: Any) -> int | None: + """Resolve the seed the acoustic stage draws its noise from. + + A seed published by the talker wins, because a talker that already + consumed the request seed knows which value it used. Otherwise fall back + to the stage-0 request's own sampling seed. + """ + published = _payload_section(payload, "meta").get("audio_seed") + if isinstance(published, int): + return published + seed = getattr(getattr(request, "sampling_params", None), "seed", None) + return int(seed) if isinstance(seed, int) else None + + +def _forwarder_state(transfer_manager: Any, request_id: str) -> dict[str, Any]: + """Return (creating if needed) the per-request forwarder state.""" + request_state = transfer_manager.request_payload.get(request_id) + if not isinstance(request_state, dict) or _STATE_KEY not in request_state: + request_state = {_STATE_KEY: {"chunk_seq": 0, "terminal_sent": False}} + transfer_manager.request_payload[request_id] = request_state + return request_state[_STATE_KEY] + + +def ar2acoustic_async_chunk( + transfer_manager: Any, + multimodal_output: dict[str, Any] | None, + request: Any, + is_finished: bool = False, +) -> OmniPayloadStruct | None: + """Forward one AR frame span to the acoustic stage. + + The talker publishes whole spans, so this is a relay plus bookkeeping + rather than an accumulator: it stamps the emission index, the span's global + start frame, its length and the terminal flags, then hands the tensors to + the connector. + + This hook is the flush point for the request's tail. The AR scheduler calls + ``save_async`` on the step that stops the request, including when that + step produced no inter-stage output, so this function is guaranteed one + call with ``is_finished=True`` (omni_ar_scheduler.py:551, base + chunk_transfer_adapter.py:341). Whatever the talker published on that step + therefore reaches stage 1 in the same hop. + + Returns ``None`` on steps where the talker released no span and the request + is still running. + """ + request_id = request.external_req_id + finished = bool(is_finished or request.is_finished()) + + state = _forwarder_state(transfer_manager, request_id) + + latent = _window_latent(multimodal_output) + codes = _window_codes(multimodal_output) + published_meta = _payload_section(multimodal_output, "meta") + published_last = published_meta.get("last_chunk") is True + + if state["terminal_sent"]: + if latent is not None: + logger.error( + "MiniMax Music 3 talker emitted %d more frames for request %s after the " + "terminal span; the downstream stream is already closed and they are dropped", + int(latent.shape[0]), + request_id, + ) + return None + + if latent is None: + if not finished: + return None + # Terminal step with nothing left to send: still deliver the marker so + # stage 1 can flush and complete instead of polling until timeout. + state["terminal_sent"] = True + # The unconditioned twin never publishes a span: the pair shares one + # song and the conditioned row emits it. Its silence is the contract, + # not a failure. + if state["chunk_seq"] == 0 and not request_id.endswith(CFG_UNCOND_SUFFIX): + logger.error("MiniMax Music 3 talker produced no audio for request %s", request_id) + return OmniPayloadStruct( + meta=MetaStruct( + chunk_seq=state["chunk_seq"], + last_chunk=True, + stream_finished=torch.tensor(True, dtype=torch.bool), + finished=torch.tensor(True, dtype=torch.bool), + req_id=[request_id], + ) + ) + + span_frames = int(latent.shape[0]) + if codes is not None and codes.shape[0] != span_frames: + raise ValueError( + f"MiniMax Music 3 span mismatch for request {request_id}: latent has " + f"{span_frames} frames but codes have {codes.shape[0]}" + ) + + chunk_seq = int(state["chunk_seq"]) + state["chunk_seq"] = chunk_seq + 1 + + # A span is terminal when the engine stopped the request or the talker said + # so. Only the engine's verdict closes the stream, so a talker that flags + # a span last while generation continues is reported rather than obeyed. + terminal = bool(finished or published_last) + if published_last and not finished: + logger.warning( + "MiniMax Music 3 talker flagged a terminal span for request %s while the " + "engine still has the request running", + request_id, + ) + state["terminal_sent"] = terminal + + span_start = published_meta.get("num_processed_tokens") + if not isinstance(span_start, int): + # Fallback for talkers that publish one hop-aligned window per step. + # It is only sound before the terminal span, which may cover more than + # one window and therefore break the running count. + span_start = chunk_seq * AR_CHUNK_HOP_FRAMES + if terminal and span_frames > AR_CHUNK_HOP_FRAMES: + logger.warning( + "MiniMax Music 3 terminal span for request %s carries %d frames but no " + "meta.num_processed_tokens; the derived start frame %d assumes a single " + "hop-aligned window and may misplace the tail", + request_id, + span_frames, + span_start, + ) + + finished_flag = torch.tensor(terminal, dtype=torch.bool) + meta = MetaStruct( + chunk_seq=chunk_seq, + num_processed_tokens=int(span_start), + codec_chunk_frames=span_frames, + last_chunk=terminal, + # meta.finished is stripped before the payload reaches the stage-1 + # model; stream_finished is what survives and drives the final flush. + stream_finished=finished_flag, + finished=finished_flag, + req_id=[request_id], + ) + audio_seed = _request_audio_seed(request, multimodal_output) + if audio_seed is not None: + meta.audio_seed = audio_seed + + payload = OmniPayloadStruct( + latent=latent.detach().to(_LATENT_TRANSPORT_DTYPE).contiguous(), + meta=meta, + ) + if codes is not None: + payload.codes = CodesStruct(audio=codes.detach().contiguous()) + return payload + + +# --------------------------------------------------------------------------- +# Stage 1: non-streaming engine-input builder +# --------------------------------------------------------------------------- + + +def _empty_acoustic_prompt() -> OmniTokensPrompt: + """Zero-length stage-1 request. + + Stage 1 decodes it to empty audio, which the serving layer's no-audio + guard turns into a request-level error. That is preferable to raising + here, where the failure would surface as a stalled request. + """ + return OmniTokensPrompt( + prompt_token_ids=[], + multi_modal_data=None, + mm_processor_kwargs=None, + additional_information=None, + ) + + +def ar2acoustic( + source_outputs: list[Any], + prompt: Any = None, + _requires_multimodal_data: bool = False, +) -> list[OmniTokensPrompt]: + """Build stage-1 engine inputs for the non-streaming path. + + With ``async_chunk: false`` the whole song arrives once stage 0 has + finished, so the song is handed over as a single logical window and stage + 1 does its own windowing through ``chunking.chunk_windows``. The payload + layout matches the streaming one key for key, so the acoustic stage reads + one contract in both modes. + + The prompt is ``[0] * num_windows``: the generation scheduler sizes the + request from ``prompt_token_ids``, and one placeholder per acoustic window + keeps that budget proportional to the work without shipping the + conditioning through the token path. + """ + del prompt, _requires_multimodal_data + + acoustic_inputs: list[OmniTokensPrompt] = [] + for ar_output in source_outputs: + if not ar_output.finished: + continue + output = ar_output.outputs[0] + multimodal_output = getattr(output, "multimodal_output", None) + + try: + latent = _window_latent(multimodal_output) + codes = _window_codes(multimodal_output) + except ValueError as exc: + logger.warning("Skipping malformed MiniMax Music 3 stage-0 output: %s", exc) + acoustic_inputs.append(_empty_acoustic_prompt()) + continue + + if latent is None: + logger.error("MiniMax Music 3 talker produced no audio for request %s", ar_output.request_id) + acoustic_inputs.append(_empty_acoustic_prompt()) + continue + + total_frames = int(latent.shape[0]) + meta: dict[str, Any] = { + "chunk_seq": 0, + "num_processed_tokens": 0, + "codec_chunk_frames": total_frames, + "last_chunk": True, + "stream_finished": torch.tensor(True, dtype=torch.bool), + "req_id": [ar_output.request_id], + } + published_seed = _payload_section(multimodal_output, "meta").get("audio_seed") + if isinstance(published_seed, int): + meta["audio_seed"] = published_seed + + additional_information: dict[str, Any] = { + "latent": latent.detach().to(_LATENT_TRANSPORT_DTYPE).contiguous(), + "meta": meta, + } + if codes is not None: + additional_information["codes"] = {"audio": codes.detach().contiguous()} + + num_windows = max(1, -(-total_frames // AR_CHUNK_HOP_FRAMES)) + acoustic_inputs.append( + OmniTokensPrompt( + prompt_token_ids=[0] * num_windows, + multi_modal_data=None, + mm_processor_kwargs=None, + additional_information=additional_information, + ) + ) + return acoustic_inputs diff --git a/vllm_omni/transformers_utils/configs/__init__.py b/vllm_omni/transformers_utils/configs/__init__.py index cdf105d69..2a31b461c 100644 --- a/vllm_omni/transformers_utils/configs/__init__.py +++ b/vllm_omni/transformers_utils/configs/__init__.py @@ -23,6 +23,7 @@ _CLASS_TO_MODULE: dict[str, str] = { "DotsTTSConfig": "vllm_omni.transformers_utils.configs.dots_tts", "VoxtralTTSConfig": "vllm_omni.transformers_utils.configs.voxtral_tts", "CosyVoice3Config": "vllm_omni.transformers_utils.configs.cosyvoice3", + "MiniMaxMusic3Config": "vllm_omni.transformers_utils.configs.minimax_music3", "OmniVoiceConfig": "vllm_omni.transformers_utils.configs.omnivoice", "BailingMoeV2Config": "vllm_omni.transformers_utils.configs.ming_flash_omni", "BailingMM2Config": "vllm_omni.transformers_utils.configs.ming_flash_omni", @@ -46,6 +47,7 @@ __all__ = [ "DotsTTSConfig", "VoxtralTTSConfig", "CosyVoice3Config", + "MiniMaxMusic3Config", "OmniVoiceConfig", "BailingMoeV2Config", "BailingMM2Config", @@ -78,6 +80,7 @@ from vllm_omni.transformers_utils.configs import glm_tts as _glm_tts # noqa: F4 from vllm_omni.transformers_utils.configs import higgs_audio_v3 as _higgs_audio_v3 # noqa: F401, E402 from vllm_omni.transformers_utils.configs import mammoth_moda2 as _mammoth_moda2 # noqa: F401, E402 from vllm_omni.transformers_utils.configs import ming_flash_omni as _ming_flash_omni # noqa: F401, E402 +from vllm_omni.transformers_utils.configs import minimax_music3 as _minimax_music3 # noqa: F401, E402 from vllm_omni.transformers_utils.configs import omnivoice as _omnivoice # noqa: F401, E402 from vllm_omni.transformers_utils.configs import voxcpm2 as _voxcpm2 # noqa: F401, E402 from vllm_omni.transformers_utils.configs import voxtral_tts as _voxtral_tts # noqa: F401, E402 diff --git a/vllm_omni/transformers_utils/configs/minimax_music3.py b/vllm_omni/transformers_utils/configs/minimax_music3.py new file mode 100644 index 000000000..be951055d --- /dev/null +++ b/vllm_omni/transformers_utils/configs/minimax_music3.py @@ -0,0 +1,65 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""HuggingFace config for the MiniMax Music 3 composite repository. + +The checkpoint root carries only:: + + {"architectures": ["MiniMaxMusic3ForConditionalGeneration"], + "model_type": "minimax_music3"} + +and ships no modelling code, so ``trust_remote_code`` has nothing to load and +``AutoConfig`` rejects the repository outright. Stage 0 sidesteps this by +declaring ``model_subdir="language_model"`` and reading the plain +``Qwen3ForCausalLM`` config, but the acoustic stage must open the root to reach +``condition_encoder/``, ``transformer/`` and ``vocoder/``. Registering this +class is what lets that stage build a ``ModelConfig`` at all. + +The generic transformer fields describe the AR backbone, because the composite +root has no transformer of its own: they are metadata for vLLM's config +plumbing, not a description of the acoustic stage. The acoustic stage holds no +vLLM attention layers, so no KV cache is derived from them; it reads its real +geometry from ``models/minimax_music3/constants.py`` and its weights from the +component folders. +""" + +from __future__ import annotations + +from typing import Any + +from transformers import AutoConfig +from transformers.configuration_utils import PretrainedConfig + + +class MiniMaxMusic3Config(PretrainedConfig): + """Root config for the multi-component MiniMax Music 3 repository.""" + + model_type = "minimax_music3" + + def __init__(self, **kwargs: Any) -> None: + # AR backbone geometry (language_model/config.json). Kept in sync with + # the checkpoint so max_model_len validation sees the real context. + kwargs.setdefault("hidden_size", 4096) + kwargs.setdefault("num_hidden_layers", 36) + kwargs.setdefault("num_attention_heads", 32) + kwargs.setdefault("num_key_value_heads", 8) + kwargs.setdefault("head_dim", 128) + kwargs.setdefault("intermediate_size", 12288) + kwargs.setdefault("vocab_size", 200000) + kwargs.setdefault("max_position_embeddings", 10240) + kwargs.setdefault("tie_word_embeddings", False) + super().__init__(**kwargs) + + # Audio contract, mirroring models/minimax_music3/constants.py. These + # are fixed by the checkpoint; they are exposed here only so a caller + # inspecting the config sees the same numbers the model uses. + self.audio_frame_rate = 25 + self.num_codebooks = 8 + self.audio_vocab_size = 1024 + self.c0_vocab_size = 16384 + self.sample_rate = 32000 + self.audio_channels = 2 + + +AutoConfig.register("minimax_music3", MiniMaxMusic3Config) + +__all__ = ["MiniMaxMusic3Config"]