`core.inference.engines.dynamic_engine`#

Module Contents#

Classes#

`EngineState`	State machine for the inference engine.
`_VisionCacheEntry`	Projected embedding and the preprocessed media needed to reuse it.
`RequestEntry`	Entry in the engine’s `self.requests` dict.
`DynamicInferenceEngineStepResult`	Result returned by modern dynamic-engine step APIs.
`DynamicInferenceEngine`	The dynamic inference engine.

Functions#

`_take_expanded_prefix_stitching_metadata`	Remove and return engine-private expanded-prefix stitching metadata.
`_prepend_expanded_multimodal_prefix`	Prepend exact prior model tokens to an expanded current-turn suffix.
`_slice_suffix_media_metadata`	Select metadata for media placeholders occurring in the compact suffix.
`format_mem_bytes`	Convert a byte count to a human-readable string in tb, gb, mb, kb, or bytes.
`_weight_scoped_salt`	Scope a request’s block hashes to the weight generation that will serve it.
`_engine_reply_frames`	Frame finished requests as [metadata, body, body, …] for the coordinator.
`_get_decode_only_log_state`	Build the console transition label and color state for one inference step.
`_cuda_graph_mempool_bytes`	Return (reserved, allocated) bytes belonging to the global CUDA graph mempool.
`_slice_merged_segments`	Slice `attr` over the concatenation of request checkpoint segments.

Data#

`_PROMPT_PREPARATION_ERROR_FIELD`
`_MULTIMODAL_STITCHING_FIELDS`
`_PACKED_NONE`
`_SUBMIT_REQUEST_FRAMES`
`_SUBMIT_REQUEST_METADATA_FIELDS`
`logger`
`DEPRECATED_ARGS`

API#

core.inference.engines.dynamic_engine._PROMPT_PREPARATION_ERROR_FIELD#: ‘_request_prompt_preparation_error’

core.inference.engines.dynamic_engine._MULTIMODAL_STITCHING_FIELDS#: None

core.inference.engines.dynamic_engine._take_expanded_prefix_stitching_metadata(offload_params)#: Remove and return engine-private expanded-prefix stitching metadata.

core.inference.engines.dynamic_engine._prepend_expanded_multimodal_prefix( prefix_tokens, suffix_tokens, suffix_mask, inference_wrapper, modality, )#: Prepend exact prior model tokens to an expanded current-turn suffix.

core.inference.engines.dynamic_engine._slice_suffix_media_metadata( prefix_media_count, *, num_tiles, imgs_sizes, num_frames, video_frame_indices, video_fps, )#: Select metadata for media placeholders occurring in the compact suffix.

core.inference.engines.dynamic_engine._PACKED_NONE#: b’\xc0’

core.inference.engines.dynamic_engine._SUBMIT_REQUEST_FRAMES#: 4

core.inference.engines.dynamic_engine._SUBMIT_REQUEST_METADATA_FIELDS#: 4

core.inference.engines.dynamic_engine.logger#: ‘getLogger(…)’

core.inference.engines.dynamic_engine.DEPRECATED_ARGS#: [‘enable_cuda_graph’, ‘random_seed’, ‘track_paused_request_events’, ‘enable_chunked_prefill’, ‘infer…

class core.inference.engines.dynamic_engine.EngineState(*args, **kwds)#

Bases: enum.Enum

State machine for the inference engine.

Initialization

RUNNING#: ‘auto(…)’

PAUSING#: ‘auto(…)’

PAUSED#: ‘auto(…)’

UNPAUSING#: ‘auto(…)’

SUSPENDING#: ‘auto(…)’

SUSPENDED#: ‘auto(…)’

RESUMING#: ‘auto(…)’

RESUMED#: ‘auto(…)’

STOPPING#: ‘auto(…)’

STOPPED#: ‘auto(…)’

exception core.inference.engines.dynamic_engine.EngineSuspendedError#

Bases: Exception

Engine is currently suspended and not performing steps.

Initialization

Initialize self. See help(type(self)) for accurate signature.

class core.inference.engines.dynamic_engine._VisionCacheEntry#

Projected embedding and the preprocessed media needed to reuse it.

embedding: torch.Tensor#: None

modality: str#: None

imgs: torch.Tensor#: None

num_tiles: Optional[torch.Tensor]#: None

num_img_embeddings_per_tile: int#: None

imgs_sizes: Optional[torch.Tensor]#: None

num_frames: Optional[torch.Tensor]#: None

video_frame_indices: Optional[List[List[int]]]#: None

video_fps: Optional[List[float]]#: None

core.inference.engines.dynamic_engine.format_mem_bytes(mem_bytes)#: Convert a byte count to a human-readable string in tb, gb, mb, kb, or bytes.

core.inference.engines.dynamic_engine._weight_scoped_salt( weight_epoch: int, media_cache_key: Optional[str], ) → Optional[str]#

Scope a request’s block hashes to the weight generation that will serve it.

Under KVCacheManagementMode.PERSIST the prefix cache survives a refit: reinitialize_inference_state_buffers() only resets metadata on the RECOMPUTE path, so the KV and Mamba hash tables outlive suspend/resume and a request admitted afterwards can match blocks whose KV the previous weights computed. Staleness is then bounded only by eviction pressure, since the engine-side allocator has no TTL.

Mixing the weight generation into the salt makes chains from different generations disjoint, so stale blocks become unmatchable rather than being freed – nothing is mutated at refit time, leaving live requests, chunked prefill and pending Mamba restores untouched.

Composed with the media key rather than replacing it: multimodal requests still need equal token placeholders backed by different media to stay unmatchable. Epoch 0 returns the media key unchanged, so an engine that never resumes hashes exactly as before.

core.inference.engines.dynamic_engine._engine_reply_frames( finished_requests: List[dict], ) → List[bytes]#

Frame finished requests as [metadata, body, body, …] for the coordinator.

The metadata frame carries only what the coordinator needs to route each reply: the request id, and whether it must detokenize into the body. Every body stays a separate opaque frame so the coordinator can forward it without decoding – a finished request echoes the prompt back, so decoding it costs more than the inbound submission did.

Parameters:: finished_requests – Serialized requests, in the order their frames follow.
Returns:: The frames to send, metadata first.

core.inference.engines.dynamic_engine._get_decode_only_log_state( mode: megatron.core.inference.config.AsyncScheduleMode, decode_only: megatron.core.inference.text_generation_controllers.text_generation_controller.DecodeOnly, ) → Tuple[str, Optional[bool]]#

Build the console transition label and color state for one inference step.

Parameters:

mode (AsyncScheduleMode) – Active scheduling mode.
decode_only (DecodeOnly) – Decode-only state for the consumed and launched forwards.

