`core.inference.text_generation_controllers.text_generation_controller`#

Module Contents#

Classes#

`AsyncScheduleLogitsState`	Track logits submitted for the next async-scheduling sample.
`_AsyncScheduleSampleResult`	GPU samples, reusable CPU views, and readiness events for one async step.
`DecodeOnly`	Decode-only state for the consumed and launched forwards.
`DynamicBatchControllerStepResult`	Result of one dynamic-batching controller step.
`_AsyncScheduleRequestResult`	Request state produced by async scheduling bookkeeping.
`_AsyncScheduleLogProbsGPUResult`	GPU logprob outputs awaiting transfer to CPU.
`_AsyncScheduleLogProbsTransfer`	Transient CPU views retaining their GPU sources until D2H completes.
`TextGenerationController`	The text generation controller (the main sampling loop)

API#

class core.inference.text_generation_controllers.text_generation_controller.AsyncScheduleLogitsState#

Track logits submitted for the next async-scheduling sample.

is_valid: bool#: False

cuda_graph_request_count: Optional[int]#: None

token_row_indices: Optional[torch.Tensor]#: None

set_pending( cuda_graph_request_count: Optional[int], token_row_indices: Optional[torch.Tensor] = None, ) → None#

Record logits submitted for the next sample.

Parameters:

cuda_graph_request_count (Optional[int]) – CUDA graph request count for the pending logits, or None when CUDA graphs were not used.
token_row_indices (Optional[Tensor]) – Original GPU input row for each logical token row in the pending forward.

clear() → None#: Clear the pending logits state.

class core.inference.text_generation_controllers.text_generation_controller._AsyncScheduleSampleResult#

GPU samples, reusable CPU views, and readiness events for one async step.

sampled_tokens_gpu: torch.Tensor#: None

sampled_tokens_cpu_view: torch.Tensor#: None

sampled_mtp_tokens_gpu: Optional[torch.Tensor]#: None

sampled_mtp_tokens_cpu_view: Optional[torch.Tensor]#: None

accepted_tokens_cpu_view: Optional[torch.Tensor]#: None

accepted_counts_gpu: Optional[torch.Tensor]#: None

accepted_counts_cpu_view: Optional[torch.Tensor]#: None

accepted_counts_cpu_ready_event: Optional[torch.cuda.Event]#: None

sample_cpu_ready_event: Optional[torch.cuda.Event]#: None

class core.inference.text_generation_controllers.text_generation_controller.DecodeOnly#

Decode-only state for the consumed and launched forwards.

.. attribute:: consumed

Whether the consumed output came from a decode-only forward, or None when no output was consumed.

.. attribute:: launched

Whether the launched forward is decode-only, or None when no real forward was launched.

consumed: Optional[bool]#: None

launched: Optional[bool]#: None

__bool__() → bool#

Return the shared decode-only state when both forwards agree.

Returns:: The common consumed and launched decode-only state.
Return type:: bool
Raises:: ValueError – If either forward is absent or the two states differ.

class core.inference.text_generation_controllers.text_generation_controller.DynamicBatchControllerStepResult#

Result of one dynamic-batching controller step.

.. attribute:: decode_only

Decode-only state for the consumed and launched forwards.

.. attribute:: output

Sampled-step output, or None when no output was produced.

.. attribute:: primer_only

Whether the step launched only an async-scheduling primer.

decode_only: core.inference.text_generation_controllers.text_generation_controller.DecodeOnly#: None

output: Optional[Dict]#: None

primer_only: bool#: False

class core.inference.text_generation_controllers.text_generation_controller._AsyncScheduleRequestResult#

Request state produced by async scheduling bookkeeping.

sampled_tokens_cpu: torch.Tensor#: None

accepted_tokens_cpu: Optional[torch.Tensor]#: None

active_request_ids: torch.Tensor#: None

finished_request_ids: torch.Tensor#: None

survivor_idxs: Optional[torch.Tensor]#: None

newly_paused_request_ids: Optional[torch.Tensor]#: None

evict_request_ids: Optional[torch.Tensor]#: None

class core.inference.text_generation_controllers.text_generation_controller._AsyncScheduleLogProbsGPUResult#

GPU logprob outputs awaiting transfer to CPU.

selected_log_probs: torch.Tensor#: None

top_n_log_probs: Optional[torch.Tensor]#: None

top_n_token_ids: Optional[torch.Tensor]#: None

row_counts: List[int]#: None

top_n_counts: List[int]#: None

skip_prompt_log_probs: List[bool]#: None

num_decode_requests: int#: None

gpu_ready_event: Optional[torch.cuda.Event]#: None

class core.inference.text_generation_controllers.text_generation_controller._AsyncScheduleLogProbsTransfer#

Transient CPU views retaining their GPU sources until D2H completes.

selected_log_probs_cpu_view: torch.Tensor#: None

top_n_log_probs_cpu_view: Optional[torch.Tensor]#: None

