`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

class core.inference.disaggregation.inference_state_handoff.InferenceStateHandoffMixin#

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

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

_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]]#

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, ) → 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.

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

schedule_waiting_requests() → None#

Reject prompt scheduling on a dedicated disaggregated decode engine.

Side: decode engine; pull and push transport paths.

setup_kv_transfer(role: str, backend: str = 'nixl') → 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.
backend – transfer backend name, resolved through the explicit registry (“nixl”; “nccl” selects the two-sided push family).

push_handoff_kv(request_id: int, decode_metas: list) → 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 posted its matching receives when SUBMIT_REQUEST_WITH_KV arrived; the sends are reaped asynchronously and the pins stay until the coordinator’s 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]], ) → 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, ) → 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.

_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 transport path only.

_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.

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.

_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 be active.

_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.

_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`#