Returns:

Current step label, including the previous step when it differs, and whether to use decode coloring.

Return type:

Tuple[str, Optional[bool]]

core.inference.engines.dynamic_engine._cuda_graph_mempool_bytes() → Tuple[int, int]#

Return (reserved, allocated) bytes belonging to the global CUDA graph mempool.

PyTorch’s torch.cuda.memory_stats() reports process-wide totals that mix in every other allocation (KV cache, NCCL workspaces, layer scratch). To isolate growth caused by graph capture, we walk torch.cuda.memory_snapshot() and filter segments by their segment_pool_id against the graph pool handle. Returns (0, 0) if the pool hasn’t been created yet.

core.inference.engines.dynamic_engine._slice_merged_segments( segments, attr: str, start: int, end: int, ) → list#

Slice attr over the concatenation of request checkpoint segments.

Equivalent to merged[start:end] where merged is the per-attribute concatenation produced by DynamicInferenceRequestRecord.merge(), without materializing the full merged list.

class core.inference.engines.dynamic_engine.RequestEntry#

Entry in the engine’s self.requests dict.

record: megatron.core.inference.inference_request.DynamicInferenceRequestRecord#: None

future: asyncio.Future[megatron.core.inference.inference_request.DynamicInferenceRequest]#: None

class core.inference.engines.dynamic_engine.DynamicInferenceEngineStepResult#

Bases: typing.TypedDict

Result returned by modern dynamic-engine step APIs.

Initialization

Initialize self. See help(type(self)) for accurate signature.

active_request_ids: list[int]#: None

finished_requests: list[megatron.core.inference.inference_request.DynamicInferenceRequest]#: None

step_time: float#: None

cuda_graph_request_count: int | None#: None

class core.inference.engines.dynamic_engine.DynamicInferenceEngine( controller: megatron.core.inference.text_generation_controllers.text_generation_controller.TextGenerationController, context: megatron.core.inference.contexts.dynamic_context.DynamicInferenceContext, )#

Bases: megatron.core.inference.engines.abstract_engine.AbstractEngine

The dynamic inference engine.

This engine allows requests of varying length to be dynamically added and removed in each inference step. In contrast to the static engine that has a set batch size and sequence length during the forward pass, each request in the dynamic engine can have different current prompt and output length at any given step, and the processing is restricted only by a max number of total tokens across all requests.

Parameters:

text_generation_controller (TextGenerationController) – A text generation controller that will be used to define how to preprocess prompts, generate output tokens, and apply token-level generation policy.
inference_context (DynamicInferenceContext) – Context for managing in-flight batching and a dynamic block-level KV cache (similar to paged attention).

Initialization

requires_recurrent_state_dummy_slot#: False

_STATE_EVENTS#: ()

_weight_epoch: int#: 0

payload_stager: Optional[megatron.core.inference.inference_request.RequestPayloadStager]#: None

prompt_preparer: Optional[megatron.core.inference.inference_request.RequestPromptPreparer]#: None

_staging_tasks: Set[asyncio.Task]#: ‘frozenset(…)’

_staging_tasks_peak: int#: 0

_initialize_disaggregation_state() → None#: Hook overridden by the KV-handoff engine composition.

_get_disaggregation_config() → megatron.core.inference.disaggregation.coordinator_setup.NativeDisaggregationConfig | None#: Return native disaggregation setup, if configured.

_is_disaggregated_role(role: str) → bool#: Return whether this engine has the given disaggregated role.

_reset_pending_kv_imports() → None#: Hook overridden by the KV-handoff engine composition.

property pending_kv_import_count: int#: Number of decode requests awaiting a KV import (none here).

property has_admittable_kv_import: bool#: Whether a completed KV import is eligible for admission (false here).

_poll_pending_kv_imports() → int#

_admit_pending_kv_imports() → int#

_setup_handoff_completion_tracking( hostname: str | None = None, ) → None#: Hook overridden by the KV-handoff engine composition.

_drain_handoff_completion_notifications() → list[tuple[int, bool, bool]]#: Hook overridden by the KV-handoff engine composition.

_record_handoff_completion_notification( request_id: int, failed: bool, source_safe: bool, ) → None#: Hook overridden by the KV-handoff engine composition.

_prepare_handoff_metadata_batch( requests_and_state: list[tuple], decode_tokens_by_request: Dict[int, list[int]], decode_log_probs_by_request: Dict[int, list[float]], ) → dict#: Hook overridden by the KV-handoff engine composition.

_notify_request_error( request_id: int, error: BaseException, *, source_safe: bool = False, ) → None#: Report a terminal request failure from the model-parallel coordinator.

_notify_request_aborted(request_id: int, *, source_safe: bool) → None#: Report that cancellation can safely complete.

_capture_handoff_meta(request, prepared) → None#

_release_pinned_handoff_blocks(block_ids: list) → int#

_release_pinned_handoff_ssm_slot(ssm_slot: int | None) → None#

setup_kv_transfer(role: str) → None#: Raising stub; the hand-off engine composition overrides it.

push_handoff_kv( request_id: int, decode_metas: list, cached_prefix_blocks: int = 0, ) → None#: Raising stub; the hand-off engine composition overrides it.

_handoff_push_error_source_safe(error: Exception) → bool#: Return whether source storage is reusable after a push setup error.

_poll_pending_kv_pushes() → int#

property pending_kv_push_count: int#: Number of prefill sends awaiting completion (none here).

property pinned_handoff_block_count: int#: Number of source KV blocks retained by handoff, if enabled.

add_request_with_kv_handoff( request_id, prompt, sampling_params, kv_meta, src_block_ids, *, offload_params=None, ) → asyncio.Future[megatron.core.inference.inference_request.DynamicInferenceRequest]#: Raising stub; the hand-off engine composition overrides it.

release_handoff_blocks(request_id: int) → None#: Raising stub; the hand-off engine composition overrides it.

abort_kv_handoff(request_id: int) → bool#: Return whether a composed handoff engine consumed the cancellation.

static _raise_kv_handoff_not_enabled(operation: str) → None#

reset() → None#: Reset per-run state; the caller must first drain all requests and payload staging.

static _tensor_nbytes(tensor: torch.Tensor) → int#

_vision_cache_entry_nbytes( entry: core.inference.engines.dynamic_engine._VisionCacheEntry, ) → int#

clear_vision_embedding_cache() → None#: Release all projected embeddings and reusable media retained by this engine.

_invalidate_vision_state() → None#: Mark cached and request-local projected media as weight-stale.

_refresh_vlm_request_data( request: megatron.core.inference.inference_request.DynamicVLMInferenceRequest, ) → None#: Rebuild missing media state for a known-multimodal request.