top_n_token_ids_cpu_view: Optional[torch.Tensor]#: None

row_counts: List[int]#: None

top_n_counts: List[int]#: None

skip_prompt_log_probs: List[bool]#: None

num_decode_requests: int#: None

cpu_ready_event: Optional[torch.cuda.Event]#: None

gpu_result: core.inference.text_generation_controllers.text_generation_controller._AsyncScheduleLogProbsGPUResult#: None

class core.inference.text_generation_controllers.text_generation_controller.TextGenerationController( inference_wrapped_model: megatron.core.inference.model_inference_wrappers.abstract_model_inference_wrapper.AbstractModelInferenceWrapper, tokenizer, )#

The text generation controller (the main sampling loop)

This class tokenizes the input, runs inference, samples from logits, and detokenizes the output.

Parameters:

inference_wrapped_model (AbstractModelInferenceWrapper) – A model that is wrapped using the specs given in the abstract_model_inference_wrapper.py
tokenizer (type) – Tokenizer used for tokenizing and detokenizing the prompts

Initialization

set_stop_word_finished_ids_callback(callback)#

Set a callback to get request IDs that should be marked as finished due to stop words.

The callback should have signature: callback(active_request_ids: List[int]) -> Set[int] Returns a set of request IDs from active_request_ids that should be marked as finished.

Parameters:: callback – Function that returns request IDs to mark as finished.

_init_dynamic_sampling_tensors()#: Initialize tensors needed for dynamic sampling.

_init_mtp_sampling_tensors()#

Pre-allocate MTP sampling tensors.

Addresses must be stable across steps for CUDA graph capture.

static tokenize_prompt( tokenizer, prompt: str, add_BOS: bool = False, ) → List[int]#

Utility to tokenize the input prompts.

Parameters:

tokenizer – The tokenizer to use.
prompt (str) – The input prompt.
add_BOS (bool) – Whether to add a BOS token.

Returns:

Returns the tokenized prompt.

Return type:

List[int]

static detokenize( tokenizer, tokens: List[int], remove_EOD: bool = True, skip_special_tokens: bool = True, ) → str#

Detokenize a sequence of token IDs, optionally removing trailing EOD tokens and handling skip_special_tokens for different tokenizer APIs.

Parameters:

tokenizer – The tokenizer to use for detokenization.
tokens (List[int]) – The token IDs to convert back to text.
remove_EOD (bool) – Whether to remove trailing EOD tokens before detokenization. Defaults to True.
skip_special_tokens (bool) – Whether to remove special tokens (e.g. BOS/EOS) during detokenization. Only passed through if the tokenizer supports it.

Returns:

The detokenized string.

Return type:

str

detokenize_generations( tokens_gpu_tensor: torch.Tensor, lengths_gpu_tensor: torch.Tensor, detokenize_segments: bool, skip_special_tokens: bool = True, ) → tuple[str, Optional[List[List[str]]]]#

Detokenize the generated tokens.

Parameters:

tokens_gpu_tensor (torch.Tensor) – Tensor containing the tokens
lengths_gpu_tensor (torch.Tensor) – Tensor containing the lengths of each sequence
detokenize_segments (bool) – If True, returns individually detokenized tokens. If False,
in (returns None as second element. Helpful for understanding per-token boundaries)
text. (generated)
skip_special_tokens (bool) – If True removes special tokens like bos
detokenization. (during)

Returns:

A tuple containing:

str: The complete detokenized text
List[str] | None: List of segmented tokens if detokenize_segments is True, else None

Return type:

tuple[str, List[str] | None]

sample_from_logits(

last_token_logits: torch.Tensor,

sampling_params: Optional[megatron.core.inference.sampling_params.SamplingParams] = None,

vocab_size: Optional[int] = None,

generation_started: Optional[torch.Tensor] = None,

top_n_logprobs_dict: Dict[int, List[Dict[str, float]]] = None,

logits: Optional[torch.Tensor] = None,

**kwargs,

) → torch.Tensor#

Samples the logits to generate outputs

Given the logits of the last token, this function samples it according to the parameters defined in sampling_params and returns the samples. If sampling parameters top_n_logprobs > 0 at each step it also updates the top_n_logprobs dict.

Parameters:

last_token_logits (torch.Tensor) – The last token logits. A tensor of size [batch_size, vocab_size]
sampling_params (SamplingParams) – The parameters to use for inference.
vocab_size (int) – Obtained from the tokenizer. Defaults to None
generation_started (torch.Tensor) – A boolean tensor of shape [batch_size]. True indicates the prompt at that index has started generating tokens.
top_n_logprobs_dict (top_n_logprobs_dict) – The dict to be updated

Returns:

1D tensor with [batch_size] elements top_n_logprobs_this_step (torch.return_types.topk): a topk tensor with values as logits and indices as the top k elements. None if sampling params top_n_logprobs is 0.

