Several codec frames inside one stage-1 execute_model.
One MiniCPM-o Talker decode step is almost all vLLM host work wrapped around a small device forward: a 190M-parameter model that emits exactly one codec frame, but pays scheduler, metadata, sampling and IPC costs per step. Over the ~118 frames of a typical request that host work is the largest single term after TTFT.
None of it is per frame. The scheduler step, input preparation, attention metadata, sampling, output assembly and engine-core IPC are per step, and a step could just as well carry K frames. This module runs them.
How a frame advances without the host¶
The step is scheduled with K query positions per request (stage 1's speculative_config in the deploy YAML). vLLM prepares all K of them -- positions, KV slots, block table, sequence lengths -- and the captured graph is a K-token uniform decode. What vLLM cannot do is fill in the inputs: frame k+1's embedding is the codec token sampled at frame k, so the K positions are sequential, not speculative.
So the loop replays the same captured K-token graph K times. Before replay k it writes frame k's embedding into row k of the persistent inputs_embeds buffer the graph captured; after replay k it samples frame k from row k of the hidden states. Rows above k still hold stale embeddings during replay k, and their KV is stale too -- but the attention bias is causal, so row k never reads them, and the next replay overwrites them with the real thing. After the K'th replay every row and every KV entry is what a K-step decode would have written.
A replay of the 20-layer forward is 0.83 ms of device and microseconds of host, so K replays cost what K frames cost and the per-step host work is paid once.
Nothing in the loop touches the host: the embedding lookup, the codec sample (the captured codec-sampling step), the stop row and the emitted delta are all device tensors enqueued in order on one stream, and the runner's existing coalesced D2H carries them out at the end of the step.
What makes it correct¶
- Stop is still per frame.
make_omni_outputevaluates the codec EOS and the frame ceiling on every frame exactly as it does today. A request that ends at frame j reports its stop row there; frames j+1..K-1 take the "already finished" early return, emit an empty delta and report a stop row too. vLLM's rejection sampler sees the first stop and truncates the request to j+1 tokens, which is the same sequence a one-frame-per-step run would produce. The ceiling is not only the stage-resolved codec budget: the runner also passes each request's remaining output budget (SamplingParams.max_tokens - num_output_tokens) and the codec merge clamps it intostate["max_tokens"], so a request limit that falls inside a K-frame step ends the loop where the engine stops accepting ids instead of emitting frames past it (PR #7929 review). - The emitted codes are unchanged. The extra frames after a stop contribute nothing: their deltas are empty or the
-1sentinel the connector already drops (_codec_scalars). So the audio is expected to be bit-identical, not merely close -- which is a far stronger check than the CER gate.
Off when stage 1's deploy config carries no speculative_config.
dataclass ¶
What one codec frame has to change before the one-query graph replays.
A K-query graph replay costs 1.33-1.47 ms of device time against the one-query graph's 0.83 ms, and it is a step change rather than a per-row cost -- a two-query graph costs the same as a four-query one. The loop only ever reads one row of the wide graph, so it replays the narrow one instead and moves the five things a frame changes:
inputs_embeds[r]-- request r's embedding for this framepositions[r]-- its position, +frameslot_mapping[r]-- the KV slot vLLM reserved for this frameseq_lens[r]-- its length, +frame+1, in the buffer the graph reads- the block table -- unchanged for the whole step
The step's own K-wide layouts are cloned at begin because the narrow views alias them: slot_mapping[:rows] is the first row of slot_mapping.view(rows, K).
Frames this step should run, or 0 when the loop must not engage.
Returns the uniform per-request query length of a pure decode step over a Talker that can build its own decode embedding. Anything else -- a prefill, a mixed step, a batch whose requests were scheduled different token counts -- returns 0 and the caller takes the ordinary single-forward path.
Every refusal says so once. A step that schedules several tokens per request and then takes the single-forward path is not merely slow: that path samples one frame from the last row of each span and reports one stop row for a step that needs one per position, so it corrupts the codec stream and the scheduler's accounting alike. _model_forward turns those into an error rather than letting them run.
begin_narrow_step(
*,
runner: Any,
forward_context: Any,
frames: int,
positions: Any,
inputs_embeds: Any,
spans: list[tuple[int, int]],
) -> NarrowStep | None
Point the forward context at the one-query graph, or decline.
Declines -- and the caller falls back to the wide replay -- whenever the narrow graph for this row count is not captured, which is what makes this safe to land before every path that reaches it has been exercised.
constant_drafts(
valid_sampled_token_ids: Any,
frames: int,
num_reqs: int,
scheduled_draft_counts: list[int] | None = None,
fold_short_rows: bool = False,
) -> list[list[int]]
continue repeated -- for the whole batch, or for none of it.
The drafts are how the next step gets K query positions per request, and the multi-frame loop can only run a step whose requests all have the same span: it replays one captured uniform-decode graph and samples one codec frame per position. So drafting per request is not an option.
scheduled_draft_counts[i] is how many draft positions the scheduler actually gave request i in the step that produced row i -- the authoritative split between the two one-token shapes that must not be confused:
- a row shorter than the step's own schedule (1 sampled +
scheduleddrafts) on a step that did carry drafts is a stop row firing mid-span: the rejection sampler truncated it to the frames it accepted, so its next-step schedule is j+1 wide against the neighbours' K -- non-uniform, fold the batch. A row exactly as wide as the step scheduled is not a stop even belowframes: the scheduler itself shortened the span (the uniform rewrite inomni_ar_scheduler), and this step must emit drafts or the scheduler's booked frames never reach stage 2 and the requests starve; - a row of length 1 on a step that carried no drafts is the expected shape of a prefill result or of an ordinary one-frame step (an intentional single-frame fallback, or the step right after such a fallback). Folding here would starve the K-step loop forever: the next step schedules one token again, produces another one-token row, and no drafts are ever emitted (Buildkite three-tier bench: 31/32 requests failed on exactly this stall). A draftless step is precisely where the K-step span starts or resumes, so those rows draft.
Mixed batches -- some requests drafted, some not -- cannot share a span and fold. So: every request drafts, or nobody does.
fold_short_rows gates the stop-truncated-row folding above. The caller (npu_ar_model_runner.propose) is the only place that can tell a live short row -- a request that already ran multi-frame and fired its codec stop mid-span, which must fold or the neighbours schedule against a dead speaker -- from a row that merely looks short because its step's drafts were scheduler padding: that request re-drafted is what keeps the K-step alive across the padding step. Passing the flag False keeps the pre-split behaviour of never folding on row length.
When scheduled_draft_counts is None the caller predates the scheduler split and every row is assumed to sit on a drafted step.
Frames the next step should be scheduled for, or 0 to stay generic.
The count travels to the scheduler as speculative tokens because that is vLLM V1's only way of saying "this request advanced by more than one token", which is what grows the block table and reserves the KV slots the loop writes into. Nothing about it is speculative: the drafts are always continue, and the frames are generated sequentially rather than verified.
Let vLLM's rejection-sampler output survive the Talker's narrow head.
RejectionSampler.parse_output keeps a sampled token only while it is both not the placeholder and < vocab_size, and stage 1's input_batch reports a vocab size of 0 -- the Talker's vLLM-level head is the two-wide continue/stop row compute_logits builds, not a tokenizer vocabulary, and nothing had ever needed the number. With 0 the mask is empty, every accepted token is dropped, and the request comes back having generated nothing: the scheduler then never rolls back the tokens it advanced for the step and schedules a negative count two steps later.
A single-frame step never reached this. It returns one token per request, which _bookkeeping_sync takes down the max_gen_len == 1 branch, and that branch does not consult vocab_size at all.
Raising it back to the width of the row the model actually emits is the whole fix -- and that width is STOP_ROW_WIDTH (two), not the width of whatever tensor the caller happens to be holding. The Talker's vLLM-level head is the two-wide continue/stop row compute_logits builds, so two is what every vocab_size reader has to see.
The width has to be the row the model emits, not the width of whatever tensor the caller happens to be holding: InputBatch.add_request stores top_k = vocab_size as its "no top-k" sentinel, which only holds for 0 and 2. A hidden width (768) instead passes the parse_output filter by accident but leaves top_k_reqs non-empty, and the Ascend enable_reduce_sample branch then treats the two-column distribution as a 768-way one -- the stop row does not survive into the request's token list and the request runs to max_tokens.
True when this step schedules several tokens for a request that is decoding.
Scoped to the Talker. Stage 0 drafts the Thinker's text with n-grams and so schedules multi-token decode steps of its own, which vLLM handles perfectly well -- it is only the Talker's one-frame-per-position sampler that cannot.
Whether this worker should capture and replay the one-query graph.
Off on a part whose one-query capture is known to fault (_NARROW_BLOCKED_SOC_PREFIXES), so a capture can never brick a deployment; those parts keep the K-query captures and the wide replay.
neutralize_kstep_min_tokens(logitsprocs: Any) -> None
Stop vLLM's min_tokens layer from censoring the stop signal.
The yaml min_tokens: 50 is enforced by MinTokensLogitsProcessor writing -inf into the stop-token logits. It unmasks once len(output_token_ids) >= min_tokens, a counter owned by vLLM's spec bookkeeping; while that counter lags, the request's only stop signal stays masked and finished_reason can only be length (observed on stage 1: all length, zero stop).
The real minimum-length guard lives elsewhere: talker_codec_sample.greedy_codec_sample masks the codec EOS itself via state.step < min_tokens, using the model's own frame counter. This vLLM layer is therefore redundant and, under multi-frame decode, also swallows the stop signal -- clear its list and leave the minimum to the model-side guard.
min_toks is re-registered by update_state every step, so this must be called every step (the caller places it right before sampling). Never raises.
run(
*,
model: Any,
run_model: Callable[[], Any],
after_forward: Callable[[], None],
inputs_embeds: Tensor | None,
input_ids: Tensor | None,
frames: int,
model_kwargs: dict[str, Any],
model_kwargs_extra: dict[str, Any],
narrow: NarrowStep | None = None,
) -> Any
Run frames codec frames and return the merged step output.
after_forward is the runner's post-forward graph bookkeeping. It is called after every replay because that is where it sits in the single-frame path; under static-shape decode it is two pointer comparisons, and the frames of one step share their sequence lengths so there is nothing for it to rebind.