static _can_defer_image_encode( *, imgs: torch.Tensor, num_tiles: Optional[torch.Tensor], imgs_sizes: Optional[torch.Tensor], num_frames: Optional[torch.Tensor], ) → bool#

Whether a request’s image encode can wait until prefill scheduling.

Only dynamic-resolution images with a concatenable patch layout ([P, F] or [1, P, F]) qualify; videos, static tiling and other layouts keep the per-request encode in _build_vlm_request.

_encode_waiting_images() → None#

Encode the images of waiting requests whose admission deferred the encode.

_build_vlm_request admits batchable image requests without running the vision encoder. This runs right before waiting requests are scheduled for prefill, so the images of every such request that arrived since the last prefill go through the encoder together: a request with several images already hands _forward_vision_encoder their patches concatenated with one imgs_sizes row per image, so the images of several requests are concatenated the same way and the output is split back by each request’s embedding count. One call covers at most context.max_tokens embeddings. Cache hits are filled from the vision cache, requests sharing a media_cache_key are encoded once, a group whose shared encode fails falls back to per-request encodes, and a request whose encode fails is failed on its own and removed from the waiting queue.

_encode_deferred_image_requests( deferred: List[megatron.core.inference.inference_request.DynamicVLMInferenceRequest], ) → None#: Encode, register or fail the given waiting requests (see _encode_waiting_images).

_batched_request_image_encode( group: List[Tuple[megatron.core.inference.inference_request.DynamicVLMInferenceRequest, int]], ) → Dict[int, torch.Tensor]#

Run one vision-encoder call over the images of several requests and split it back.

Parameters:: group – (request, embedding_count) pairs whose summed count fits the budget.
Returns:: Each request’s slice of the encoder output, keyed by request id.
Return type:: Dict[int, Tensor]

_encode_image_request( request: megatron.core.inference.inference_request.DynamicVLMInferenceRequest, ) → torch.Tensor#: Run the vision encoder for one request’s images (the per-request path).

_register_image_embeddings( request: megatron.core.inference.inference_request.DynamicVLMInferenceRequest, embeddings: torch.Tensor, mask: torch.Tensor, *, expected_count: int, modality: str, imgs: torch.Tensor, num_tiles: Optional[torch.Tensor], imgs_sizes: Optional[torch.Tensor], num_frames: Optional[torch.Tensor], error_subject: str, cache: bool = True, ) → None#: Check the embedding count against the mask, then store the embeddings on the request, in the vision cache and on the context (shared by the post-admission paths).

_get_cached_vision_entry( cache_key: Optional[str], modality: Optional[str] = None, ) → Optional[core.inference.engines.dynamic_engine._VisionCacheEntry]#: Return and promote a complete reusable vision-cache entry.

_get_cached_vision_embedding( cache_key: Optional[str], ) → Optional[torch.Tensor]#

_cache_vision_embedding( cache_key: Optional[str], embedding: torch.Tensor, *, modality: str, imgs: torch.Tensor, num_tiles: Optional[torch.Tensor] = None, num_img_embeddings_per_tile: int = 0, imgs_sizes: Optional[torch.Tensor] = None, num_frames: Optional[torch.Tensor] = None, video_frame_indices: Optional[List[List[int]]] = None, video_fps: Optional[List[float]] = None, ) → None#

async wait_until(state: core.inference.engines.dynamic_engine.EngineState)#

Wait until the engine reaches the given state.

Only stable states (RUNNING, PAUSED, SUSPENDED, RESUMED, STOPPED) are supported. Transient states (PAUSING, SUSPENDING, RESUMING, STOPPING) are not directly waitable.

create_cuda_graphs(reset_context: bool = True)#

Create cuda graphs.

This method iterates the dynamic context’s cuda_graph_request_counts to record and capture cuda graphs.

Parameters:: reset_context (bool) – Whether to reset the context after building cuda graphs.

async start_listening_to_data_parallel_coordinator( inference_coordinator_port: int | None = None, launch_inference_coordinator: bool = True, *, hostname: str | None = None, coordinator_schedule_output_path: str | None = None, loop: Optional[asyncio.AbstractEventLoop] = None, )#

Initializes ZMQ communication to connect the engine with an inference coordinator.

This asynchronous method sets up the distributed communication infrastructure that allows this inference engine to act as a worker under a central InferenceCoordinator. It configures different ZMQ socket patterns based on the rank’s role within the distributed topology.

Note that this method must be called on all ranks, as it uses blocking torch broadcasts.

The setup involves two primary roles within each data-parallel group:

MP Coordinator (TP_rank=0, PP_rank=0): This rank connects directly to the central coordinator via a ZMQ DEALER socket. It receives requests and uses a ZMQ PUB (publisher) socket to broadcast them to all other ranks within its model-parallel (MP) group.
MP Workers (all other ranks): These ranks use ZMQ SUB (subscriber) sockets to listen for requests broadcast by their local MP Coordinator.

This architecture uses TCP sockets for both inter-node and intra-node broadcasts within an MP group.

Finally, after setting up the communication channels and ensuring all ranks are synchronized, this method starts the main engine processing loop (self.run_engine) as a background asyncio task.

Parameters:

inference_coordinator_port (int | None) – The network port where the central InferenceCoordinator is or will be listening. If None, a random available port will be selected. If not None, the coordinator will attempt to bind to this port, but should it not succeed (e.g., if the port is already in use), it may bind to a different port. The actual port used is returned by this method.
launch_inference_coordinator (bool, optional) – If True, the global rank 0 process will spawn and manage the InferenceCoordinator process. Defaults to True.
hostname (str | None) – Hostname or IP address to use for ZMQ socket binding. If None, defaults to socket.gethostname(). Should be set to a routable address in multi-node settings where gethostname() may return 127.0.0.1.

Returns:

The network address of the central InferenceCoordinator, which may not have the same port as what the user requested with inference_coordinator_port.

Return type:

inference_coordinator_addresss (str)

static suspend_resume_ctx(key: str, *, unified_memory_level: int) → None#

Context manager for of suspending and resuming the engine.

This context manager records the time and memory usage when suspending and resuming the context. TODO(@lmcafee): add argument to optionally return nullcontext, to avoid overhead.

Parameters:: key (str) – Key that identifies caller (e.g., ‘suspend’ or ‘resume’).
Returns:: None.

suspend()#: Suspend engine by deallocating context’s GPU state.

resume()#: Resume engine by reallocating context’s GPU state.

_resalt_unprefilled_requests() → None#

Re-salt requests that hold no KV yet under the current weight epoch.