Return type:

sampled_logits (torch.Tensor)

update_generation_status( updated_prompts_tokens: torch.Tensor, generation_started: torch.Tensor, current_context_end_position: int, is_generation_done_tensor: torch.Tensor, generated_sequence_lengths: torch.Tensor, termination_id: Optional[int] = None, ) → Tuple[torch.Tensor, torch.Tensor]#

Checks which prompts have reached an end condition

We check which prompts have reached an end condition and set the corresponding flags of the is_generation_done_tensor to True. The generated sequence lengths increase as we keep generating, until that prompts hits an end condition. The generation_started tensor determines which prompts have started generating.

Parameters:

updated_prompts_tokens (torch.Tensor) – The prompts tokens updated with the latest generated tokens. A tensor of shape [batch_size, max_seq_len] (i.e max_seq_len = max_prompt_len + tokens_to_generate)
generation_started (torch.Tensor) – A boolean tensor of shape [batch_size]. True indicates the prompt at that index has started generating tokens.
current_context_end_position (int) – An integer indicating which position to extract from the prompts tokens to get the latest generated tokens.
is_generation_done_tensor (torch.Tensor) – A boolean tensor of shape [batch_size]. True indicates the prompt at that index has reached end condition.
generated_sequence_lengths (torch.Tensor) – A int tensor of shape [batch_size]. Each value represents the generated sequence lengths for that prompt.

Returns:

Returns the boolean is_generation_done_tensor and the generated_sequence_lengths after updating it

Return type:

Tuple[torch.Tensor, torch.Tensor]

pad_input_prompt_tokens( batch_prompt_tokens_list: List[List[int]], padded_batch_size: int, padded_sequence_length: int, ) → torch.Tensor#

Method to pad input prompts

Given a list of prompts, pad them all to uniform length

Parameters:

batch_prompt_tokens_list (List[List[int]]) – A list containing the prompt tokens
padded_batch_size (int) – The maximum number of requests for this batch
padded_sequence_length (int) – The maximum number of input + output tokens for this batch

Returns:

A torch tensor of shape [padded_batch_size, padded_sequence_length]

Return type:

torch.Tensor

unpad_input_prompt_tokens( padded_batch_prompt_tokens: torch.Tensor, original_batch_size: int, )#

Truncates the given input tensor back to the original prompt size before padding.

Parameters:

padded_batch_prompt_tokens (torch.Tensor) – The padded tokens tensor
original_batch_size (int) – The original batch size before padding

_dynamic_step_context_init( construct_graph_dimensions: Optional[megatron.core.inference.batch_dimensions_utils.InferenceBatchDimensions] = None, is_dummy_forward: bool = False, transfer_bookkeeping_to_gpu: bool = True, record_bookkeeping_done_event: bool = False, ) → Tuple[torch.Tensor, torch.Tensor, Optional[torch.cuda.Event]]#

Initializes the inference context for dynamic batching.

Parameters:

construct_graph_dimensions (Optional[InferenceBatchDimensions]) – The graph config to use for constructing the cuda graphs.
is_dummy_forward (bool) – Whether we are running an expert parallel dummy forward pass
transfer_bookkeeping_to_gpu (bool) – Whether to publish the prepared CPU bookkeeping snapshot to GPU before returning.
record_bookkeeping_done_event (bool) – Whether to record an event after the bookkeeping H2D transfer.

Returns:

The active input IDs, position IDs, and optional bookkeeping H2D completion event.

Return type:

Tuple[Tensor, Tensor, Optional[torch.cuda.Event]]

_dynamic_step_forward_logits( input_ids: torch.Tensor, position_ids: torch.Tensor, )#

Forward step the model to get logits for dynamic batching.

This also handles logits-broadcasting for pipeline parallelism.

Parameters:

input_ids (Tensor) – The input token IDs.
position_ids (Tensor) – The position IDs.

_rewind_kv_cache( accepted_counts_cpu: Optional[torch.Tensor] = None, ) → tuple#

Update the KV cache bookkeeping for speculative decoding.

After forward pass with speculative tokens, some tokens may be rejected. This function “rewinds” the KV cache bookkeeping to reflect only the accepted tokens. The core bookkeeping rewind runs on CPU (mutating the CPU source-of-truth tensors in place); the Mamba hybrid-model state update stays on GPU because it operates on GPU-resident state buffers.

Parameters:: accepted_counts_cpu (Optional[Tensor]) – Accepted MTP draft counts already copied to CPU. When omitted, this method performs the legacy D2H copy.
Returns:: Blocks detached by rewind and the mask selecting valid block IDs.
Return type:: tuple

_sample_from_logits_2d(logits_2d: torch.Tensor) → torch.Tensor#

Sample tokens from 2D logits using existing sampling parameters.

Parameters:: logits_2d (Tensor) – Logits of shape [num_requests, vocab_size].
Returns:: Sampled tokens of shape [num_requests].
Return type:: Tensor

