core.inference.text_generation_controllers.mtp_controller_mixin

Multi-token-prediction (MTP) speculative decoding for dynamic inference.

MTPControllerMixin holds the MTP half of TextGenerationController: the draft-KV commit pass, the serial draft loop that produces speculative tokens, and the dummy forwards that keep idle expert-parallel ranks in lockstep with them. It is a mixin rather than a standalone helper because these paths read a large amount of controller state (sampling buffers, the wrapped model, SP/PP topology); the split is for readability, and mixing in preserves the previous behaviour exactly.

The only state the mixin writes is its own: the sampling/draft buffers allocated by _init_mtp_sampling_tensors. Everything else it touches belongs to the inference context or the model – including the chunked-prefill boundary carry, which lives on the context’s MTPMetadata so that it is invalidated alongside every other piece of MTP state.

Module Contents

Classes

_CommitSegments	One request class’s contribution to the commit pass, in active-request order.
MTPControllerMixin	MTP speculative-decoding paths for TextGenerationController.

API

class core.inference.text_generation_controllers.mtp_controller_mixin._CommitSegments

One request class’s contribution to the commit pass, in active-request order.

The commit pass concatenates the decode and prefill contributions into a single varlen forward, so both are described the same way: one (count, start) pair per request, plus the packed hidden/token rows those counts index into.

Parameters:

• counts (Tensor) – [R] tokens each request rewrites this step. May be 0 for a request with nothing to refresh; the entry is still required, because the commit pass’s segment count must stay equal to the active request count.

• starts (Tensor) – [R] MTP write position of each request’s first rewritten token.

• hidden (Tensor) – [T, 1, H] packed main hidden states; empty when T == 0.

• tokens (Tensor) – [T] next-token ids paired with hidden; empty when T == 0.

counts: torch.Tensor

None

starts: torch.Tensor

None

hidden: torch.Tensor

None

tokens: torch.Tensor

None

class core.inference.text_generation_controllers.mtp_controller_mixin.MTPControllerMixin

MTP speculative-decoding paths for TextGenerationController.

_init_mtp_sampling_tensors()

Pre-allocate MTP sampling tensors.

Addresses must be stable across steps for CUDA graph capture.

_mtp_commit_pass(

context,

unwrapped_model,

gathered_hidden,

num_decode_requests: int,

active_request_count: int,

base_position,

) → bool

Refresh every committed position’s draft KV from the MAIN hidden (a per-step first pass).

One varlen roll-by-one forward over all active requests’ committed tokens, so a committed position p always holds draft KV = f(h_p^main + emb(t_{p+1})) and never a stale chained-draft-hidden value (the cause of depth-increasing acceptance decay). K/V are pure projections of the input (no RoPE), so only the write positions matter; the output hidden is discarded. _mtp_decode_commit_segments and _mtp_prefill_commit_segments derive the per-request positions.

Returns True if it issued a forward (caller runs a dummy slot otherwise for EP balance).

_mtp_decode_commit_segments(

context,

gathered_hidden,

num_decode_requests: int,

base_position,

stride: int,

) → core.inference.text_generation_controllers.mtp_controller_mixin._CommitSegments

Refresh the accepted-draft positions of every decode request from this step’s hiddens.

Request r accepted a_r drafts, whose KV was written during the previous step’s draft loop from chained draft hiddens. Rewrite the earlier a_r of those positions from the main hiddens now that they are committed; the last accepted position is (re)written by decode depth 0, so it is excluded here.

_mtp_prefill_commit_segments(

context,

gathered_hidden,

decode_len: int,

num_decode_requests: int,

active_request_count: int,

) → core.inference.text_generation_controllers.mtp_controller_mixin._CommitSegments

Seed each prefill request’s prompt positions from main hiddens (roll-by-one).

Entry p is f(h_p, emb(t_{p+1})): its two inputs are one position apart. With off = request_kv_length_offsets and q = this chunk’s request_query_lengths, pairing each of the chunk’s hiddens with the token after it seeds positions off .. off+q-2. The chunk’s last hidden is dropped because its partner token is not in this chunk – the next chunk writes it as a seam, or draft depth 0 does on the final chunk.

SEAM: position off-1 needs the PREVIOUS chunk’s last hidden, so only this chunk can write it, and only if that hidden was carried over – i.e. this same request computed off-1. At most one request per step qualifies, since chunked prefill admits one in-flight request.

A prefix-cache hit also gives off > 0 but writes no seam: the skipped prefix came from a DIFFERENT request, so no carry matches. None is needed – _compute_prefix_match drops the last hash-matched block, which leaves off-1 already correct in the inherited block and moves the divergent entry into one the body rows below write. Removing that back-off silently breaks this.

Every request contributes exactly ONE segment (q-1 entries, or q with a seam), so the segment count stays active_request_count and the fixed-size MHA buffers never overflow. Also refreshes the carried chunk-boundary hidden for the in-flight chunked request.

_mtp_commit_forward(

context,

unwrapped_model,

packed_hidden: torch.Tensor,

packed_tokens: torch.Tensor,

append_counts: torch.Tensor,

request_start_positions: torch.Tensor,

total: int,

active_request_count: int,

num_decode_requests: int,

) → None

Run the packed commit-pass rows through the MTP layer to populate draft K/V.

Pads to a TP multiple and scatters for sequence parallelism, publishes the varlen write metadata on the context, and issues the forward. The output hidden is discarded – only the KV write matters.

_mtp_dummy_prefill_forward(context, unwrapped_model) → None

Issue one MTP-layer forward with dummy tensors and NO KV append.

Every rank must run exactly one MTP “prefill-slot” forward per step so the MoE/EP all-to-alls stay balanced when some ranks seed a real prompt and others (idle ranks, or active ranks with no prefill this step) do not. Uses inference_context=None so the attention runs cache-free (no child append); only the MoE all-to-all matters here. Mirrors the real seed’s single forward_single_position call.

_compute_serial_mtp_and_sample(

base_position: Optional[torch.Tensor] = None,

) → None

Run the MTP phase with the main forward’s execution state held across it.

_compute_serial_mtp_and_sample_impl(

base_position: Optional[torch.Tensor],

) → None

Compute MTP logits serially after verification and sample speculative tokens.

This ensures that MTP predictions are always conditioned on verified tokens. Each MTP depth receives the correctly sampled token from the previous depth (or the base token for depth 0) rather than stale speculative tokens from the previous step.

When sequence parallelism is active, hidden states are kept in SP format (scattered along the first dimension) between MTP depths to avoid a redundant gather + scatter round-trip per depth.

Parameters:

base_position (Optional[Tensor]) – GPU position of the first new MTP draft for each request. Legacy scheduling derives it from rewound CPU state.

_run_dummy_serial_mtp_forward() → None

Run dummy MTP forward passes to participate in EP collectives.

When speculative decoding is active and MTP layers contain MoE sublayers (inherited from the decoder layer spec), each serial MTP step triggers EP all-to-all collectives. The dummy EP rank must issue matching collective calls so the real ranks do not hang.

This mirrors the structure of _compute_serial_mtp_and_sample:

• On the last PP stage (where MTP resides): run compute_mtp_single_step with dummy tensors so the MoE all-to-all is executed.

• When PP > 1: participate in the broadcast_from_last_pipeline_stage that the real ranks also perform.