Such a request has its whole prompt computed by the weights that serve it after resume, so it must hash like a request admitted now. A request with any prefill or decode progress keeps its salt: its leading blocks were computed by the previous weights.

async _notify_cond_for_new_request()#: Helper function to notify condition variable when a new request is added.

_finalize_request( request_entry: core.inference.engines.dynamic_engine.RequestEntry, ) → megatron.core.inference.inference_request.DynamicInferenceRequest#: Record the end of a request whose status is terminal and merge its record.

_complete_request( request_entry: core.inference.engines.dynamic_engine.RequestEntry, finished_request: Optional[megatron.core.inference.inference_request.DynamicInferenceRequest] = None, *, reply_now: bool = False, ) → megatron.core.inference.inference_request.DynamicInferenceRequest#

Resolve a request future, finalizing its record unless already finalized.

Parameters:

request_entry – Engine-owned record and completion future.
finished_request – Result finalized before context cleanup, when available.
reply_now – Send the result to the coordinator now. Used for a request that ends outside a step: one that failed before it ran a step, a resumed generation that arrived finished, or a KV import already satisfied by its transferred token.

_send_requests_to_coordinator( requests: List[megatron.core.inference.inference_request.DynamicInferenceRequest], ) → None#: Send completed or failed flat requests from model-parallel rank 0.

_record_and_serialize_finished_request( request: megatron.core.inference.inference_request.DynamicInferenceRequest, ) → Optional[Dict]#

Record a completed request, then build its reply or hand it to a staging task.

Returns None when the payload is staged: that reply is sent by the staging task (see _stage_and_reply) once the stager’s write returns.

async _stage_and_reply( stager: megatron.core.inference.inference_request.RequestPayloadStager, request: megatron.core.inference.inference_request.DynamicInferenceRequest, payload: megatron.core.inference.inference_request.OffloadedRequestPayload, finished_metadata: megatron.core.inference.inference_request.FinishedRequestRecord, ) → None#

Stage one completed payload on a worker thread, then send its reply.

The payload stager’s write may be a slow network round trip (e.g. media). Running it here keeps the engine loop decoding the other requests; the reply still waits for this request’s write, so its stage metadata describes a durable write. The reply is sent after the await, back on the engine loop’s thread, which is the only thread that uses the ZMQ socket. A failed write keeps the payload on the reply, and a reply that cannot be built or sent is replaced by a FAILED reply. Nothing awaits this task, so any other error exits the process (like the engine loop) instead of vanishing and leaving the client waiting forever.

async _drain_staging_tasks() → None#: Wait for in-flight staging tasks so their replies are sent before teardown.

_staging_log_suffix() → str#

Step-log fragment for the staging backlog; resets the peak to the current size.

A backlog that keeps growing means the stager’s writes cannot keep up with the rate requests finish.

_serialize_finished_request( request: megatron.core.inference.inference_request.DynamicInferenceRequest, stage_result: Optional[megatron.core.inference.inference_request.RequestPayloadStageResult], ) → Dict#: Construct a completed request’s coordinator reply from its staging result.

_handle_failed_request( request_id: int, *, stage: str = 'be added to the engine', )#

Handle a failed request by sending the reply immediately.

The request is added to failed_request_ids so that the next bookkeeping pass can return it.

Parameters:

request_id (int) – The failed request.
stage (str) – What the request failed to do, for the operator warning.

_fail_submission( request_id: int, sampling_params: Optional[megatron.core.inference.sampling_params.SamplingParams], exc: BaseException, ) → None#

Called from the SUBMIT_REQUEST handler when image preprocessing or add_request raises. Registering a placeholder record with Status.FAILED lets _handle_failed_request publish the ENGINE_REPLY without leaving the client hanging or killing the engine loop.

_collect_failed_requests( request_ids: Optional[set[int]] = None, ) → List[megatron.core.inference.inference_request.DynamicInferenceRequest]#

Remove and return a snapshot of synchronously failed requests.

Parameters:: request_ids – Optional ownership filter. Failed requests outside this set remain queued for their caller.
Returns:: Failed requests selected from the current queue snapshot.

async snapshot_live_requests() → List[megatron.core.inference.inference_request.LiveRequestSnapshot]#: Snapshot every non-failed request for later resumption; must run on the engine loop whenever that loop is running (any loop may run it while the engine loop is stopped, as between a suspend and a resume).

_routing_from_blocks( request: megatron.core.inference.inference_request.DynamicInferenceRequest, ) → Optional[numpy.ndarray]#: A resident request’s expert routing, read from its blocks as its completion reads it.

has_unfinished_requests() → bool#: Test if context contains unfinished requests.

get_request( request_id: int, ) → megatron.core.inference.inference_request.DynamicInferenceRequest#

Get most recent request from a request record.

Parameters:: request_id (int) – Request id.
Returns:: (DynamicInferenceRequest) The most recent request in the record.

_validate_async_sched_support_for_config() → None#

Validate config-level restrictions for async scheduling.

Raises if the config does not support async scheduling.

_resolve_generation_budget( request: megatron.core.inference.inference_request.DynamicInferenceRequest, ) → bool#: Resolve and clamp the generation budget; False when the request overflows the engine.

_add_request( request: megatron.core.inference.inference_request.DynamicInferenceRequest, *, is_resume: bool = False, record: Optional[megatron.core.inference.inference_request.DynamicInferenceRequestRecord] = None, ) → asyncio.Future[megatron.core.inference.inference_request.DynamicInferenceRequest]#

Add a request to the engine.

Parameters:

request (DynamicInferenceRequest) – Request to add.
is_resume (bool) – Whether the newest segment of a record is being re-admitted, as after suspend/resume, rather than a new record opened.
record (Optional[DynamicInferenceRequestRecord]) – With is_resume, the record of a generation resumed from another engine.

Returns:

Future completed when the request finishes.

Return type:

asyncio.Future[DynamicInferenceRequest]

add_request( request_id: int, prompt: Union[str, List[int], torch.Tensor], sampling_params: Optional[megatron.core.inference.sampling_params.SamplingParams] = None, precomputed_block_hashes: Optional[List[int]] = None, *, imgs: Optional[torch.Tensor] = None, num_tiles: Optional[torch.Tensor] = None, num_img_embeddings_per_tile: int = 0, imgs_sizes: Optional[torch.Tensor] = None, num_frames: Optional[torch.Tensor] = None, video_frame_indices: Optional[List[List[int]]] = None, video_fps: Optional[List[float]] = None, media_tokens_preexpanded: bool = False, offload_params: Optional[Dict] = None, media_cache_key: Optional[str] = None, ) → asyncio.Future[megatron.core.inference.inference_request.DynamicInferenceRequest]#