_compute_serial_mtp_and_sample( base_position: Optional[torch.Tensor] = None, ) → None#

Compute MTP logits serially after verification and sample speculative tokens.

This ensures that MTP predictions are always conditioned on verified tokens. Each MTP depth receives the correctly sampled token from the previous depth (or the base token for depth 0) rather than stale speculative tokens from the previous step.

When sequence parallelism is active, hidden states are kept in SP format (scattered along the first dimension) between MTP depths to avoid a redundant gather + scatter round-trip per depth.

Parameters:: base_position (Optional[Tensor]) – GPU position of the first new MTP draft for each request. Legacy scheduling derives it from rewound CPU state.

_verify_speculative_tokens( output_tokens: torch.Tensor, input_tokens_required: torch.Tensor, num_decode_requests: int, num_prefill_requests: int, active_request_count: int, ) → tuple#: Verify speculative tokens against input tokens (Triton kernel).

_dynamic_step_sample_logits_and_verify_tokens( input_ids: torch.Tensor, token_row_indices: Optional[torch.Tensor] = None, ) → None#

Sample MTP logits and verify pending draft tokens.

Parameters:

input_ids (Tensor) – Input token storage used by the pending forward.
token_row_indices (Optional[Tensor]) – Original GPU input row for each current logical token row after survivor compaction.

_prepare_speculative_tokens_for_next_forward_pass( num_decode_requests: int, output_tokens: torch.Tensor, required_logit_indices: torch.Tensor, last_one_indices: torch.Tensor, accepted_tokens_mask: torch.Tensor, input_tokens_required: torch.Tensor, )#

Prepare accepted speculative tokens for the next forward pass (Triton kernel).

.. rubric:: Example

_dynamic_step_sample_logits()#: Sample tokens from logits for dynamic batching.

_active_requests_sampling_filter_flags( active_request_count: Optional[int] = None, ) → Tuple[bool, bool]#

Return (no_top_k, no_top_p) batch-level escape hatches for the active batch.

These drive the FlashInfer sampler’s dispatch (top-p-only / top-k-only / joint) and are read from the pinned CPU sampling metadata, so they incur no GPU sync. A filter is “absent” only when NO active request uses it. Padded rows carry a neutral 0 and never flip a flag.

_dynamic_step_log_probs_bookkeeping() → Tuple[bool, bool]#

Perform bookkeeping necessary to compute log probs for dynamic batching.

Returns:: Whether to return the sampled log_probs. return_top_n_logprobs (bool): Whether to return top-n log_probs.
Return type:: return_log_probs (bool)

_router_record_bookkeeping() → Optional[numpy.ndarray]#

Collect flat routing indices for MoE router recording.

Retrieves recorded routing decisions via the context’s routing_metadata (which handles CUDA graph static buffers), performs the TP all-gather when sequence parallelism is active, strips CUDA padding, and returns a flat CPU numpy array aligned with the context’s active-token layout. Must be called while context attributes are still valid (before request transitions).

Returns:: Flat routing array of shape [active_token_count, num_layers, topk], or None if routing replay is disabled or no routing data was recorded.
Return type:: Optional[np.ndarray]

_dynamic_step_calculate_log_probs() → Optional[torch.Tensor]#: Calculate log probs from logits.

_dynamic_step_calculate_log_probs_speculative() → Tuple[List[List[float]], torch.Tensor]#

Calculate log probs from logits for speculative decoding.

For decode requests, computes log probs for each accepted speculative token and the newly sampled token using the main model logits. For prefill requests, handles prompt log probs the same way as non-speculative decoding.

The main model logits at position j predict the token at position j+1. So:

log_prob(accepted_token[j]) comes from logits at position j
log_prob(newly_sampled_token) comes from logits at position accepted_count

Returns:

 log_probs_list: List of lists, one per active request, containing
     log probs for the tokens emitted in this step.
 log_probs_tensor: Full log_softmax tensor for top-n computation.

Return type:

Tuple of (log_probs_list, log_probs_tensor)

_dynamic_step_calculate_top_n_logprobs_speculative( log_probs_tensor: torch.Tensor, ) → Optional[Dict[int, List[Tuple[torch.Tensor, torch.Tensor]]]]#

Calculate top-n log probs for speculative decoding.

For decode requests, computes top-n at each position that produced an emitted token (accepted speculative positions + the newly sampled position). For prefill requests, behaves identically to the non-speculative path.

Parameters:: log_probs_tensor (Tensor) – Pre-computed log_softmax tensor from _dynamic_step_calculate_log_probs_speculative.
Returns:: A dictionary mapping request_idx to list of (top_n_values, top_n_indices) tuples, one per emitted token position.

_dynamic_step_calculate_top_n_logprobs( log_probs_tensor: Optional[torch.Tensor] = None, ) → Optional[Dict[int, List[Tuple[torch.Tensor, torch.Tensor]]]]#

