`core.inference.disaggregation.inference_state_handoff`#

Engine-side lifecycle for disaggregated prefill/decode state handoff.

Module Contents#

Classes#

`_PreparedHandoffMetadata`	Per-request metadata assembled once for a completed prefill batch.
`InferenceStateHandoffMixin`	Optional KV/SSM handoff behavior composed into the dynamic engine.

Data#

_SSM_STATE_KINDS

API#

core.inference.disaggregation.inference_state_handoff._SSM_STATE_KINDS#: (‘conv’, ‘recurrent’)

class core.inference.disaggregation.inference_state_handoff._PreparedHandoffMetadata#

Per-request metadata assembled once for a completed prefill batch.

local_blocks: list[int]#: None

local_ssm_slot: int | None#: None

kv_meta: Any#: None

ssm_meta: Any#: None

resume_tokens: list[int]#: None

resume_log_probs: list[float]#: None

class core.inference.disaggregation.inference_state_handoff.InferenceStateHandoffMixin#

Optional KV/SSM handoff behavior composed into the dynamic engine.

requires_recurrent_state_dummy_slot#: True

_finalize_request( request_entry: megatron.core.inference.engines.dynamic_engine.RequestEntry, ) → megatron.core.inference.inference_request.DynamicInferenceRequest#: Attach merged RL history before the completion is consumed or sent.

_initialize_disaggregation_state() → None#: Initialize state without importing or constructing a transfer backend.

_get_disaggregation_config() → megatron.core.inference.disaggregation.coordinator_setup.NativeDisaggregationConfig | None#

_is_disaggregated_role(role: str) → bool#

_handoff_push_error_source_safe(error: Exception) → bool#

_notify_kv_read_done(request_id: int) → None#: Hook for control planes that release source storage after transfer completion.

_notify_kv_transfer_ready( request_id: int, cached_prefix_blocks: int, ) → None#: Hook used when decode commits storage for a two-sided transfer.

_setup_handoff_completion_tracking( hostname: str | None = None, ) → None#: Create the CPU path used to aggregate model-parallel transfer completion.

_drain_handoff_completion_notifications() → list[tuple[int, bool, bool]]#

Collect decisions that the existing MP schedule broadcast must distribute.

Side: decode MP coordinator; pull and push transport paths.

_record_handoff_completion_notification( request_id: int, failed: bool, source_safe: bool, ) → None#

Record the coordinator’s shared admission decision on every MP rank.

Side: decode engine; pull and push transport paths.

property pending_kv_import_count: int#

Number of decode requests waiting for capacity or transfer completion.

Side: decode engine; pull and push transport paths.

property has_admittable_kv_import: bool#: Whether a completed import can be handled at the next admission point.

property pending_kv_push_count: int#

Number of prefill sends waiting for transport completion.

Side: prefill engine; push transport path only.

property pinned_handoff_block_count: int#: Number of prefill KV blocks retained for unfinished handoffs.

_reset_pending_kv_imports() → None#: Drain and release pending handoff transfers before an engine reset.

_pending_prefill_block_hashes() → set[int]#: Prevent recovery prefill from duplicating blocks reserved by imports.

schedule_waiting_requests() → None#

Recover evicted decodes locally, but reject new prompt submissions.

Side: decode engine; pull and push transport paths.

setup_kv_transfer( role: str, *, require_handoff_claim: bool = False, ) → None#

Bring up the KV transfer agents for this engine.

This method must be called collectively by every model-parallel rank in the engine because each rank participates in PP and TP metadata gathers.

Parameters:

role – “prefill” or “decode”; used to name the local transfer agent.
require_handoff_claim – The external control plane guarantees a successful source-coordinator claim before every remote read. Enables expiration of unclaimed offers; never expires readers.

push_handoff_kv( request_id: int, decode_metas: list, cached_prefix_blocks: int = 0, ) → None#

Push a pinned hand-off’s KV and live SSM state to the decode instance described by decode_metas (two-sided transports only).

The decode committed its destination storage before requesting these sends. Sends are reaped asynchronously and pins stay until RELEASE_KV.

Side: prefill engine; push transport path only.

_poll_pending_kv_pushes() → int#

Reap completed push sends; unfinished ones stay pending.

Side: prefill engine; push transport path only.

_prepare_handoff_metadata_batch( requests_and_state: list[tuple[megatron.core.inference.inference_request.DynamicInferenceRequest, list[int], int | None]], decode_tokens_by_request: Dict[int, list[int]], decode_log_probs_by_request: Dict[int, list[float]], ) → dict[int, core.inference.disaggregation.inference_state_handoff._PreparedHandoffMetadata]#

Assemble metadata for all handoffs completed by one engine step.

Side: prefill engine; pull and push transport paths.

static _ssm_stage_transfer_meta( static_metas: dict, selected_slot: int, ) → dict#: Attach the final recurrent-state slot to one stage’s descriptors.

_capture_handoff_meta( request: megatron.core.inference.inference_request.DynamicInferenceRequest, prepared: core.inference.disaggregation.inference_state_handoff._PreparedHandoffMetadata | None, ) → None#

