991 lines
45 KiB
Python
991 lines
45 KiB
Python
# 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
|