Calculate top-n log probs from logits for dynamic batching.

Parameters:: log_probs_tensor (Optional[Tensor]) – Pre-computed log probabilities tensor. If provided, avoids recomputing log_softmax. Should be the tensor returned by calculate_log_probs.
Returns:: A dictionary mapping request_idx to list of (top_n_logprobs, top_n_indices) tuples. Each tuple in the list represents one token position.

_run_dummy_base_forward( input_ids: torch.Tensor, position_ids: torch.Tensor, ) → None#

Run the base-model portion of an expert-parallel dummy step.

Parameters:

input_ids (Tensor) – Dummy input token IDs.
position_ids (Tensor) – Dummy input position IDs.

_run_dummy_serial_mtp_forward() → None#

Run dummy MTP forward passes to participate in EP collectives.

When speculative decoding is active and MTP layers contain MoE sublayers (inherited from the decoder layer spec), each serial MTP step triggers EP all-to-all collectives. The dummy EP rank must issue matching collective calls so the real ranks do not hang.

This mirrors the structure of _compute_serial_mtp_and_sample:

On the last PP stage (where MTP resides): run compute_mtp_single_step with dummy tensors so the MoE all-to-all is executed.
When PP > 1: participate in the broadcast_from_last_pipeline_stage that the real ranks also perform.

_run_dummy_legacy_step( input_ids: torch.Tensor, position_ids: torch.Tensor, ) → None#

Run a legacy dummy step in base-forward then MTP order.

Parameters:

input_ids (Tensor) – Dummy input token IDs.
position_ids (Tensor) – Dummy input position IDs.

_run_dummy_async_sched_step( input_ids: torch.Tensor, position_ids: torch.Tensor, ) → None#

Run an async-scheduling dummy step in MTP then base-forward order.

Parameters:

input_ids (Tensor) – Dummy input token IDs.
position_ids (Tensor) – Dummy input position IDs.

dummy_forward() → None#: Run the mode-specific dummy step used by idle expert-parallel ranks.

_transfer_samples_to_cpu(active_request_count: int) → tuple#

Batch GPU-to-CPU transfer of sampled tokens.

Called at the boundary between GPU sampling and CPU bookkeeping. After this returns, all sampled data is on CPU and the remainder of the step is 100% CPU.

Returns:: (sampled_tokens_cpu, sampled_mtp_tokens_cpu) where sampled_mtp_tokens_cpu is None when speculative decoding is off.
Return type:: tuple

_apply_stop_word_finished_ids( active_request_ids: torch.Tensor, active_request_mask: torch.Tensor, ) → None#

Mark requests whose generated output matched a stop word as finished.

Parameters:

active_request_ids (Tensor) – IDs for requests active during the current step.
active_request_mask (Tensor) – Mask updated in place for requests that remain active.

_dynamic_step_context_bookkeeping() → Dict[str, torch.Tensor]#

Update the dynamic inference context after sampling.

Parameters:

new_sample (Tensor) – The newly sampled tokens.
request_metadata (Optional[Dict[str, Tensor]]) – An override for the tensors that manage request metadata, such as sampling parameters. By default, this metadata is retrieved from the context.

Returns:

A dictionary containing: active_request_ids (Tensor): Current active request IDs. newly_paused_request_ids (Tensor): Newly paused request IDs. finished_request_ids (Tensor): Finished request IDs.

Return type:

Dict [str, Tensor]

_validate_async_sched_support_for_step( run_async_overlap: bool, ) → None#

Validate controller/context state for async scheduling.

Parameters:: run_async_overlap (bool) – Whether this step uses overlap ordering.

Raises if the current step does not support async scheduling.

_compact_async_sched_logits(survivor_idxs: torch.Tensor) → None#

Compact pending logits and sampling metadata into survivor order.

Parameters:: survivor_idxs (Tensor) – Active-row indices for requests that remain active after async scheduling.

static _synchronize_async_sched_event( event: Optional[torch.cuda.Event], ) → None#

Block the host until an async-scheduling CUDA event completes.

Parameters:: event (Optional[torch.cuda.Event]) – CUDA event to synchronize, or None when no CUDA work was recorded.

_copy_async_sched_accepted_counts_to_cpu( accepted_counts_gpu: torch.Tensor, ) → Tuple[torch.Tensor, Optional[torch.cuda.Event]]#

Start copying MTP acceptance counts into their reusable CPU buffer.

Parameters:: accepted_counts_gpu (Tensor) – Accepted MTP draft count per active request.
Returns:: Transient CPU view and its copy-completion event.
Return type:: Tuple[Tensor, Optional[torch.cuda.Event]]

_copy_async_sched_sample_to_cpu( sampled_tokens_gpu: torch.Tensor, sampled_mtp_tokens_gpu: Optional[torch.Tensor] = None, accepted_tokens_gpu: Optional[torch.Tensor] = None, ) → Tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.cuda.Event]]#

