Is this page helpful?

core.inference.disaggregation.engine#

Dynamic inference engine composed with the KV/state hand-off behavior.

Module Contents#

Classes#

StateHandoffDynamicInferenceEngine

Dynamic engine with control-plane-neutral KV/state handoff behavior.

DisaggDynamicInferenceEngine

Dynamic inference engine with prefill/decode state hand-off support.

API#

class core.inference.disaggregation.engine.StateHandoffDynamicInferenceEngine(
controller: megatron.core.inference.text_generation_controllers.text_generation_controller.TextGenerationController,
context: megatron.core.inference.contexts.dynamic_context.DynamicInferenceContext,
)#

Bases: megatron.core.inference.disaggregation.inference_state_handoff.InferenceStateHandoffMixin, megatron.core.inference.engines.dynamic_engine.DynamicInferenceEngine

Dynamic engine with control-plane-neutral KV/state handoff behavior.

External control planes can use this class directly and configure the transfer role with :meth:setup_kv_transfer. It intentionally does not install the coordinator-native registration protocol below.

Initialization

_record_and_serialize_finished_request(
request: megatron.core.inference.inference_request.DynamicInferenceRequest,
) → dict | None#

A prefill handoff is intermediate, not a ledger/stager completion.

_send_requests_to_coordinator(
requests: list[megatron.core.inference.inference_request.DynamicInferenceRequest],
) → None#

Register external handoffs before publishing their source descriptors.

class core.inference.disaggregation.engine.DisaggDynamicInferenceEngine(
controller: megatron.core.inference.text_generation_controllers.text_generation_controller.TextGenerationController,
context: megatron.core.inference.contexts.dynamic_context.DynamicInferenceContext,
)#

Bases: core.inference.disaggregation.engine.StateHandoffDynamicInferenceEngine

Dynamic inference engine with prefill/decode state hand-off support.

This subclass adds the coordinator-native registration callbacks. External control planes should use :class:StateHandoffDynamicInferenceEngine.

Initialization

_record_and_serialize_finished_request(
request: megatron.core.inference.inference_request.DynamicInferenceRequest,
) → dict | None#

Send handoff metadata and canonical tokens on the internal decode hop.

set_disaggregation_config(
*,
role,
identity,
spawn_coordinator,
coordinator_group,
) → None#

Configure one coordinator-native prefill or decode instance.

_build_instance_transfer_meta() → list[dict]#

Flatten this instance’s PP/TP transport descriptors for registration.

_notify_kv_read_done(request_id: int) → None#

Tell the native coordinator that decode no longer needs source storage.

_notify_kv_transfer_ready(
request_id: int,
cached_prefix_blocks: int,
) → None#

Tell the native coordinator that NCCL destination storage is committed.