Is this page helpful?

`core.inference.contexts.dynamic_context`#

Module Contents#

Classes#

`ContextErrorFactory`	Factory class for serializing/deserializing context errors.
`DynamoHelper`	Manage KV-cache lifecycle events consumed by the Dynamo frontend.
`DynamicInferenceContext`	Inference context that is passed to the main model in order to efficiently calculate and store the KV cache during inference.

Functions#

get_mem_size_str

Convert number of bytes to human-readable string.

Data#

`KVEventListener`
`DEPRECATED_ARGS`

API#

core.inference.contexts.dynamic_context.KVEventListener#: None

core.inference.contexts.dynamic_context.DEPRECATED_ARGS#: [‘params_dtype’, ‘num_layers’, ‘kv_channels’, ‘num_attention_heads’, ‘max_sequence_length’, ‘buffer_…

exception core.inference.contexts.dynamic_context.ContextOverflowError( request_id: Optional[int], message: Optional[str] = None, *, is_transient: bool = True, )#

Bases: Exception

Base exception for when a new request does not fit.

Parameters:: is_transient (bool) – Flag marking whether error is transient (i.e., may work if we try again, but fails due to the current context state), or permanent (i.e., request will never fit in this context).

Initialization

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

exception core.inference.contexts.dynamic_context.RequestOverflowError( request_id: Optional[int], message: Optional[str] = None, *, is_transient: bool = True, )#

Bases: core.inference.contexts.dynamic_context.ContextOverflowError

Adding request would overflow max request count.

Initialization

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

exception core.inference.contexts.dynamic_context.TokenOverflowError( request_id: Optional[int], message: Optional[str] = None, *, is_transient: bool = True, )#

Bases: core.inference.contexts.dynamic_context.ContextOverflowError

Adding request would overflow max token count.

Initialization

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

exception core.inference.contexts.dynamic_context.MaxSequenceLengthOverflowError( request_id, message: Optional[str] = None, )#

Bases: core.inference.contexts.dynamic_context.ContextOverflowError

Adding request would overflow max sequence length.

Initialization

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

exception core.inference.contexts.dynamic_context.BlockOverflowError( request_id: Optional[int], message: Optional[str] = None, *, is_transient: bool = True, )#

Bases: core.inference.contexts.dynamic_context.ContextOverflowError

Adding request would overflow available memory blocks.

Initialization

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

exception core.inference.contexts.dynamic_context.ActiveRequestCountOverflowError( max_request_count, active_request_count, )#

Bases: core.inference.contexts.dynamic_context.ContextOverflowError