Start copying async sampling outputs into reusable CPU buffers.

Parameters:

sampled_tokens_gpu (Tensor) – Sampled base token IDs for active requests.
sampled_mtp_tokens_gpu (Optional[Tensor]) – Generated MTP draft token IDs.
accepted_tokens_gpu (Optional[Tensor]) – Accepted pending MTP draft token IDs.

Returns:

Transient CPU views for base, draft, and accepted tokens plus the copy-completion event.

Return type:

Tuple[Tensor, Optional[Tensor], Optional[Tensor], Optional[torch.cuda.Event]]

_build_async_sched_request_state( sampled_tokens_cpu: torch.Tensor, resolved_sequence_lengths: torch.Tensor, ) → Tuple[torch.Tensor, torch.Tensor, torch.Tensor]#

Build request IDs and the active/finished mask for resolution.

Parameters:

sampled_tokens_cpu (Tensor) – Sampled CPU token IDs for active requests.
resolved_sequence_lengths (Tensor) – Sequence lengths after accepting current output and before preparing unverified successor tokens.

Returns:

Active request IDs, finished request IDs, and the active-request mask.

Return type:

Tuple[Tensor, Tensor, Tensor]

_run_async_sched_sample() → core.inference.text_generation_controllers.text_generation_controller._AsyncScheduleSampleResult#

Sample active requests and start transferring their tokens to CPU.

Returns:: Base-token samples and transfer state.
Return type:: _AsyncScheduleSampleResult

_run_async_sched_sample_mtp() → core.inference.text_generation_controllers.text_generation_controller._AsyncScheduleSampleResult#

Verify pending MTP logits and generate the next draft tokens.

Returns:: Base, draft, accepted-token, and transfer state.
Return type:: _AsyncScheduleSampleResult

_run_async_sched_mtp_rewind( sample_result: core.inference.text_generation_controllers.text_generation_controller._AsyncScheduleSampleResult, ) → None#

Rewind rejected MTP KV state before preparing the successor.

Parameters:: sample_result (_AsyncScheduleSampleResult) – Verified MTP sampling state.

_run_async_sched_log_probs( sample_result: core.inference.text_generation_controllers.text_generation_controller._AsyncScheduleSampleResult, ) → Optional[core.inference.text_generation_controllers.text_generation_controller._AsyncScheduleLogProbsGPUResult]#

Calculate selected and top-n log probabilities on the GPU.

Parameters:: sample_result (_AsyncScheduleSampleResult) – Sampled and accepted tokens for active requests.
Returns:: GPU logprob outputs and their completion event, or None when no request needs logprobs.
Return type:: Optional[_AsyncScheduleLogProbsGPUResult]

_copy_async_sched_log_probs_to_cpu( gpu_result: Optional[core.inference.text_generation_controllers.text_generation_controller._AsyncScheduleLogProbsGPUResult], ) → Optional[core.inference.text_generation_controllers.text_generation_controller._AsyncScheduleLogProbsTransfer]#

Start selected and top-n logprob transfers to reusable CPU buffers.

Parameters:: gpu_result (Optional[_AsyncScheduleLogProbsGPUResult]) – GPU outputs produced by the current sampling step.
Returns:: Transient CPU views, transfer-completion event, and retained GPU sources, or None.
Return type:: Optional[_AsyncScheduleLogProbsTransfer]

static _materialize_async_sched_log_probs( transfer: Optional[core.inference.text_generation_controllers.text_generation_controller._AsyncScheduleLogProbsTransfer], accepted_counts_cpu: Optional[torch.Tensor] = None, ) → Tuple[Optional[List[List[float]]], Optional[Dict[int, List[Tuple[torch.Tensor, torch.Tensor]]]]]#

Convert completed CPU transfer views to the legacy result format.

Parameters:

transfer (Optional[_AsyncScheduleLogProbsTransfer]) – Completed logprob transfer for the current step.
accepted_counts_cpu (Optional[Tensor]) – Accepted MTP draft count per active request, or None for one-token decoding.

Returns:

Tuple containing selected logprobs per request and optional top-n values/token IDs per request.

_run_async_sched_prepare() → Tuple[torch.Tensor, torch.Tensor]#

Prepare decode requests and return live GPU forward-input views.

The returned views have their final shape and stable backing storage, but their contents are populated later. Sampling updates the input-ID view, and deferred bookkeeping publication updates the position-ID view.

Returns:: Live GPU input-ID and position-ID views for the speculative forward.
Return type:: Tuple[Tensor, Tensor]

_run_async_sched_publish_bookkeeping() → Optional[torch.cuda.Event]#

Publish prepared bookkeeping without overwriting GPU input token IDs.

Returns:: Event marking bookkeeping H2D completion.
Return type:: Optional[torch.cuda.Event]

_commit_mamba_intermediate_states() → None#: Commit prefix-cacheable Mamba states produced by the current forward.