Add request to inference context.

Supports both text-only and multimodal requests. For text-only, call with just (request_id, prompt, sampling_params). For multimodal, also pass imgs and either (num_tiles + num_img_embeddings_per_tile) for static resolution or imgs_sizes for dynamic resolution.

When multimodal kwargs are provided the method will:

Expand compact media tokens, or derive a mask from pre-expanded tokens.
Run the vision encoder to produce image embeddings.
Store the embeddings and mask in the context for later use by the controller’s forward step.

Parameters:

request_id (int) – Unique ID of request.
prompt (Union[str, Tensor]) – Prompt as either a text string or token IDs.
sampling_params (Optional[SamplingParams]) – Sampling parameters for the request.
precomputed_block_hashes (Optional[List[int]]) – Prefix-cache hashes already computed for the prompt’s complete blocks. Values must match compute_block_hashes_batched(prompt_tokens, block_size_tokens).
imgs (Optional[Tensor]) – Image tensor [num_tiles, C, H, W] or [1, total_patches, patch_features] (or None).
num_tiles (Optional[Tensor]) – Number of tiles per image (1-D tensor, or None). Static resolution.
num_img_embeddings_per_tile (int) – Number of image embeddings per tile. Static resolution.
imgs_sizes (Optional[Tensor]) – Per-image sizes [N, 2] with [H, W]. Dynamic resolution.
num_frames (Optional[Tensor]) – Number of frames per image/video item.
video_frame_indices (Optional[List[List[int]]]) – Source indices for each sampled video frame.
video_fps (Optional[List[float]]) – Source FPS for each video item.
media_tokens_preexpanded (bool) – Whether prompt token IDs already contain one model token per projected media embedding.
offload_params (Optional[Dict]) – Opaque metadata forwarded to the payload stager.
media_cache_key (Optional[str]) – Media identity computed by the submitting inference client. Direct callers may omit it and let the engine derive an identity from the resolved media tensors.

Returns:

Returns an asyncio Future[DynamicInferenceRequest] for the user to wait on.

_admit_request( request: megatron.core.inference.inference_request.DynamicInferenceRequest, ) → asyncio.Future[megatron.core.inference.inference_request.DynamicInferenceRequest]#: Admit a built request, resuming the generation its preparer handed back, if any.

_build_vlm_request( *, request_id: int, prompt_str: Optional[str], tokens: torch.Tensor, sampling_params: Optional[megatron.core.inference.sampling_params.SamplingParams], imgs: Optional[torch.Tensor], num_tiles: Optional[torch.Tensor], num_img_embeddings_per_tile: int, imgs_sizes: Optional[torch.Tensor], precomputed_block_hashes: Optional[List[int]] = None, num_frames: Optional[torch.Tensor] = None, video_frame_indices: Optional[List[List[int]]] = None, video_fps: Optional[List[float]] = None, media_tokens_preexpanded: bool = False, offload_params: Optional[Dict] = None, media_cache_key: Optional[str] = None, ) → megatron.core.inference.inference_request.DynamicVLMInferenceRequest#: Prepare media tokens, run the vision encoder, register per-request media data on the context, and return a DynamicVLMInferenceRequest.

_place_cached_prompt_log_probs( request: megatron.core.inference.inference_request.DynamicInferenceRequest, ) → None#

Append the cached prompt log-probs that continue a request’s prompt scores.

A chunk that skips cached tokens is given their rows at admission. Under async scheduling that precedes post-processing of the previous chunk, so each cached run is placed only once every score before it has arrived.

Parameters:: request (DynamicInferenceRequest) – Request whose prompt scores to extend.

_promote_recomputed_prompt_scores( entry: core.inference.engines.dynamic_engine.RequestEntry, ) → None#

Restore the original prompt scores from the first complete recomputation.

A request checkpointed before its prompt scores arrive recomputes them into a later record segment, but completion and streaming read them from the first.

Parameters:: entry (RequestEntry) – Entry whose first record segment to repair.

static _find_step_stop_match( generated_tokens: List[int], step_tokens: List[int], stop_word_ids: Optional[List[List[int]]], ) → Optional[Tuple[int, int]]#

Find the earliest stop-sequence match ending in the current burst.

Parameters:

generated_tokens – Tokens retained before the current burst.
step_tokens – Untrimmed tokens produced by the current step.
stop_word_ids – Configured tokenized stop sequences.

Returns:

The zero-based position in step_tokens of the stop sequence’s final token and the matched sequence length, or None when no stop sequence ends in this burst.

static _get_usable_speculative_proposal_count( num_speculative_tokens: int, accepted_token_count: int, terminal_token_position: Optional[int], ) → int#

Count proposal positions usable before a semantic terminal boundary.

A terminal accepted draft makes only proposals through that draft usable. A replacement-token boundary still required verification of the entire speculative group, as did a step with no terminal boundary.

Parameters:

num_speculative_tokens – Configured proposals per decode request.
accepted_token_count – Accepted draft prefix length for this request.
terminal_token_position – Earliest logical output-burst boundary, -1 for a boundary before this burst, or None for no boundary.

Returns:

Number of leading proposal positions to include in acceptance metrics.

_process_request_outputs( request_ids: torch.Tensor, finished_request_ids: torch.Tensor, sample: torch.Tensor, accepted_tokens: torch.Tensor, log_probs: torch.Tensor, consumed_chunked_prefill_request_id: int, top_n_logprobs: Optional[Dict[int, List[Tuple[torch.Tensor, torch.Tensor]]]] = None, pre_fwd_active_token_count: Optional[int] = None, pre_fwd_step_count: Optional[int] = None, finished_routing_block_ids: Optional[Dict[int, list[int]]] = None, finished_handoff_block_ids: Optional[Dict[int, list[int]]] = None, finished_handoff_ssm_slots: Optional[Dict[int, int]] = None, finished_handoff_decode_tokens: Optional[Dict[int, list[int]]] = None, process_request_ids: Optional[set[int]] = None, record_step_latency: bool = False, ) → Tuple[List[int], List[megatron.core.inference.inference_request.DynamicInferenceRequest], Dict[int, Tuple[int, int]]]#

Apply sampled output and assemble results while request storage is valid.

Parameters:

