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

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

Data#

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

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.

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

_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

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

_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]]#: Hook overridden by the KV-handoff engine composition.

_record_handoff_completion_notification( request_id: int, failed: 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]], ) → dict#: Hook overridden by the KV-handoff engine composition.

_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, backend: str = 'nixl') → None#: Raising stub; the hand-off engine composition overrides it.

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

_poll_pending_kv_pushes() → int#

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

add_request_with_kv_handoff( request_id, prompt, sampling_params, kv_meta, src_block_ids, ) → 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.

static _raise_kv_handoff_not_enabled(operation: str) → None#

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

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.

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

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

static _complete_request( request_entry: core.inference.engines.dynamic_engine.RequestEntry, ) → megatron.core.inference.inference_request.DynamicInferenceRequest#: Merge an engine-owned record once and resolve its completion future.

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

_serialize_finished_request( request: megatron.core.inference.inference_request.DynamicInferenceRequest, finished_metadata: Optional[megatron.core.inference.inference_request.FinishedRequestRecord], ) → Dict#: Stage a non-streaming accepted payload before constructing its coordinator reply.

_handle_failed_request(request_id: int)#

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.

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

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.

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

Add a request to the engine.

Parameters:

request (DynamicInferenceRequest) – Request to add.
is_resume (bool) – Whether an existing record is being explicitly re-admitted after suspend/resume.

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

_apply_prompt_preparation_error(request, offload_params) → None#: Fail a request whose prompt preparer reported an error on MP rank zero.

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

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, ) → Tuple[List[int], List[megatron.core.inference.inference_request.DynamicInferenceRequest]]#

Handles post-processing for requests after a step.

Parameters:

request_ids (torch.Tensor) – A list of request_ids
finished_request_ids (torch.Tensor) – A list of finished request ids
evict_request_ids (torch.Tensor) – A list of evicted request ids.
step_time (float) – The latency of the last step
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.

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.

_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( req: megatron.core.inference.inference_request.DynamicInferenceRequest, capacity: int, ) → int#

Raw prefill length that computes an aligned chunk within capacity.

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

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 first 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 by deferring the admission rather than scheduling eagerly.

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

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