_run_async_sched_forward( input_ids_gpu_view: torch.Tensor, position_ids_gpu_view: torch.Tensor, ) → None#

Run one dynamic forward pass and cache logits for async scheduling.

Parameters:

input_ids_gpu_view (Tensor) – Live GPU view of the input token IDs.
position_ids_gpu_view (Tensor) – Live GPU view of the position IDs.

_run_dummy_async_sched_base_step() → None#: Run the base-forward half of an async EP step after local work finishes.

_run_async_sched_forward_primer() → Tuple[bool, Optional[torch.cuda.Event]]#

Launch the initial forward when no valid logits state exists.

Returns:: Whether this call launched the forward primer and its bookkeeping H2D completion event.
Return type:: Tuple[bool, Optional[torch.cuda.Event]]

_run_async_sched_resolve( sample_result: core.inference.text_generation_controllers.text_generation_controller._AsyncScheduleSampleResult, resolved_sequence_lengths: torch.Tensor, ) → core.inference.text_generation_controllers.text_generation_controller._AsyncScheduleRequestResult#

Resolve request state and compact speculative forward logits.

Parameters:

sample_result (_AsyncScheduleSampleResult) – Sampling outputs in reusable CPU views.
resolved_sequence_lengths (Tensor) – Sequence lengths after accepting current output and before preparing unverified successor tokens.

Returns:

Sampled tokens, resolved request row sets, and survivor indices.

Return type:

_AsyncScheduleRequestResult

_run_async_sched_update_requests( sample_result: core.inference.text_generation_controllers.text_generation_controller._AsyncScheduleSampleResult, resolved_sequence_lengths: torch.Tensor, ) → core.inference.text_generation_controllers.text_generation_controller._AsyncScheduleRequestResult#

Run complete request lifecycle bookkeeping for a no-overlap step.

Parameters:

sample_result (_AsyncScheduleSampleResult) – Sampling outputs in reusable CPU views.
resolved_sequence_lengths (Tensor) – Sequence lengths after accepting the current output.

Returns:

Stable sampled output and lifecycle results.

Return type:

_AsyncScheduleRequestResult

_build_async_sched_step_result( request_result: core.inference.text_generation_controllers.text_generation_controller._AsyncScheduleRequestResult, cuda_graph_request_count: Optional[int], decode_only: core.inference.text_generation_controllers.text_generation_controller.DecodeOnly, log_probs: Optional[List[List[float]]], top_n_logprobs: Optional[Dict[int, List[Tuple[torch.Tensor, torch.Tensor]]]], *, count_compaction: bool, ) → core.inference.text_generation_controllers.text_generation_controller.DynamicBatchControllerStepResult#

Build the public result and update async-scheduling counters.

Parameters:

request_result (_AsyncScheduleRequestResult) – Completed request bookkeeping.
cuda_graph_request_count (Optional[int]) – CUDA graph request count used by the consumed forward.
decode_only (DecodeOnly) – Decode-only state for the consumed and launched forwards.
log_probs (Optional[List[List[float]]]) – Selected-token log probabilities grouped by active request.
top_n_logprobs (Optional[Dict[int, List[Tuple[Tensor, Tensor]]]]) – Top-n log probabilities and token IDs grouped by active request.
count_compaction (bool) – Whether finished requests discarded successor rows.

Returns:

Completed sampled-step result.

Return type:

DynamicBatchControllerStepResult

async _run_async_sched_step_no_overlap( *, schedule_waiting_requests: Optional[Callable[[], None]], ) → core.inference.text_generation_controllers.text_generation_controller.DynamicBatchControllerStepResult#

Run sample/MTP -> update -> admit -> forward.

The first call in an active chain has no pending output. It skips the first two phases, admits requests, and launches a primer-only forward.

Parameters:: schedule_waiting_requests (Optional[Callable[[], None]]) – Engine callback that admits eligible non-chunked prefill requests.
Returns:: Primer-only state or sampled output.
Return type:: DynamicBatchControllerStepResult

async _run_async_sched_step_overlap() → core.inference.text_generation_controllers.text_generation_controller.DynamicBatchControllerStepResult#

Run prepare -> sample -> forward -> resolve with one token per request.

Returns:: Completed sampled-step result.
Return type:: DynamicBatchControllerStepResult

async _run_async_sched_step_overlap_mtp() → core.inference.text_generation_controllers.text_generation_controller.DynamicBatchControllerStepResult#

Run sample/MTP -> prepare -> forward -> resolve with MTP.

Returns:: Completed sampled-step result.
Return type:: DynamicBatchControllerStepResult

async _run_legacy_step( skip_bookkeeping: Optional[bool] = False, ) → core.inference.text_generation_controllers.text_generation_controller.DynamicBatchControllerStepResult#

Forward step the model and update the inference context.