request_ids (torch.Tensor) – A list of request_ids
finished_request_ids (torch.Tensor) – A list of finished request ids
sample – Tensor: The newly generated token for each request
accepted_tokens – Tensor: The additional accepted tokens for each request
log_probs – (List): Log probs for each request
consumed_chunked_prefill_request_id (int) – Chunked-prefill request ID associated with the consumed forward, or -1 if it had no partial chunk.
top_n_logprobs – (Dict): Top-n log probs for each request. Maps request_idx to list of (top_n_logprobs, top_n_indices) tuples.
pre_fwd_active_token_count (Optional[int]) – Active token count for the consumed forward.
pre_fwd_step_count (Optional[int]) – Step count for the consumed forward.
finished_routing_block_ids – (Dict[int, List[int]]): Block IDs for finished requests, saved before update_requests released them. Used for per-block routing reconstruction.
finished_handoff_block_ids – Prompt KV block IDs retained for state handoff.
finished_handoff_ssm_slots – Live SSM slots detached for state handoff.
finished_handoff_decode_tokens – First sampled token and optional MTP proposals needed to resume directly from imported prefill state on decode.
process_request_ids – Requests to process in this phase, or None for the full batch.
record_step_latency – Reserve this step’s latency entries before stop-word trimming.

Returns:

Active request IDs, completed requests, and latency-entry starting offsets with pre-trim token counts for the later timing update.

_finalize_finished_requests(

completed_requests: Dict[int, Tuple[megatron.core.inference.inference_request.DynamicInferenceRequest, Optional[Tuple[int, int]]]],

context_state: Dict,

record_step_latency: bool,

**output,

) → None#

Finalize only completed rows before their context storage is released.

Parameters:

completed_requests – Step-owned results and pending latency entries.
context_state – Metadata describing the consumed forward.
record_step_latency – Whether this step will have a measured latency.
**output – Sampled tokens, scores, finished IDs, and retained handoff state.

post_process_requests( request_ids: torch.Tensor, finished_request_ids: torch.Tensor, evict_request_ids: torch.Tensor, step_time: float, sample: torch.Tensor, accepted_tokens: torch.Tensor, log_probs: torch.Tensor, consumed_chunked_prefill_request_id: int, top_n_logprobs: Optional[Dict[int, List[Tuple[torch.Tensor, torch.Tensor]]]] = None, pre_fwd_active_token_count: Optional[int] = None, pre_fwd_step_count: Optional[int] = None, finished_routing_block_ids: Optional[Dict[int, list[int]]] = None, finished_handoff_block_ids: Optional[Dict[int, list[int]]] = None, finished_handoff_ssm_slots: Optional[Dict[int, int]] = None, finished_handoff_decode_tokens: Optional[Dict[int, list[int]]] = None, completed_requests: Optional[Dict[int, Tuple[megatron.core.inference.inference_request.DynamicInferenceRequest, Optional[Tuple[int, int]]]]] = None, ) → Tuple[List[int], List[megatron.core.inference.inference_request.DynamicInferenceRequest]]#

Process survivors and publish completed results after a controller step.

Parameters:

request_ids – IDs in the consumed forward’s row order.
finished_request_ids – IDs completed by the controller.
evict_request_ids – IDs to checkpoint and requeue.
step_time – Measured step latency, or zero on non-logging steps.
sample – Sampled tokens in consumed row order.
accepted_tokens – Accepted MTP tokens, or None.
log_probs – Selected-token scores per request, or None.
consumed_chunked_prefill_request_id – Partial-chunk ID, or -1.
top_n_logprobs – Optional top-n scores by consumed row.
pre_fwd_active_token_count – Consumed forward’s token count.
pre_fwd_step_count – Consumed forward’s step count.
finished_routing_block_ids – Block lists for results not finalized early.
finished_handoff_block_ids – Retained handoff KV blocks.
finished_handoff_ssm_slots – Detached handoff SSM slots.
finished_handoff_decode_tokens – Tokens for resuming the handed-off request.
completed_requests – Results assembled by the pre-release callback and pending latency entries. Futures are published here, not in the callback.

Returns:

Active request IDs and completed requests.

_get_and_clear_stop_word_finished_ids( active_request_ids: list[int], ) → set[int]#

Get and clear the set of request IDs that should be finished due to stop words.

This callback is called from the controller during bookkeeping to get request IDs that were detected as hitting stop words in the previous step’s post_process_requests.

Parameters:: active_request_ids – List of currently active request IDs.
Returns:: Set of request IDs from active_request_ids that should be marked as finished.

_set_finish_reason( request: megatron.core.inference.inference_request.DynamicInferenceRequest, ) → None#: Record final termination semantics without re-inferring them in clients.

_terminating_token_ids( request: megatron.core.inference.inference_request.DynamicInferenceRequest, ) → frozenset#

Token ids that end generation for this request.

The CPU-side counterpart of the controller’s per-step tensor check, resolved through the controller, including explicit request stop ids. Disabling model EOS (ignore_eos) does not disable explicit stops.

Parameters:: request (DynamicInferenceRequest) – Request to resolve ids for.
Returns:: Model EOS ids (unless disabled) and explicit request stop ids.
Return type:: frozenset

_truncate_at_mid_block_eos( request: megatron.core.inference.inference_request.DynamicInferenceRequest, tokens: list[int], request_log_probs: Optional[list], top_n_logprobs: Optional[Dict[int, List[Tuple[torch.Tensor, torch.Tensor]]]], req_idx: int, ) → Tuple[list, Optional[list], bool]#

Drop everything a speculative step emitted after an EOS.

An EOS can land on an accepted speculative token rather than on the step’s last token, which is the only one the controller’s termination check sees (sampled_tokens_cpu is the target model’s own sample at the last accepted position). Accepted tokens are verified target tokens, not rolled-back drafts, so an EOS among them is real and must end the request – otherwise the rest of the block, plus every later step, is emitted after EOS. This mirrors the stop-word check, which truncates a stop sequence that ends mid-block.

Trailing log probs / top-n are trimmed to match, keeping them aligned with tokens. top_n_logprobs is trimmed in place; log probs are returned.

Parameters:

request (DynamicInferenceRequest) – Request the tokens belong to.
tokens (list[int]) – Tokens emitted by this step, in order.
request_log_probs (Optional[list]) – This step’s log probs, or None.
top_n_logprobs (Optional[Dict]) – Per-request top-n log probs, or None.
req_idx (int) – This request’s index into top_n_logprobs.

Returns:

Truncated tokens, truncated log probs, and whether an EOS was found (the caller defers the finish by a step).

Return type:

Tuple[list, Optional[list], bool]

_find_mid_block_eos( request: megatron.core.inference.inference_request.DynamicInferenceRequest, tokens: list[int], ) → Optional[int]#

Locate an explicit stop token or a speculative-block model EOS.

Model EOS is scanned when a step emits more than one token. Explicit request stop tokens are also checked for single-token steps, since the controller does not hold their variable-length sets. Its EOS check inspects just the step’s last token, so an EOS on an accepted draft position is otherwise missed and generation runs past it. A single-token step was already checked by the controller, and a request the controller has already finished does not reach here as unfinished.