Used when initialize_attention_state() is called with `num_warmup_requests > max_requests.

Initialization

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

exception core.inference.contexts.dynamic_context.TensorStateDeallocatedError( request_id: Optional[int], message: Optional[str] = None, *, is_transient: bool = True, )#

Bases: core.inference.contexts.dynamic_context.ContextOverflowError

Context’s tensor state is currently deallocated, such as when the engine has been suspended.

Initialization

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

class core.inference.contexts.dynamic_context.ContextErrorFactory#

Factory class for serializing/deserializing context errors.

classmethod serialize( error: core.inference.contexts.dynamic_context.ContextOverflowError, ) → dict#

Serialize error.

Parameters:: error (ContextOverflowError) – Error.
Returns:: (dict) Serialized error data.

classmethod deserialize( obj: dict, ) → core.inference.contexts.dynamic_context.ContextOverflowError#

Deserialize error.

Parameters:: obj (dict) – Serialized error data.
Returns:: (ContextOverflowError) Deserialized error.

core.inference.contexts.dynamic_context.get_mem_size_str(n_bytes: int) → str#: Convert number of bytes to human-readable string.

class core.inference.contexts.dynamic_context.DynamoHelper#

Manage KV-cache lifecycle events consumed by the Dynamo frontend.

Initialization

property has_kv_event_listeners: bool#: Return whether any KV-event listeners are registered.

add_kv_event_listener( listener: core.inference.contexts.dynamic_context.KVEventListener, ) → None#

Register a KV-cache lifecycle listener.

Parameters:: listener – Callback invoked with the event kind and payload.

queue_kv_stored_event(payload: dict[str, Any]) → None#

Queue a stored event for publication after a successful forward pass.

Parameters:: payload – Stored-event payload.

publish_pending_kv_stored_events() → None#: Publish blocks whose KV contents were produced by a successful forward pass.

discard_pending_kv_stored_events() → None#: Discard registrations left by an interrupted or failed forward pass.

on_kv_blocks_deregistered( _block_ids: list[int], hashes: set[int], ) → None#

Publish removal events for deregistered KV blocks.

Parameters:

_block_ids – Deregistered block IDs, unused by Dynamo.
hashes – Hashes of the deregistered blocks.

notify_kv_cache_cleared() → None#: Notify listeners that no previously advertised block is routable.

_emit_kv_event(kind: str, payload: dict[str, Any]) → None#: Notify Dynamo listeners without allowing frontend failures to stop inference.

class core.inference.contexts.dynamic_context.DynamicInferenceContext( model_config: megatron.core.transformer.TransformerConfig, inference_config: megatron.core.inference.config.InferenceConfig, )#

Bases: core.inference.contexts.base_context.BaseInferenceContext

Inference context that is passed to the main model in order to efficiently calculate and store the KV cache during inference.

The dynamic inference context manages both: 1) in-flight batching, and 2) a memory buffer for the block-level KV cache. For in-flight batching, requests of arbitrary sequence length may be added, paused, or removed from the context at any step. The only constraint is the maximum number of requests or tokens that the context is defined to support. For the block-level KV cache, a memory buffer is allocated up front (size buffer_size_gb if unified_memory_level == 0, or buffer_size_gb + paused_buffer_size_gb if unified_memory_level == 1), that is divided into blocks and dynamically assigned to requests. At any given step, any unassigned blocks equate to unused space.

Parameters:

model_config (TransformerConfig) – Model config.
inference_config (InferenceConfig) – Inference config.

Initialization

Args:

DEFAULT_MAX_TOKENS#: 16384

TOKEN_ROUNDER#: 64

REQUEST_ROUNDER#: 4

TMS_TAG#: ‘inference_context’

_allocate_memory_buffer()#: Allocate the KV cache memory buffer.

_allocate_mamba_states()#: Allocate Mamba states for hybrid models.

initialize_all_tensors() → None#: Allocate all GPU state during initial construction.

reinitialize_inference_state_buffers()#

Restore large tensors (KV cache, Mamba states) after a suspend.

Called by the engine during resume(). Initial allocation is in initialize_all_tensors().

deallocate_inference_state_buffers()#

Deallocate large tensors (KV cache, Mamba states) during suspend.

Called by the engine during suspend(). Mirror to reinitialize_inference_state_buffers().

classmethod round_up_tokens(value, tp_size=None)#: Round up to nearest multiple of TOKEN_ROUNDER that is also divisible by tensor model parallel size.

classmethod round_up_requests(value, tp_size=None)#: Round up to nearest multiple of REQUEST_ROUNDER that is also divisible by tensor model parallel size.

is_static_batching() → bool#: Is static batching? False.

is_decode_only() → bool#

Return if this iteration we run decode only implementation.

When CUDA graphs are active, uses padded_batch_dimensions because it reflects the post-expert-parallel sync state. Otherwise falls back to num_prefill_requests which is always up-to-date regardless of where we are in the step lifecycle.

using_cuda_graph_this_step() → bool#: Returns True if cuda graphs are being used for this step.

has_unfinished_requests() → bool#: Test if any requests remain.

cu_query_lengths() → Tuple[torch.Tensor, int]#: Cumulative query sequence lengths.

cu_kv_lengths() → Tuple[torch.Tensor, torch.Tensor, int]#: Cumulative key/value sequence lengths.

get_active_sequence_lengths() → torch.Tensor#: Total sequence length (query + key) for active requests.

get_max_sequence_lengths() → torch.Tensor#: Maximum sequence length for active requests.

get_active_request_count()#: Returns the current number of active requests.

build_active_slices(batch_size: int)#

Build the active slices of specific tensors. This is run on every forward step.

If the context is reordered to active -> paused -> finished, this can be graphed.

pad_active_slices()#: Pad the active slices of specific tensors.

append_key_value_cache( layer_number: int, key: torch.Tensor, value: torch.Tensor, ) → None#

Append to KV cache.

Parameters:

layer_number (int) – Layer number.
key (Tensor) – Key tensor.
value (Tensor) – Value tensor.

key_value_cache( layer_number: int, ) → Tuple[torch.Tensor, Optional[torch.Tensor], torch.Tensor]#

Read from KV cache.

Parameters:: layer_number (int) – Layer number.
Returns:: (Tuple[Tensor, Tensor, Tensor]) The key and value pointer tensors that point to blocks within the block-level memory buffer as well as the block table.

mamba_states_cache( layer_number: int, intermediate: bool = False, ) → Tuple[torch.Tensor, torch.Tensor]#: Returns the Mamba state tensors for the given layer.

_allocate_mamba_cache(mamba_gb: float) → None#

Allocate the Mamba state cache for prefix caching.

Parameters:: mamba_gb – GPU memory budget in GB for the cache.

apply_fused_qk_rotary_emb( query: torch.Tensor, key: torch.Tensor, cos_sin_emb: torch.Tensor, config: megatron.core.transformer.TransformerConfig, ) → Tuple[torch.Tensor, torch.Tensor]#

Apply rotary embedding to query and key tensors using flashinfer’s fused rope.

Parameters:

query (Tensor) – Query tensor.
key (Tensor) – Key tensor.
cos_sin_emb (Tensor) – Rotary embeddings.
config (TransformerConfig) – Transformer config.

Returns:

(Tuple[Tensor, Tensor]) Query and Key tensors after applying rotary embeddings.

apply_rotary_emb_query( query: torch.Tensor, query_emb: torch.Tensor, config: megatron.core.transformer.TransformerConfig, cu_seqlens_q: torch.Tensor, cp_group: torch.distributed.ProcessGroup, mscale: float = 1.0, ) → torch.Tensor#

Apply rotary embedding to query tensor.

Parameters:

query (Tensor) – Query tensor.
query_emb (Tensor) – Query rotary embeddings.
config (TransformerConfig) – Transformer config.
cu_seqlens_q (Tensor) – Cumulative sequence lengths.
cp_group (torch.distributed.ProcessGroup) – Process group for context parallel.

Returns:

(Tensor) Query tensor after applying rotary embeddings.

apply_rotary_emb_key( key: torch.Tensor, key_emb: torch.Tensor, config: megatron.core.transformer.TransformerConfig, cp_group: torch.distributed.ProcessGroup, mscale: float = 1.0, ) → torch.Tensor#

Apply rotary embedding to key tensor.

Parameters:

key (Tensor) – Key tensor.
key_emb (Tensor) – Key rotary embeddings.
config (TransformerConfig) – Transformer config.
cp_group (torch.distributed.ProcessGroup) – Process group for context parallel.

Returns:

(Tensor) Key tensor after applying rotary embeddings.

set_ep_zmq_communicator(communicator) → None#

Attach an EP-group ZMQ communicator for CPU-side sync collectives.

When set, match_graph_config() uses this communicator’s sync_all_reduce_max() to perform the EP batch-dimension MAX reduction on the CPU instead of launching a NCCL AllReduce kernel on the compute stream. Expected to be called once by the inference engine after both the context and the communicator have been created.

Parameters:: communicator – AsyncZMQCommunicator over the EP process group.

reset_attention_state() → None#: Reset state used within attention, after each step.

reset_mamba_state() → None#: Reset state used within Mamba layers.

add_dummy_requests_parallel( requests: Sequence[megatron.core.inference.inference_request.DynamicInferenceRequest], *, count_as_prefill: bool = True, ) → None#: Fast path to add dummy requests without allocating real KV blocks.

add_dummy_requests_for_cudagraph_capture( graph_dimensions: megatron.core.inference.batch_dimensions_utils.InferenceBatchDimensions, ) → None#: Adds dummy requests to reflect the number of prefill and decode requests in the graph config. These are using during cuda graph captures.

property num_decode_requests: int#: Returns the number of decode requests.

add_dummy_requests_for_expert_parallel_step( graph_dimensions: megatron.core.inference.batch_dimensions_utils.InferenceBatchDimensions, ) → None#

Minimal context setup so an EP rank with no real requests can replay an already-captured cuda graph without crashing or corrupting memory.

This is the fast alternative to add_dummy_requests_for_cudagraph_capture (which goes through the heavyweight add_dummy_requests_parallel path).

We setup minimal state such that initialize_attention_state and the forward pass can run without error.

Called AFTER the EP sync so graph_dimensions reflects the agreed-upon graph.

initialize_attention_state( *, construct_graph_dimensions: Optional[megatron.core.inference.batch_dimensions_utils.InferenceBatchDimensions] = None, is_expert_parallel_dummy_cuda_graph_step: bool = False, transfer_bookkeeping_to_gpu: bool = True, record_bookkeeping_done_event: bool = False, ) → Optional[torch.cuda.Event]#

Initialize attention state so that every layer can use it.

Parameters:

construct_graph_dimensions (Optional[InferenceBatchDimensions]) – The graph config to use for constructing the cuda graphs.
is_expert_parallel_dummy_cuda_graph_step (bool) – Whether this is a dummy expert model parallel step.
transfer_bookkeeping_to_gpu (bool) – Whether to publish the prepared CPU bookkeeping snapshot to GPU before returning. Legacy callers publish immediately; async scheduling binds the GPU views here and publishes their values later.
record_bookkeeping_done_event (bool) – Whether to record an event after the bookkeeping H2D transfer.

Returns:

Event marking bookkeeping H2D completion, or None when no event was requested or no transfer was performed.

Return type:

Optional[torch.cuda.Event]

_execute_pending_mamba_ops() → None#

Execute Mamba GPU operations deferred from add_request() / update_requests().

This runs at the start of initialize_attention_state() so that all GPU Mamba state is correct before the forward pass.

transfer_bookkeeping_to_gpu( skip_token_input_ids: bool = False, record_done_event: bool = False, ) → Optional[torch.cuda.Event]#

Batch transfer CPU bookkeeping state to GPU staging buffers.

Legacy steps call this from initialize_attention_state(). Async scheduling instead delays publication until after preparation and the GPU sample-to-input copy. Legacy transfers block because the pinned CPU source is re-staged in place. Async scheduling requests an event-tracked non-blocking copy and synchronizes that event before reusing the source.

The bookkeeping fields are backed by one contiguous pinned CPU buffer and one contiguous GPU buffer; a single memcpy covers the whole transfer. Request-level staging slots are refreshed from the persistent CPU tensors immediately before the H2D (GPU reads them at [:n_active] while CPU bookkeeping keeps them at [paused_count:total_count)).

Parameters:

skip_token_input_ids (bool) – If true, leave gpu_view.token_to_input_ids unchanged while copying the rest of the bookkeeping buffer.
record_done_event (bool) – Whether to record and return an event after an asynchronous bookkeeping transfer.

Returns:

Event marking H2D completion, or None when no event was requested.

Return type:

Optional[torch.cuda.Event]

copy_async_sched_sample_to_forward( sampled_tokens_cuda: torch.Tensor, sampled_mtp_tokens_cuda: Optional[torch.Tensor] = None, ) → None#

Populate GPU input token IDs from sampled CUDA tokens for async scheduling.

Async scheduling keeps sampled tokens GPU-resident for the next decode forward. CPU bookkeeping is prepared independently and published later; this direct GPU copy populates the live input-ID view without waiting for the sample’s CPU copy.

Parameters:

sampled_tokens_cuda (Tensor) – 1D CUDA tensor containing one sampled token per active decode request.
sampled_mtp_tokens_cuda (Optional[Tensor]) – MTP draft tokens with shape [num_speculative_tokens, active_request_count].

reset_tensors() → None#: Fill all bookkeeping tensors with sentinel values.

reset_metadata( preserve_prefix_cache: bool = False, *, preserve_counters: bool = False, ) → None#

Reset all bookkeeping state: counters, block allocator, attention/mamba state.

This must be called after initialize_all_tensors() and after any suspend/resume cycle to bring the context back to a clean state.

Parameters:

preserve_prefix_cache – When True, keep the KV block allocator’s prefix-cache state (hash index, ref counts, cached blocks) intact. Used by the idle dummy_forward path, which only needs to clear the transient one-token step state – wiping the allocator there would destroy cross-request prefix reuse for any subsequent request (the engine idles between requests at low concurrency, especially with EP > 1).
preserve_counters – When True, keep engine-step, prefix-cache clock, prefill-token, and async-scheduling counters intact.

reset( preserve_prefix_cache: bool = False, *, preserve_counters: bool = False, ) → None#

Reset entire context.

This method does:

Fill all GPU tensors with sentinel values.
Reset active/paused request/token counts to zero.
Reset available blocks to entire memory.

This method is useful after cuda graph warmup iterations, where the context’s memory buffer is referenced by the cuda graph system and cannot be deallocated.

Parameters:

preserve_prefix_cache – When True, keep the KV and Mamba prefix-cache state (hash indices and cached blocks/slots) intact. Used by the idle dummy_forward path so an idle step between requests does not destroy cross-request prefix reuse.
preserve_counters – When True, keep engine-step, prefix-cache clock, prefill-token, and async-scheduling counters intact.

current_input_and_position_ids( *, num_warmup_tokens: Optional[int] = None, ) → Tuple[torch.Tensor, torch.Tensor]#

Flattened input and position IDs for forward pass.

Parameters:: num_warmup_tokens (Optional[int]) – Number of tokens to return for warming up cuda graphs. Must be less than or equal to max_tokens.
Returns:: (Tuple[Tensor, Tensor]) Flattened active input and position IDs.

speculative_required_logit_indices() → torch.Tensor#

Token-level indices needed for speculative decode verification.

Returns all decode token positions (base + speculative) concatenated with the last token position of each prefill request.

Returns:: (Tensor) 1-D indices into the packed token sequence, length num_decode_requests * (num_speculative_tokens + 1) + num_prefill_requests in eager, or the equivalent padded count under non-eager.

property num_last_token_logits: int#

Number of rows produced by last_token_logits for the current step.

Single source of truth for the bound: one row per request, with (num_speculative_tokens + 1) rows per decode request when MTP is active.

last_token_logits(logits: torch.Tensor) → torch.Tensor#

Select the logit positions needed for token generation.

When speculative decoding is active, decode requests need logits for all their tokens (base + speculative) for verification, while prefill requests only need the last token logit. This avoids materializing the full vocab-sized logits for every prefill token, which causes large memory spikes during prefill-heavy batches.

Parameters:: logits (Tensor) – Output logits of forward pass, shape [1, S, H].
Returns:: (Tensor) Selected logits, shape [N, H], where N == num_last_token_logits.

_find_mamba_match_count( req: megatron.core.inference.inference_request.DynamicInferenceRequest, start_block: int, end_block: int, ) → int#

Find the farthest cached Mamba state within a chunk-local block range.

Mamba state restore is only valid for blocks that the current chunk also assigns from the KV cache. Chunked prefill can schedule a prompt prefix that is shorter than the farthest cached full-prompt Mamba boundary, so this helper intentionally uses the same block domain as KV matching.

_compute_prefix_match( req: megatron.core.inference.inference_request.DynamicInferenceRequest, prefill_chunk_length: int, record_mamba_match: bool = False, ) → Tuple[list, int, int, int, int, int]#

Compute prefix match results and skip counts for a request chunk.

Shared by check_availability (budget checks) and add_request (execution).

Parameters:

req – Request being scheduled.
prefill_chunk_length – Number of prompt tokens considered in this chunk.
record_mamba_match – If True, store the chunk-local executable Mamba match count on the request for diagnostics/tests.

Returns:

Tuple of (matched_block_ids, num_blocks_from_pool, already_allocated_blocks, overall_required_blocks, prefix_skip_tokens, effective_prefill_chunk_length).

check_availability( req: megatron.core.inference.inference_request.DynamicInferenceRequest, ) → Tuple[bool, bool, bool]#: Check if the request can be added to the context.

_find_kv_match_count( req: megatron.core.inference.inference_request.DynamicInferenceRequest, start_block: int, end_block: int, ) → tuple[list[int], int]#

Find cached blocks matching a range of the prompt using precomputed hashes.

Looks up hashes in req.precomputed_block_hashes[start_block:end_block] against the block allocator’s hash-to-block mapping. Stops at the first non-match.

Parameters:

req – The inference request with precomputed_block_hashes set.
start_block – First block index to match (inclusive).
end_block – Last block index to match (exclusive); clamped to hash count.

Returns:

List of matched block IDs (consecutive from start_block)
Parent hash of the last matched block (0 if no matches)

Return type:

Tuple of

add_request( req: megatron.core.inference.inference_request.DynamicInferenceRequest, prefill_chunk_length: Optional[int] = None, ) → None#

Add request to context. At this stage, we assume that the request is valid and can be added, as the checks are done in the schedule function.

Parameters:

req (DynamicInferenceRequest) – Request to add.
prefill_chunk_length (Optional[int]) – Length of prefill chunk to add. If None, the request will be fully added.

Returns:

None

_move_book_keeping_tensors( src_idxs, dst_idxs, next_tokens, new_speculative_tokens=None, )#: Move all the relevent booking tensors with src idxs to dst idxs

_swap_book_keeping_tensors( src_idxs, dst_idxs, next_tokens=None, new_speculative_tokens=None, )#: Swaps all the relevent booking tensors with src idxs to dst idxs

get_index_of_chunked_prefill_request(safe: bool = True) → int#

Get the index of the chunked prefill request in the context.

If safe is True, then clamp the search space to the current total request count. Otherwise, expand the search beyond the current total request count.

Returns:: (int) Index of the chunked prefill request, or -1 if none exists.

is_chunked_prefill_enabled() → bool#: Returns whether chunked prefill is enabled.

release_memory_blocks_from_request_indexes(request_indexes) → None#

Release memory blocks used by the given request idxs.

Parameters:: request_indexes (torch.Tensor) – Request indexes. (Note, NOT request ids.)

_get_paused_request_count_within_block_budget() → int#: Count the left-most paused requests whose blocks fit the paused budget.

_get_releasable_block_counts( request_start_idx: int, request_end_idx: int, ) → list[int]#

Count blocks made allocatable by each right-most request suffix.

The returned list has one entry for every possible suffix length, including zero. For prefix caching, a shared block is credited only once all selected request references account for its current allocator reference count.

resume_paused_requests( active_request_count: int, newly_paused_request_ids: Optional[torch.Tensor], ) → tuple[int, Optional[torch.Tensor]]#

Resume as many paused requests as compute and KV capacity permit.

Parameters:

active_request_count (int) – Number of active requests.
newly_paused_request_ids (Optional[torch.Tensor]) – Newly paused request ids.

Returns:

(tuple[int, Optional[torch.Tensor]]) Updated active count and newly paused ids.

evict_overflow_paused_requests( active_request_count: int, next_tokens: torch.Tensor, new_speculative_tokens: Optional[torch.Tensor] = None, ) → Optional[torch.Tensor]#

Evict requests that overflow the paused-block retention budget.

Parameters:

active_request_count (int) – Number of active requests.
next_tokens (torch.Tensor) – Sampled tokens.
new_speculative_tokens (Optional[torch.Tensor]) – Speculative tokens.

Returns:

(Optional[torch.Tensor]) Evicted request ids.

_get_async_sched_rows_requiring_new_block() → torch.Tensor#

Return active request rows that need a block during the next prepare.

Returns:: Boolean mask over active request rows.
Return type:: Tensor

can_prepare_requests() → bool#

Return whether requests can be prepared without lifecycle changes.

Returns:: Whether all requests are active decode requests and the shared KV-block pool can satisfy the exact next-step allocation demand.
Return type:: bool

prepare_requests() → None#

Speculatively prepare active decode requests for the next forward pass.

Async scheduling only supports decode-only steps with no pause, evict, or resume lifecycle changes. If preparation cannot allocate the required KV blocks without a lifecycle change, this method raises. The prepared decode layout establishes the active token count.

commit_sampled_tokens( sampled_tokens_cpu: torch.Tensor, sampled_mtp_tokens_cpu: Optional[torch.Tensor] = None, ) → None#

Commit sampled CPU token IDs to the prepared request state.

This establishes the post-resolution active token count and populates the CPU input-ID staging rows in survivor order. Overlapped async scheduling has already copied the same samples into the live GPU input view for the speculative forward.

Parameters:

sampled_tokens_cpu (Tensor) – Sampled CPU token for each active request.
sampled_mtp_tokens_cpu (Optional[Tensor]) – MTP draft tokens with shape [num_speculative_tokens, active_request_count].

resolve_requests( active_requests_mask: torch.Tensor, ) → Tuple[torch.Tensor, torch.Tensor]#

Resolve finished requests after an async scheduling forward pass.

Prefill requests transition to decode during resolution. Request rows use the same hole-filling order as update_requests so seeded sampling stays consistent with legacy scheduling. Token tensors and the active token count are left untouched; prepare rebuilds derived token metadata and the controller commits sampled input IDs after resolution.

Parameters:: active_requests_mask (Tensor) – 1D mask marking requests that remain active.
Returns:: Request IDs that finished and source row indices for surviving requests in their resolved destination order.
Return type:: Tuple[Tensor, Tensor]

update_requests( active_requests_mask: torch.Tensor, new_tokens: torch.Tensor, new_speculative_tokens: torch.Tensor = None, ) → torch.Tensor#

Update context state after calling engine.step().

This method is responsible for:

Update prefill requests to decode requests.
Persist decode requests as decode requests.
Terminate requests by length or termination id.

Note: All bookkeeping tensors (i.e., self.request_*) are laid out contiguously, with a conceptual division between paused requests on the ‘left’ (or, lower indices) and active requests in the ‘middle’ (or, middle indices) and completed requests on the ‘right’ (or, higher indices). The integers paused_request_count and total_request_count are used to track the boundaries between these request groups.

0:paused_request_count -> paused requests
paused_request_count:total_request_count -> active requests
total_request_count:max_requests -> completed requests are moved here. The reason for maintaining contiguous tensors rather than multiple smaller (e.g., per-group or per-request) tensors is for both 1) speed (avoid unnecessary tensor allocations), and 2) compatibility with the Flash Attention kernels, which packed contiguous tensors.

The following happens in this code :

The active token mask tells us which requests are still active and which are completed
If no paused requests are present and no active requests we release all memory and reset.
Concatenate the paused tokens to the active tokens
For the finished requests we release memory blocks and move them to the right
We identify requests that require a new block and add them to the paused requests (i.e move them left)
Resume paused requests & evict overflowing paused requests.
We make changes to the request book keeping tesnsors and setup the tokens for next iteration
We make relevant changes to the token bookkeeping tensors

Parameters:

active_requests_mask (Tensor) – 1D Mask tensor marking active requests. (Active request length)
new_tokens (Tensor) – Newly sampled tokens, with one token per active request. (Active request length)
new_speculative_tokens (Tensor) – Newly sampled speculative tokens, with num_speculative tokens per active request. (num_speculative_tokens, active_request_length)

Returns:

(Tensor) Newly paused request IDs.

_processed_log_probs( logits: torch.Tensor, n_active: int, active_query_lengths: Optional[torch.Tensor], sampling: Optional[megatron.core.inference.sampling.base.Sampling], row_to_request: Optional[torch.Tensor] = None, ) → torch.Tensor#

Calculate raw or sampling-processed per-row log probabilities.

Parameters:

logits (Tensor) – Raw logits with shape [num_rows, vocab_size].
n_active (int) – Number of active requests represented by the rows.
active_query_lengths (Optional[Tensor]) – CPU token counts used to map prefill rows to active requests, or None for decode.
sampling (Optional[Sampling]) – Backend providing processed logprobs.
row_to_request (Optional[Tensor]) – Explicit CPU mapping from each logit row to an active request.

Returns:

Per-row log probabilities over the vocabulary.

Return type:

Tensor

calculate_log_probs_tensors( logits: torch.Tensor, new_tokens: torch.Tensor, only_last_token_logits: Optional[bool] = False, sampling: Optional[megatron.core.inference.sampling.base.Sampling] = None, row_to_request: Optional[torch.Tensor] = None, ) → Tuple[torch.Tensor, torch.Tensor]#

Calculate selected-token and full-distribution log probabilities.

Parameters:

logits (Tensor) – Raw model output logits with shape [1, sequence_length, vocab_size].
new_tokens (Tensor) – Newly sampled tokens for active requests.
only_last_token_logits (Optional[bool]) – Whether logits contain only each request’s final token row.
sampling (Optional[Sampling]) – Sampling backend used for processed log probabilities.
row_to_request (Optional[Tensor]) – Explicit CPU mapping from each logit row to an active request.

Returns:

Selected-token log probabilities flattened in active-token order and the full per-row log-probability tensor.

Return type:

Tuple[Tensor, Tensor]

calculate_log_probs( logits: torch.Tensor, new_tokens: torch.Tensor, only_last_token_logits: Optional[bool] = False, sampling: Optional[megatron.core.inference.sampling.base.Sampling] = None, ) → Tuple[List[List[float]], torch.Tensor]#

Calculate log probs for all active requests and return them.

TODO: @wdykas support top-n log probs.

Parameters:

logits (Tensor) – Raw model output logits with shape [1, sequence_length, vocab_size].
new_tokens (Tensor) – The newly sampled tokens.
only_last_token_logits (bool) – If set, the logits are from only the last token in each request
sampling (Optional[Sampling]) – Backend used to optionally modify log-probs.

Returns:

List of lists where each inner list contains log probs for a request in the same order as the active requests (from paused_request_count to total_request_count). log_probs (Tensor): Used to compute top n logprobs later if required.

get_kvcache_utilization_stats() → dict#

Compute KV cache buffer utilization stats for the current step.

Returns a dictionary with counts and percentages for both allocated block usage (overall buffer occupancy) and active usage (blocks referenced by currently active requests this step).

Returns:: { ‘total_blocks’: int, ‘allocated_blocks’: int, ‘active_unique_blocks’: int, ‘allocated_utilization’: float, ‘active_utilization’: float, ‘active_request_count’: int, ‘paused_request_count’: int, ‘gtd_block_count’: int, }