Attach prepared transfer metadata and retain the request’s blocks.

Side: prefill engine; pull and push transport paths.

release_handoff_blocks(request_id: int) → None#

Release blocks pinned by a previous do_kv_handoff completion.

Side: prefill engine; pull and push transport paths.

_release_pinned_handoff_blocks(block_ids: list) → int#

Release this request’s ownership of its pinned handoff blocks.

Side: prefill engine; pull and push transport paths.

_release_pinned_handoff_ssm_slot(ssm_slot: int | None) → None#: Release a prefill live-state slot after its handoff ownership ends.

add_request_with_kv_handoff( request_id: int, prompt: list, sampling_params: megatron.core.inference.sampling_params.SamplingParams, kv_meta: dict, src_block_ids: list, *, offload_params: dict | None = None, ) → asyncio.Future[DynamicInferenceRequest]#

Start or capacity-queue a KV import and return its completion future.

Side: decode engine; pull and push transport paths. A pull backend starts the read here, while a push backend posts the matching receive.

_try_start_kv_handoff_import( handoff: megatron.core.inference.disaggregation.pending_handoff_imports.DeferredKvHandoff, ) → bool#

Reserve cache state and start one import, or return without mutation.

Side: decode engine; pull and push transport paths.

_reserve_handoff_hashes(request_id: int, hashes: list[int]) → bool#: Reserve unregistered hashes until one decode import publishes them.

_release_handoff_hash_reservations( request_id: int, hashes: list[int], ) → None#: Release every pending hash owned by one decode import.

_find_cached_handoff_prefix( hashes: list[int], num_blocks: int, ) → list[int]#

Find the contiguous handoff prefix already cached on decode.

Side: decode engine; pull and push transport paths.

_handoff_capacity_available( num_blocks: int, cached_blocks: list[int], ) → bool#

Check capacity in the rank-local mirror of model-parallel allocator state.

Side: decode engine; pull and push transport paths.

_drain_deferred_kv_handoffs() → int#

Start queued handoffs in FIFO order while the queue head fits.

Side: decode engine; pull and push transport paths.

abort_kv_handoff(request_id: int) → bool#: Cancel a deferred or transferring decode handoff.

static _pending_transfer_handles( pending: megatron.core.inference.disaggregation.pending_handoff_imports.PendingKvImport, ) → list#

Return this decode import’s active KV and SSM transfer handles.

Side: decode engine; pull and push transport paths.

static _pending_destinations_safe( pending: megatron.core.inference.disaggregation.pending_handoff_imports.PendingKvImport, ) → bool#: Return whether a failed transfer has stopped accessing its destinations.

_wait_for_pending_storage( pending: megatron.core.inference.disaggregation.pending_handoff_imports.PendingKvImport, ) → bool#: Wait for trackable operations and return whether storage can be reused.

_validate_decode_ready_handoff( pending: megatron.core.inference.disaggregation.pending_handoff_imports.PendingKvImport, ) → None#: Validate that transferred state can start decode without prompt execution.

_finalize_kv_handoff_import( pending: megatron.core.inference.disaggregation.pending_handoff_imports.PendingKvImport, ) → None#

Side: decode engine; pull and push transport paths.

_complete_handoff_request_without_forward(request_id: int) → None#: Complete an imported request whose transferred token already satisfies it.

_release_pending_kv_import( pending: megatron.core.inference.disaggregation.pending_handoff_imports.PendingKvImport, ) → None#

Release storage owned by an unadmitted decode import.

Side: decode engine; pull and push transport paths.

static _wait_for_transfer_handles(*handles) → bool#: Wait for known handles; return false if any may still access storage.

_report_completed_kv_imports() → None#

Report locally terminal imports without synchronizing the compute ranks.

Side: decode engine; pull and push transport paths.

_poll_pending_kv_imports() → int#

Progress imports without mutating the active decode batch.

Side: decode engine; pull and push transport paths.

_quarantine_kv_import( pending: megatron.core.inference.disaggregation.pending_handoff_imports.PendingKvImport, ) → None#: Retain unsafe storage, but stop reserving hashes this import cannot publish.

_poll_quarantined_kv_imports() → None#: Reclaim failed imports once their transfer handles become terminal.

_admit_pending_kv_imports() → int#

Handle completed imports at an engine scheduling admission point.

Side: decode engine; pull and push transport paths.

_reserve_ssm_handoff_import() → megatron.core.inference.disaggregation.pending_handoff_imports.PendingSSMImport#: Reserve the live request slot that receives exact SSM state.

_start_ssm_handoff_import( request_id: int, ssm_meta: dict, pending: megatron.core.inference.disaggregation.pending_handoff_imports.PendingSSMImport, ) → None#: Post transfers into slots already reserved for one handoff.

core.inference.disaggregation.inference_state_handoff#

Module Contents#

Classes#

Data#

API#

`core.inference.disaggregation.inference_state_handoff`#