Parameters:

request (DynamicInferenceRequest) – Request the tokens belong to.
tokens (list[int]) – Tokens emitted by this step, in order.

Returns:

Index of the first EOS in tokens, or None if there is none.

Return type:

Optional[int]

_check_stop_words_for_request_post_append( request: megatron.core.inference.inference_request.DynamicInferenceRequest, *, record: Optional[megatron.core.inference.inference_request.DynamicInferenceRequestRecord] = None, num_new_tokens: Optional[int] = None, ) → Tuple[bool, int, int]#

Check if a request should stop due to stop words (after token is appended).

This method is called from post_process_requests after the token has already been appended to request.generated_tokens. In the speculative decoding case, multiple tokens may have been appended at once. If a stop word is found in the middle of the speculative tokens, the trailing tokens after the stop word are truncated from generated_tokens.

With speculative decoding, multiple tokens are appended at once. The stop word may end before the last appended token, leaving extra tokens that must be trimmed. When this happens, generated_tokens is truncated in-place and the number of trimmed tokens is returned so the caller can also trim log probs.

Parameters:

request – The request to check.
record – Full checkpoint history for the request. Supplying the record lets stop sequences span checkpoint boundaries.
num_new_tokens – Number of tokens appended in the current step. The returned trim count is limited to these tokens so the caller can trim pending log-probability results without deleting prompt data.

Returns:

Tuple of (stop_word_hit, num_new_tokens_trimmed, num_recomputed_prompt_scores_trimmed): stop_word_hit: True if the generated sequence contains a stop word. num_new_tokens_trimmed: Number of current-step tokens removed from the end of generated_tokens. num_recomputed_prompt_scores_trimmed: Number of prompt-score entries corresponding to stripped tokens from older checkpoint segments.

get_prefix_coordination_metrics() → dict#

Return prefix caching coordination metrics.

Returns:: Dict with coordination stats including the number of scheduling waits.

_mamba_batch_invariant_prefill_chunk_length( remaining: int, capacity: int, ) → int#

Number of prompt tokens to compute this step, at most capacity.

remaining counts the prompt tokens still to compute (excluding any cached prefix). Non-final calls must start and end at SSM chunk boundaries. The final prompt call may be shorter because it seeds the decode replay tail.

schedule_waiting_requests() → None#: Try to schedule requests from the waiting pool.

_can_schedule_non_chunked_prefill( req, *, record_cg_wait: bool, ) → bool#

Return whether the queue-head request can be admitted now.

Parameters:

req – Queue-head inference request.
record_cg_wait (bool) – Whether a CUDA-graph miss should update the request’s wait counter.

Returns:

Whether all request, token, KV-cache, and CUDA-graph checks pass.

Return type:

bool

_can_schedule_chunked_prefill(req) → bool#

Return whether the queue-head request can admit at least one prompt token.

Parameters:: req – Queue-head inference request.
Returns:: Whether request, token, and KV-cache capacity permit a chunk.
Return type:: bool

_should_run_async_sched_overlap() → bool#

Return whether this step should use overlap ordering.

Returns:: Whether the next step can use overlap ordering.
Return type:: bool

_pending_prefill_block_hashes() → set[int]#: Return hashes reserved by work whose KV is not yet ready for reuse.

schedule_non_chunked_prefill() → None#: Schedule non-chunked prefill requests.

_cg_admission_gating_active() → bool#

Cudagraph-aware admission gating is active when –inference-cuda-graph-all-prefills is set, the engine has prefill/mixed CGs, and the batch-dim list is populated.

All are required so legacy tests that exercise the scheduler without intending to run on captured graphs are unaffected. Gating is opt-in via cuda_graph_all_prefills.

_find_cg_chunk_size(max_chunk_tokens: int) → Optional[int]#

Return the largest chunk size <= max_chunk_tokens where batch matches a captured graph, or None if no graph covers any chunk in the budget.

Walks the captured-CG list (sorted descending by token_count) and returns the largest chunk that falls within budget and produces an applicable batch_dim under the engine’s matching mode (strict for hybrid models). Callers must explicitly handle the None case rather than scheduling eagerly.

_select_cg_chunk_size( req: megatron.core.inference.inference_request.DynamicInferenceRequest, max_chunk_tokens: int, ) → Optional[int]#

Select a graphed chunk size, or wait for the resident batch to drain.

A deferred request remains at the queue head, so the resident batch can drain and change the graph shape. A lone in-flight continuation cannot defer, so it falls back to eager execution when no graph covers it.

_register_cg_wait(req) → None#

Track a deferred admission attempt and throw a starvation warning at the threshold.

Decode is bounded by the number of decode steps. Persistent waits past _cg_admission_warn_after consecutive steps signal a problem.

_cg_admission_check( req, candidate: megatron.core.inference.batch_dimensions_utils.InferenceBatchDimensions, ) → bool#

Return True if the candidate batch shape matches a captured cudagraph.

On miss, registers a wait + warning via _register_cg_wait. On hit, resets the counter. Caller is responsible for breaking the scheduler loop on False. Passes match_ep_token_counts=False so this local admission probe doesn’t force a per-attempt NCCL all-reduce — the step-time matcher does its own EP sync.

Parameters:

req – Request whose CUDA-graph wait state should be updated.
candidate (InferenceBatchDimensions) – Candidate batch after admission.

Returns:

Whether a compatible captured graph exists.

Return type:

bool

_matches_cg_admission( candidate: megatron.core.inference.batch_dimensions_utils.InferenceBatchDimensions, ) → bool#

Return whether a candidate batch matches a captured CUDA graph.

Parameters:: candidate (InferenceBatchDimensions) – Candidate batch after admission.
Returns:: Whether a compatible captured graph exists.
Return type:: bool

schedule_chunked_prefill()#: This function schedules chunked prefill requests. Invariant: - There are at most one chunked prefill request in the waiting pool, which should be the head - There are at most one chunked prefill request in the context, which should be the last active request - context.chunked_prefill_request_id == -1 if no chunked prefill request is scheduled, otherwise it is the request id of the chunked prefill request - For each request, finished_chunk_token_count is the number of tokens that have been prefilled for this request, non-zero means it is during a chunked prefill - For each request, remaining_prompt_tokens holds the unprefilled prompt tokens

async async_forward() → Tuple[Optional[Dict], Dict, float]#

Uses asyncio for continuous generation. Sleeps when no requests are available, until new requests have been added.

Returns:: step_result (Optional[Dict]): The result of the step. context_state (Dict): Decode-only state, total/paused request count, and active token count. step_time (float): How long this step took.
Return type:: A tuple comprised of

