Is this page helpful?

core.inference.contexts.mtp_context_mixin#

Context-side half of the MTP draft-KV cache.

MTPContextMixin owns the draft KV plane’s initialization, block reservation, prefix-match policy, and per-forward bookkeeping. DynamicInferenceContext calls these helpers from its shared request lifecycle and scheduling paths. Draft forwards publish their own write maps and attention metadata without disturbing the main forward’s state.

Module Contents#

Classes#

MTPContextMixin

MTP draft-KV bookkeeping for DynamicInferenceContext.

API#

class core.inference.contexts.mtp_context_mixin.MTPContextMixin#

MTP draft-KV bookkeeping for DynamicInferenceContext.

static should_enable_mtp_kv_cache(
model_config,
mtp_layer_type_list: Optional[List[str]],
num_speculative_tokens: int,
) → bool#

Whether this model/config can populate an MTP draft-KV plane.

Parameters:

  • model_config – Transformer config, read for the MTP head’s shape.

  • mtp_layer_type_list (Optional[List[str]]) – Layer types of one MTP draft-head depth for a hybrid model; None for a non-hybrid one.

  • num_speculative_tokens (int) – Draft depth; 0 disables speculative decoding.

Returns:

True when the draft KV plane should be reserved.

Return type:

(bool)

_mtp_initialize_kv_cache(
model_config,
mtp_layer_type_list: Optional[List[str]],
) → None#

Reserve the MTP attention plane before sizing the shared KV buffer.

_mtp_initialize_metadata() → None#

Create MTP metadata after the context’s request and block limits are known.

_mtp_allocate_buffers() → None#

Allocate fixed-address hidden-state and draft-loop buffers with the context.

_mtp_prefill_reserve_blocks(
req: megatron.core.inference.inference_request.DynamicInferenceRequest,
prefill_chunk_length: int,
) → int#

Extra blocks needed by the draft loop after the completing prefill chunk.

_mtp_trim_prefix_match(
req: megatron.core.inference.inference_request.DynamicInferenceRequest,
matched_block_ids: List[int],
already_allocated_blocks: int,
) → Tuple[List[int], int]#

Return the inheritable prefix and the number of blocks declined for MTP.

property _decode_kv_lookahead_tokens: int#

Tokens to reserve beyond the committed prefix for the next step.

Verification writes D+1 tokens. With full acceptance, the following D draft forwards write through offset 2D-1 relative to the next main input. Reserve both before scheduling that main forward, so pausing precedes either write.

_mtp_next_decode_block_counts(request_slice: slice) → torch.Tensor#

Blocks needed to cover verification and drafting after full acceptance.

_allocate_mtp_decode_blocks(request_slice: slice) → None#

Grant budgeted lookahead without advancing the main write pointer.

_mtp_map_next_decode_tokens(request_slice: slice) → None#

Map main inputs by position, leaving draft-only blocks as spare capacity.

_mtp_forward_phase()#

Scope one step’s MTP forwards and undo their effects on the context on the way out.

Every MTP forward republishes the active attention metadata, token counts and CUDA-graph flag, and leaves forward_mode set so KV append/read route to the draft plane. Log-prob computation runs after this scope and reads those fields expecting the MAIN forward’s values, and a forward_mode left set would keep is_decode_only False and KV routing on the draft plane for the rest of the run. Both are undone on exit, including when a draft forward raises.

_mtp_activate_attn_metadata(
graphed: bool,
padded_request_count: int,
max_seqlen_q: int,
max_seqlen_k: int,
token_count: int,
padded_token_count: int,
) → core.inference.contexts.attention_context.mha_metadata.MHAMetadata#

Publish the attention metadata for one MTP forward and return the active MHA object.

The buffer writes above target the fixed-address gpu_view buffers that the append/attend kernels read; those buffers are shared by the graphed and non-graphed MHA metadata objects, so a captured MTP graph replays correctly against whatever positions were just written. Only the metadata object selected here, the sequence-length bounds, and the token counts vary per step.

_mtp_begin_decode(
active_request_count: int,
padded_count: int,
start_positions: torch.Tensor,
graphed: bool = False,
) → None#

Enter MTP-forward mode.

start_positions[r] == base_position[r] - 1 is the MTP write position for depth 0 of request r (derived from the main KV offsets by the caller), advanced by one per depth.

graphed mirrors the main decode step’s CUDA-graph decision (the caller passes it from the EP-synced _mtp_resolved_padded_count, NOT the live _using_cuda_graph_this_step which the commit pass has already clobbered). When True, _mtp_setup_decode_step routes to the graph attention metadata so the captured KV-aware MTP graph is replayed.

_mtp_setup_decode_step() → None#

Populate token write maps + MHA read metadata for one MTP draft depth.

_mtp_setup_prefill_step(
append_counts: torch.Tensor,
block_table_prefill: torch.Tensor,
padded_token_count: Optional[int] = None,
padded_request_count: Optional[int] = None,
request_start_positions: Optional[torch.Tensor] = None,
total: Optional[int] = None,
) → None#

Populate token write maps + MHA metadata for a varlen roll-by-one MTP write forward.

Each request writes append_counts[r] consecutive positions starting at request_start_positions[r], or at 0 when that is None (a fresh prompt seed).

append_counts/block_table_prefill are GPU tensors for the P requests in active-slice order. total is append_counts.sum(); pass it when the caller already has it, to skip a redundant device sync.