anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[Model] Add MiniMax Music 3 text-to-music (#6186)

Browse Source 
Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
This commit is contained in:
Yueqian Lin authored and GitHub committed 2026-08-18 04:33:13 +00:00
1 parent 199ef7779c
commit b283215c28
43 files changed
+6876 -405
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

docs/models/supported_models.md
+1
View File
@@ -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` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
recipes/MiniMaxAI/MiniMax-Music3.md
+173
View File
@@ -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 |
tests/dfx/reliability/invalid_param_test/test_invalid_audio_speech.py
+65
View File
@@ -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}
)
tests/e2e/online_serving/test_minimax_music3.py
+112
View File
@@ -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)
tests/engine/test_async_omni_engine_input.py
+103
View File
@@ -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()
tests/entrypoints/openai_api/test_tts_detection.py
+1
View File
@@ -153,6 +153,7 @@ _PIPELINE_STAGES = [
"indextts2_talker",
"latent_generator",
"llm",
"minimax_music3_ar",
"ming_tts",
"moss_tts",
"moss_tts_codec",
tests/model_executor/models/test_minimax_music3_batch_payload.py
+132
View File
@@ -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]
tests/model_executor/models/test_minimax_music3_frame_state.py
+214
View File
@@ -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
tests/model_executor/models/test_minimax_music3_prompt.py
+203
View File
@@ -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
tests/model_executor/models/test_minimax_music3_row_binding.py
+217
View File
@@ -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]
tests/model_executor/test_cfg_pairing.py
+389
View File
@@ -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) == ()
vllm_omni/config/model.py
+5
View File
@@ -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):
vllm_omni/config/pipeline_registry.py
+2
View File
@@ -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,
vllm_omni/data_entry_keys.py
+6
View File
@@ -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):
vllm_omni/deploy/minimax_music3.yaml
+130
View File
@@ -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
vllm_omni/deploy/minimax_music3_2gpu.yaml
+97
View File
@@ -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
vllm_omni/engine/arg_utils.py
+2
View File
@@ -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
vllm_omni/engine/async_omni_engine.py
+69 -17
View File
@@ -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,
vllm_omni/engine/stage_engine_core_proc.py
+5 -4
View File
@@ -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,
vllm_omni/engine/stage_init_utils.py
+48 -9
View File
@@ -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
vllm_omni/entrypoints/openai/tts_adapters/__init__.py
+1
View File
@@ -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,
vllm_omni/entrypoints/openai/tts_adapters/minimax_music3.py
+163
View File
@@ -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"]
vllm_omni/model_executor/models/audex/cfg.py
+16 -375
View File
@@ -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": <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 <unk> 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()
vllm_omni/model_executor/models/common/cfg_pairing.py
+460
View File
@@ -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": <id>}
uncond: {"cfg_role": "uncond", "cfg_pair_id": <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()
vllm_omni/model_executor/models/minimax_music3/__init__.py
+34
View File
@@ -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",
]
vllm_omni/model_executor/models/minimax_music3/acoustic.py
+459
View File
@@ -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},
)
vllm_omni/model_executor/models/minimax_music3/chunking.py
+96
View File
@@ -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",
]
vllm_omni/model_executor/models/minimax_music3/constants.py
+111
View File
@@ -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",
]
vllm_omni/model_executor/models/minimax_music3/dav.py
+294
View File
@@ -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]
vllm_omni/model_executor/models/minimax_music3/depth_graph.py
+137
View File
@@ -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
vllm_omni/model_executor/models/minimax_music3/dit.py
+484
View File
@@ -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
vllm_omni/model_executor/models/minimax_music3/frame_state.py
+158
View File
@@ -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"]
vllm_omni/model_executor/models/minimax_music3/pipeline.py
+103
View File
@@ -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",
]
vllm_omni/model_executor/models/minimax_music3/prompt.py
+172
View File
@@ -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"(?<!\*)\*([^*\n]+)\*(?!\*)", r"\1", line)
lines_out.append(line.rstrip())
text = "\n".join(lines_out)
return re.sub(r"^\s*[-*_]{3,}\s*$", "", text, flags=re.MULTILINE)
def _remove_markdown_format(text: str) -> 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",
]
vllm_omni/model_executor/models/minimax_music3/rvq_decoder.py
+179
View File
@@ -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"]
vllm_omni/model_executor/models/minimax_music3/sampling.py
+223
View File
@@ -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",
]
vllm_omni/model_executor/models/minimax_music3/staging.py
+66
View File
@@ -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]
vllm_omni/model_executor/models/minimax_music3/talker.py
+990
View File
@@ -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
vllm_omni/model_executor/models/minimax_music3/weights.py
+133
View File
@@ -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",
]
vllm_omni/model_executor/models/registry.py
+11
View File
@@ -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",
),
}
vllm_omni/model_executor/stage_input_processors/minimax_music3.py
+544
View File
@@ -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
vllm_omni/transformers_utils/configs/__init__.py
+3
View File
@@ -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
vllm_omni/transformers_utils/configs/minimax_music3.py
+65
View File
@@ -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"]