_try_send_streaming_partials() → None#

Send pending token deltas to the inference coordinator.

Stream offsets (_partial_emit_lengths) are indexed in the coordinates of record.merge(), i.e. across all checkpoint segments of a request. When a request is evicted or recomputed after a suspend, DynamicInferenceRequestRecord.checkpoint() appends a new segment whose prompt includes the earlier generations and whose generated_tokens restart at 0, so the latest segment alone cannot be indexed by the stream offset.

async async_bookkeep( step_result: Optional[Dict], context_state: Dict, step_time: float, ) → core.inference.engines.dynamic_engine.DynamicInferenceEngineStepResult#

Uses asyncio for continuous bookkeeping.

Parameters:

step_result (Optional[Dict]) – The result of the step.
context_state (Dict) – Decode-only state, total/paused request count, and active token count.
step_time (float) – How long this step took.

Returns:

active_request_ids (List): IDs that ran in the last step and remain active. finished_requests (List): Flat, text-unfinalized requests that finished. step_time (float): The step time in seconds. cuda_graph_request_count (int): The CUDA graph batch size matching this step.

Return type:

A dictionary containing

async async_step() → core.inference.engines.dynamic_engine.DynamicInferenceEngineStepResult#

Wrapper for controller.generate_output_tokens_dynamic_batch(), to match vLLM API. Uses asyncio for continuous generation which allows this method to sleep and wake up when new requests are available.

Returns:: Active request IDs, finished requests, and step metadata.

_run_coroutine_sync(coro)#

Run a coroutine synchronously, handling the case when already in an event loop.

This method safely runs an async coroutine from synchronous code, even when called from within an already running event loop (e.g., when used with async frameworks like pytriton).

step_modern() → core.inference.engines.dynamic_engine.DynamicInferenceEngineStepResult#: Synchronous wrapper for self.async_step.

step_legacy( sampling_params: megatron.core.inference.sampling_params.SamplingParams, ) → Tuple[List[megatron.core.inference.inference_request.DynamicInferenceRequest], List[megatron.core.inference.inference_request.DynamicInferenceRequest], float]#: Synchronous wrapper for self.async_step.

step#: None

generate( prompts: List[str], sampling_params: Optional[megatron.core.inference.sampling_params.SamplingParams] = SamplingParams(), ) → List[megatron.core.inference.inference_request.DynamicInferenceRequest]#: Generate token-complete, text-unfinalized requests for a prompt batch.

static _pack_tp_broadcast( messages: List[List[bytes]], ) → List[bytes]#

Flatten per-message frame lists into one TP-broadcast multipart message.

A message is a list of frames – metadata first, then any payload bodies. ZMQ multipart is flat, so the frame boundaries would be lost on the wire. They are carried instead in a manifest frame holding one frame count per message, which lets peer ranks rebuild the grouping without any payload being copied, decoded, or re-packed.

Parameters:: messages – One frame list per message, in delivery order.
Returns:: [tp_broadcast_header, manifest, *flattened frames].

static _unpack_tp_broadcast( frames: List[bytes], ) → List[List[bytes]]#

Rebuild per-message frame lists from a TP broadcast.

Inverse of :meth:_pack_tp_broadcast.

Parameters:: frames – The received multipart message, header frame first.
Returns:: One frame list per message, in the order they were packed.

schedule_requests() → int#

Drains the ZMQ socket for a batch of requests and adds them to the engine.

This method is a collective and synchronous operation that must be called by all ranks in a Model Parallel (MP) group at the same time. It ensures that all ranks process the exact same batch of incoming requests and control signals.

The synchronization works as follows:

The MP rank 0 drains all pending messages from its subscriber socket in a non-blocking manner.
MP rank 0 then broadcasts the number of messages it received to all other ranks in its MP group using a dedicated publisher socket.
The other MP ranks wait to receive this count, and then receive exactly that many messages from their subscriber sockets.

Once all ranks have the same batch of messages, they are unpacked and processed. New requests are added to the engine’s queue, and control signals (PAUSE, UNPAUSE, SUSPEND, RESUME, STOP) update the engine’s internal state.

.. note::

This function is synchronous and must be called collectively by all ranks in a MP group. It should not be launched in a separate coroutine to ensure all ranks execute it in lockstep before proceeding to the next engine step.

Returns:: The number of messages that were received and processed in this batch.
Return type:: int

_prepare_submit_request_message( message: List[bytes], ) → List[bytes]#

Resolve a prompt once on MP rank zero before broadcasting the request.

The preparer only has work when the client sent offload params, so a request whose offload frame is None passes through untouched without any frame being decoded. When it runs, only the prompt frame and the offload frame are rewritten; the metadata frame is never repacked.

async shutdown()#

Shut down the engine and clean up ZMQ resources.

Called from the engine loop’s finally block after the loop exits.

async run_engine(*, loop: Optional[asyncio.AbstractEventLoop] = None)#: Continually steps the engine asynchronously.

async _ep_establish_consensus( local_work: int, signal_consensus: bool, ) → tuple[int, bool]#

EP all-reduce to share work counts and pause consensus.

All-reduces two integers at once:

local_work: actual pending request count (always >= 0).
consensus flag: -1 if this rank wants to pause, 0 otherwise.

Using max for both:

max(work) > 0 means at least one EP peer has real work.
max(consensus) == -1 means ALL peers signaled -1 (all PAUSING). Any RUNNING peer contributes 0, pulling the max to 0.

Parameters:

local_work – Pending request count for this rank.
signal_consensus – True if this rank is ready to pause.

Returns:

max work across EP, and whether all peers signaled consensus.

Return type:

(global_work, all_pausing)

async _world_barrier()#

World-wide ZMQ all-reduce barrier for global rank consensus.

Used for all state transitions that require global synchronization: PAUSING → PAUSED, UNPAUSING → RUNNING, SUSPENDING → SUSPENDED, RESUMING → PAUSED, and STOPPING → STOPPED.

No-op when world_size == 1 (communicator is not created).

async run_engine_with_coordinator( *, loop: Optional[asyncio.AbstractEventLoop] = None, )#

Continually steps the engine asynchronously.

State-dependent behavior:

RUNNING: EP all-reduce to check for work, then step or idle.
PAUSING: EP all-reduce to reach consensus, then world barrier.
PAUSED / SUSPENDED: Idle-sleep, wait for signals via schedule_requests().
UNPAUSING / SUSPENDING / RESUMING / STOPPING: World barrier, then transition.
STOPPED: Teardown and exit.

core.inference.engines.dynamic_engine#

Module Contents#

Classes#

Functions#

Data#

API#

`core.inference.engines.dynamic_engine`#