Parameters:: skip_bookkeeping (Optional[bool]) – If true, skip the context bookkeeping step.
Returns:: Legacy sampled-step output and its decode-only state.
Return type:: DynamicBatchControllerStepResult

async async_generate_output_tokens_dynamic_batch( skip_bookkeeping: Optional[bool] = False, *, run_async_overlap: bool = True, schedule_waiting_requests: Optional[Callable[[], None]] = None, ) → core.inference.text_generation_controllers.text_generation_controller.DynamicBatchControllerStepResult#

Forward step the model and update the inference context.

Parameters:

skip_bookkeeping (Optional[bool]) – If true, skip context bookkeeping on the legacy path.
run_async_overlap (bool) – Whether to run the overlap ordering.
schedule_waiting_requests (Optional[Callable[[], None]]) – Engine callback used by the no-overlap path to admit eligible prefill requests.

Returns:

One controller-step result.

Return type:

DynamicBatchControllerStepResult

generate_output_tokens_dynamic_batch( loop: Optional[asyncio.AbstractEventLoop] = None, ) → Optional[Dict]#

Synchronously run dynamic batching through any primer-only calls.

Parameters:: loop (Optional[asyncio.AbstractEventLoop]) – Event loop used to run the asynchronous controller.
Returns:: Step output, or None when no work is active.
Return type:: Optional[Dict]

_update_top_n_logprobs_dict( top_n_logprobs_this_step: torch.Tensor, top_n_logprobs_indices: torch.Tensor, mask: torch.Tensor, top_n_logprobs_dict: Dict[int, List[Dict[str, float]]], )#

Function to update the top_n_logprobs at each step

This function goes through the topn logprobs generated for each, and for whichever batch has started generating tokens, it updates the top_n_logprobs_dict with the decoded token (string) as the key and the logit as the value. top_n_logprobs_dict has as keys the batch idx, the values is a list, where each element represents a dictionary of decoded token as key and logit as value generated at each step

Parameters:

top_n_logprobs_this_step (torch.Tensor) – The top n logprob values
top_n_logprobs_indices (torch.Tensor) – The indices corresponding to the top n logprobs
mask (torch.Tensor) – A mask to indicate which requests should append to the dict
top_n_logprobs_dict (top_n_logprobs_dict) – The dict to be updated

generate_all_output_tokens_static_batch( active_requests: OrderedDict[int, megatron.core.inference.inference_request.InferenceRequest], active_streams: Optional[OrderedDict[str, megatron.core.inference.async_stream.AsyncStream]] = None, ) → OrderedDict[int, megatron.core.inference.inference_request.InferenceRequest]#

Utility to generate all the output tokens and probabilities for the prompts.

This utility generates the output tokens for a static batch. It runs the forward steps till all prompts complete generation, updates the status of these requests to completed, adds the generated result and returns these requests

Parameters:: active_requests (OrderedDict[int, InferenceRequest]) – The input active requests.
Returns:: The result for each of the incoming requests
Return type:: OrderedDict[int, InferenceRequest]

prep_inference_input( prompts_tokens: torch.Tensor, active_requests: OrderedDict[int, megatron.core.inference.inference_request.InferenceRequest], use_attention_mask: bool = False, ) → Dict[str, Any]#

Preparing input data for inference, using respective wrapper’s prep_inference_input method # pylint: disable=line-too-long

Parameters:

prompts_tokens (torch.Tensor) – A tensor of shape [batch_size, max_sequence_length]
active_requests (OrderedDict[int, InferenceRequest]) – The input active requests
use_attention_mask (bool) – Whether to use an attention mask. Should be set to True only when exclusively doing prefill (no decode) with variable prompt lengths.

Returns:

A dict of the inference input for the current batch.

stream_tokens( sampling_params: megatron.core.inference.sampling_params.SamplingParams, request_ids: List[int], requests: List[megatron.core.inference.inference_request.InferenceRequest], streams: List[megatron.core.inference.async_stream.AsyncStream], generation_started: List[bool], is_generation_done: List[bool], tokens: torch.Tensor, prompt_lengths: List[int], generated_lengths: List[int], output_log_probs: Union[torch.Tensor, None], )#

Asynchronously streams tokens for the given requests.

Parameters:

sampling_params (SamplingParams) – The sampling parameters.
request_ids (List[int]) – The request IDs.
request (List[InferenceRequest]) – The requests.
stream (List[AsyncStream]) – The streams over which to send tokens.
generation_started (List[bool]) – Whether the decode step has started.
is_generation_done (List[bool]) – Whether generation has completed.
tokens (torch.Tensor) – The tokens for this request.
prompt_lengths (List[int]) – The number of prompt tokens for each request.
generated_lengths (List[int]) – The number of output tokens for each request.
output_log_probs (torch.Tensor, optional) – The log probs for each request.

core.inference.text_generation_controllers.text_generation_controller#

Module Contents#

Classes#

API#

`core.inference.text_generation_controllers.text_generation_controller`#