core.inference.config

Module Contents

Classes

MambaInferenceStateConfig	Config for initializing recurrent mixer inference state tensors.
PrefixCachingEvictionPolicy	Eviction policy for prefix caching blocks.
PrefixCachingCoordinatorPolicy	Routing policy for the DP inference coordinator with prefix caching.
MediaCacheCoordinatorPolicy	Routing policy for the DP inference coordinator with media caching.
KVCacheManagementMode	Mode for handling large tensors (KV cache, Mamba states) during suspend/resume.
CudaGraphSizingDistribution	How CUDA graph token-count sizes are spaced when generating the captured graphs.
AsyncScheduleMode	Async scheduling mode for dynamic inference.
ImageProcessingConfig	Configuration for converting raw images into model input tensors.
VideoProcessingConfig	Configuration for decoding raw video bytes into model input tensors.
MediaPromptSpec	Map one API media type to the model’s prompt-token contract.
MultimodalPromptConfig	Prompt contracts used to lower structured image/video blocks.
InferenceConfig	Config for inference.

Functions

mtp_layer_types_from_model	Layer types of one MTP draft-head depth, or None for a non-hybrid model.
routes_on_prefix	Whether policy needs per-request block hashes to make a routing decision.

API

class core.inference.config.MambaInferenceStateConfig

Config for initializing recurrent mixer inference state tensors.

Note that we maintain separate metadata for decode, regular prefill, and chunked prefill requests because the recurrent kernels do not yet support mixing these. Once the kernels have been updated we can simplify this code.

layer_type_list: List[str]

None

A list of strings that indicates the layer type (Mamba / GDN / Attention / MLP) for each layer. See megatron/core/models/hybrid/hybrid_layer_allocation.py for the list of symbols.

conv_states_shape: Tuple[int]

None

Recurrent mixer’s conv state shape per request.

ssm_states_shape: Tuple[int]

None

Recurrent mixer state shape per request.

conv_states_dtype: torch.dtype

None

The dtype to use for the Mamba conv state tensor. Defaults to the model dtype.

ssm_states_dtype: torch.dtype

None

The dtype to use for Mamba SSM state. Batch-invariant mode requires FP32.

mamba_chunk_size: int

128

The chunk size used by the Mamba SSM Triton kernels.

ssm_chunk_alignment: Optional[int]

None

Token quantum that a prefill chunk boundary must land on for the model’s SSM mixers to see a clean chunk boundary. Defaults to mamba_chunk_size, which is correct for any Mamba-only model.

This is the mixers’ shared ssm_inference_chunk_size, which is not always their chunk_size: the forked Gated Delta Product prefill kernels run at a fixed 64 whatever chunk_size says. from_model asserts every SSM layer agrees rather than reconciling a mixed stack. Only the paths that genuinely require an aligned boundary consult it – batch-invariant chunked prefill, which replays the partial tail at decode, and recurrent-state extraction for prefix caching, which can only snapshot at a chunk boundary. Ordinary chunked prefill splits anywhere, because each step re-chunks from its own slice start.

gdp_num_householder: int

0

Number of Householder copies of the Gated Delta Product layers, or 0 if the model has none. Sizes the GDP chunk descriptors used by the forked prefill kernels, whose Householder-expanded token stream is this many times longer.

__post_init__()

classmethod from_model(

model: megatron.core.transformer.module.MegatronModule,

conv_states_dtype: Optional[torch.dtype] = None,

ssm_states_dtype: Optional[torch.dtype] = None,

) → Optional[core.inference.config.MambaInferenceStateConfig]

Return recurrent inference state config for a Mamba or GDN hybrid model.

core.inference.config.mtp_layer_types_from_model(

model: megatron.core.transformer.module.MegatronModule,

) → Optional[List[str]]

Layer types of one MTP draft-head depth, or None for a non-hybrid model.

The MTP head’s layer types come from the unified hybrid pattern (”

//…”), which only HybridModel parses. Independent of whether the MAIN decoder has recurrent layers, so it cannot be derived from MambaInferenceStateConfig.

Callers that never enable speculative decoding do not need this: the draft-KV gate requires num_speculative_tokens > 0 first, so leaving InferenceConfig.mtp_layer_type_list at None is correct for them.

class core.inference.config.PrefixCachingEvictionPolicy

Bases: str, enum.Enum

Eviction policy for prefix caching blocks.

Only applies when enable_prefix_caching is True.

Initialization

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

REF_ZERO

‘ref_zero’

Deregister blocks immediately when ref_count hits 0. No caching after release.

LRU

‘lru’

Keep released blocks in hash table. Evict oldest ref=0 blocks when space is needed.

class core.inference.config.PrefixCachingCoordinatorPolicy

Bases: str, enum.Enum

Routing policy for the DP inference coordinator with prefix caching.

Initialization

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

LONGEST_PREFIX

‘longest_prefix’

Route to the rank with the longest consecutive prefix match.

FIRST_PREFIX_BLOCK

‘first_prefix_block’

Route to the rank that has the first block hash cached. O(ranks) check.

LOAD_BALANCED

‘load_balanced’

Route to the rank with the fewest in-flight requests. Ignores prefix affinity.

core.inference.config.routes_on_prefix(policy) → bool

Whether policy needs per-request block hashes to make a routing decision.

Frontends call this to decide whether hashing a prompt is worth anything: under LOAD_BALANCED the coordinator discards the hashes, so computing them is pure overhead on the request path. Kept beside the enum so a new prefix-aware policy only has to be added in one place.

Accepts the enum, its string value, or None (no policy configured).

class core.inference.config.MediaCacheCoordinatorPolicy

Bases: str, enum.Enum

Routing policy for the DP inference coordinator with media caching.

Initialization

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

AFFINITY

‘affinity’

Prefer ranks assigned the same media key when vision embeddings are cached.

LOAD_BALANCED

‘load_balanced’

Ignore media affinity and route using prefix affinity and load.

class core.inference.config.KVCacheManagementMode

Bases: str, enum.Enum

Mode for handling large tensors (KV cache, Mamba states) during suspend/resume.

Initialization

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

PERSIST

‘persist’

Do not deallocate and reallocate large tensors; keep them on GPU.

OFFLOAD

‘offload’

Offload large tensors to CPU during deallocation; onload during allocation.

RECOMPUTE

‘recompute’

Deallocate large tensors and recompute them from scratch during allocation.

class core.inference.config.CudaGraphSizingDistribution

Bases: str, enum.Enum

How CUDA graph token-count sizes are spaced when generating the captured graphs.

EXPONENTIAL — token counts halve from cuda_graph_max_tokens down to tp_size, giving a log-spaced distribution. Bounded relative padding (~2x worst case) at every scale and log2(max_tokens) total graphs.

LINEAR — Include size-1 and size-2 graphs where applicable, linear spacing up until 256, and sparser linear spacing past 256. e.g. [1, 2, 4] + range(8, 256, 8) + range(256, max+1, 16). Higher graph density at the top end.

HYBRID (default) — EXPONENTIAL for prefill and mixed graphs, LINEAR for decode-only graphs. The two serve different ranges: prefill token counts span the whole cuda_graph_max_tokens (thousands), where log spacing keeps padding bounded at ~2x for a handful of graphs, while decode-only counts are capped at max_requests * (num_speculative_tokens + 1) (tens), where halving is far too coarse – a 33-request step would pad up to a 64-request graph. Linear spacing there covers every small request count densely for little extra capture cost.

Initialization

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

EXPONENTIAL

‘exponential’

LINEAR

‘linear’

HYBRID

‘hybrid’

class core.inference.config.AsyncScheduleMode

Bases: str, enum.Enum

Async scheduling mode for dynamic inference.

Initialization

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

LEGACY

‘legacy’

Resolve requests before preparing the next forward pass.

ASYNC

‘async’

Overlap asynchronous scheduling phases by reordering them to prepare-before-resolve.

class core.inference.config.ImageProcessingConfig

Configuration for converting raw images into model input tensors.

patch_dim: int

None

dynamic_resolution: bool

False

use_tiling: bool

False

pixel_shuffle: bool

False

spatial_merge_size: int

1

dynamic_resolution_min_patches: int

1

dynamic_resolution_max_patches: int

128

vision_model_type: str

‘radio’

pixel_mean: Optional[List[float]]

None

pixel_std: Optional[List[float]]

None

img_h: Optional[int]

None

img_w: Optional[int]

None

max_num_tiles: int

1

use_thumbnail: bool

False

num_img_embeddings_per_tile: int

0

class core.inference.config.VideoProcessingConfig

Configuration for decoding raw video bytes into model input tensors.

image_config: core.inference.config.ImageProcessingConfig

None

num_frames: int

8

temporal_patch_size: int

1

frame_manifest_magic: Optional[bytes]

None

Prefix for payloads encoded as magic + UTF-8 {"frame_paths": [...]}.

video_maintain_aspect_ratio: bool

True

class core.inference.config.MediaPromptSpec

Map one API media type to the model’s prompt-token contract.

model_token: str

‘’

prefix: str = <Multiline-String>

suffix: str = <Multiline-String>

input_marker: Optional[str]

None

class core.inference.config.MultimodalPromptConfig

Prompt contracts used to lower structured image/video blocks.

image_spec: core.inference.config.MediaPromptSpec

‘field(…)’

video_spec: core.inference.config.MediaPromptSpec

‘field(…)’

get_spec(modality: str) → core.inference.config.MediaPromptSpec

Return the prompt specification for image or video.

classmethod from_dict(value)

Build from image and video specs.

class core.inference.config.InferenceConfig

Config for inference.

NOTE: Must remain mutually exclusive with the TransformerConfig.

block_size_tokens: int

256

Size of KV cache block size.

buffer_size_gb: int

20

On-GPU portion of the shared KV cache block pool. If unified_memory_level >= 1, then CPU memory is additionally utilized, resulting in a total buffer size of buffer_size_gb + paused_buffer_size_gb.

paused_buffer_size_gb: Optional[int]

None

Memory used to derive the paused-request block retention budget. This does not reserve blocks from active requests: active requests may use the entire shared pool of usable KV cache blocks. When the pool cannot satisfy new allocations, paused requests retain blocks only within this budget and excess paused requests may be evicted. The total buffer size depends on unified_memory_level (uvm): - uvm 0: buffer_size_gb (paused buffer is inclusive) - uvm 1: buffer_size_gb + paused_buffer_size_gb

mamba_inference_state_config: Optional[core.inference.config.MambaInferenceStateConfig]

None

The Mamba inference state config if the model is a hybrid model.

mtp_layer_type_list: Optional[List[str]]

None

Layer types of one MTP draft-head depth, one symbol per layer, or None for a non-hybrid model, whose head is a single attention layer by construction. Read by DynamicInferenceContext to decide whether the MTP draft attention can be given its own KV plane.

mamba_memory_ratio: Optional[float]

None

Percentage of memory buffer to allocate for Mamba states. If not specified, allocates Mamba state tensors for each KV cache block. Only used for hybrid models.

max_requests: Optional[int]

None

Max number of active requests to use for decode-only forward passes. This is primarily limited by the combination of buffer_size_gb and max_sequence_length.

max_tokens: Optional[int]

None

Max number of tokens to use for forward passes. This is primarily limited by prefill activation memory usage. (Defaults to 16384).

unified_memory_level: int

0

Sets unified memory usage within the dynamic inference context. The levels are: 0) no unified memory (default) 1) allocate memory_buffer in unified memory. Eventually, additional levels will be included to control other tensors within the context.

kv_cache_management_mode: core.inference.config.KVCacheManagementMode

None

Mode used to determine how large tensors are handled by the allocate and deallocate methods. See KVCacheManagementMode for options.

num_cuda_graphs: Optional[int]

None

Maximum number of cuda graphs to capture. Graph token counts are spaced from 1 up to a per-graph-type budget:

• Decode-only graphs are always bounded by max_requests * (num_speculative_tokens + 1).

• Prefill/mixed graphs are bounded by cuda_graph_max_tokens by default, or extend up to max_tokens when cuda_graph_all_prefills is set. Due to rounding, the actual number of cuda graphs may not equal this argument.

cuda_graph_mixed_prefill_count: Optional[int]

16

The number of mixed prefill graphs to capture if mixed prefill/decode graphs are enabled.

cuda_graph_sizing_distribution: core.inference.config.CudaGraphSizingDistribution

None

How CUDA graph token counts are spaced. HYBRID (default) applies EXPONENTIAL to prefill and mixed graphs and LINEAR to decode-only graphs, since the two cover ranges that differ by orders of magnitude. EXPONENTIAL halves from cuda_graph_max_tokens down to tp_size (log-spaced, ~log2(max_tokens) graphs). LINEAR uses a range of linear strides (includes small graphs + mid-range linearity + a bigger step size at the top end). Set EXPONENTIAL or LINEAR explicitly to apply one distribution to both families.

use_cuda_graphs_for_non_decode_steps: bool

True

Whether to use CUDA graphs for non-decode steps.

cuda_graph_all_prefills: bool

False

Whether prefill/mixed CUDA graphs should span up to max_tokens. When False (default), prefill/mixed graphs are bounded by cuda_graph_max_tokens. When True, prefill/mixed graph capture is extended to cover the full max_tokens budget.

cuda_graph_max_tokens: int

512

Token ceiling for the largest captured prefill/mixed CUDA graph. This is a raw token count (not scaled by speculative decoding). The effective ceiling is clamped to [max_requests * (num_speculative_tokens + 1), max_tokens] so it never falls below the decode bound nor exceeds the token budget. Ignored when cuda_graph_all_prefills is set, which extends capture to the full max_tokens.

static_kv_memory_pointers: bool

False

Whether the KV cache (and Mamba states) will reside at the same memory addresses after suspend/resume as before. When True, CUDA graphs that reference these buffers remain valid across suspend/resume cycles and do not need to be recaptured. Requires either UVM or torch_memory_saver when kv_cache_management_mode is not PERSIST.

max_sequence_length: int

2560

Max possible sequence length (prompt + output) that will occur.

pg_collection: Optional[megatron.core.process_groups_config.ProcessGroupCollection]

None

A ProcessGroupCollection for distributed execution.

image_preprocessing_config: Optional[core.inference.config.ImageProcessingConfig]

None

Configuration for preprocessing raw image payloads.

video_preprocessing_config: Optional[core.inference.config.VideoProcessingConfig]

None

Configuration for decoding and preprocessing raw video payloads.

use_flashinfer_fused_rope: Optional[bool]

False

If True, use flashinfer’s fused rope implementation. If None, defaults to using flash-infer if available.

materialize_only_last_token_logits: bool

True

Whether to only materialize logits for the last token. This should be set to False if returning log probs.

enable_chunked_prefill: bool

False

Whether to enable chunked prefill.

num_speculative_tokens: int

0

The number of speculative tokens to generate for decode steps.

enable_prefix_caching: bool

False

Whether to enable prefix caching for KV cache block sharing.

vision_embedding_cache_max_bytes: int

0

Maximum GPU bytes retained for reusable vision embeddings.

A value of zero disables the cache. Cache entries use an automatically generated media-content key and, unless allow_stale_multimodal_embeddings is enabled, are discarded whenever the inference engine is suspended or its generation epoch changes.

prefix_caching_eviction_policy: core.inference.config.PrefixCachingEvictionPolicy

None

Eviction policy for prefix caching blocks. See PrefixCachingEvictionPolicy for options.

Only applies when enable_prefix_caching is True.

prefix_caching_coordinator_policy: core.inference.config.PrefixCachingCoordinatorPolicy

None

Routing policy for the DP inference coordinator. See PrefixCachingCoordinatorPolicy for options.

Only applies when enable_prefix_caching is True and using a coordinator.

prefix_caching_routing_alpha: float

1.0

How hard the coordinator penalises load when routing on prefix affinity: score = cache_score - alpha * relative_load.

relative_load is a rank’s in-flight count measured against the fleet mean, so it is zero while ranks are even and grows only as they diverge. Both terms are normalized, which makes alpha dimensionless: 0 is pure prefix affinity, and higher values divert to idle ranks more readily as the fleet becomes lopsided. Must be non-negative; it is not a blend weight and is not capped at 1.

At 1.0 a single request of imbalance across two ranks exactly cancels a full cache hit, so affinity stops being decisive as soon as the fleet is uneven at all. The default keeps a hit decisive against mild imbalance while still diverting to idle ranks once ranks genuinely diverge. Larger fleets are less sensitive, since one request moves the mean less; the 16-engine runs this was tuned on ran at 1.0.

Only applies when enable_prefix_caching is True and using a coordinator.

prefix_cache_ttl_seconds: float

300.0

How long the coordinator assumes an engine still holds a block it routed there.

The coordinator sees blocks being routed but never blocks being evicted, so its view of each engine’s cache only gets staler. Entries untouched for this long are dropped. Too long and it claims hits on blocks already evicted, routing for affinity and paying a cold prefill anyway; too short and it forgets blocks the engine still holds.

media_cache_coordinator_policy: core.inference.config.MediaCacheCoordinatorPolicy

None

Media-cache routing policy for the DP inference coordinator.

Media affinity is active only when vision_embedding_cache_max_bytes is greater than zero. Media-salted prefix affinity is controlled separately by prefix_caching_coordinator_policy.

media_cache_routing_weight: float

1.0

Estimated vision-encoder reuse cost in compact-prompt block units.

Multimodal coordinator routing combines this media-hit value with the number of matching routing-prefix blocks before blending cache affinity with load using prefix_caching_routing_alpha. The engine independently uses post-expansion hashes for authoritative KV lookup. Must be non-negative.

prefix_caching_mamba_gb: Optional[float]

None

GPU memory budget (in GB) for the Mamba state cache used by prefix caching on hybrid models. Each cache slot stores SSM and conv states for all Mamba layers at a single block boundary. When set, Mamba states at KV divergence and last-aligned block boundaries are cached and reused across requests with matching prefixes.

This budget covers both buffers allocated by MambaSlotAllocator: the durable cache (ssm_states/conv_states, max_slots slots reused across requests) and the per-step extraction scratch (intermediate_ssm_out/intermediate_conv_out). The scratch is sized to the tighter of two per-step bounds, min(ceil(max_tokens / block_size_tokens), 3 * max_requests), since a single engine step can extract at most one state per block_size_tokens of its token budget (and at most 3 per request). The scratch is reserved from this budget first, so a smaller max_tokens (or max_requests) shrinks the scratch and leaves more durable cache slots.

track_paused_request_events: bool

False

Whether to track paused request events. If True, add_event_pause() is called on requests when they are paused during bookkeeping.

track_generated_token_events: bool

False

Whether to track per-token events with timestamps for each generated token. When enabled, each generated token creates a GENERATED_TOKEN event with a timestamp, useful for per-token latency analysis.

metrics_writer: Optional[WandbModule]

None

Wandb module for writing metrics.

logging_step_interval: int

0

The step interval at which to log inference metrics to wandb. Defaults to 0, which means no logging.

sampling_backend: Literal[torch, flashinfer]

‘torch’

Which sampling kernels to use during inference. Falls back to “torch” with a warning if “flashinfer” is requested but the package is not installed.

offset_sampling_seed_by_dp_rank: bool

True

If True, offset inference_sampling_seed by the data-parallel rank when seeding the sampling RNG. This gives each DP rank a unique generation seed so that the same prompt routed to different ranks produces different samples (important for RL training). If False (or ModelParallelConfig.deterministic_mode / --deterministic-mode is enabled), then all DP ranks share the same sampling / generation seed.

async_sched_mode: core.inference.config.AsyncScheduleMode

None

Mode used to schedule dynamic batching inference work. Defaults to async scheduling; use AsyncScheduleMode.LEGACY to disable it.

logprobs_mode: Literal[raw_logprobs, processed_logprobs]

‘raw_logprobs’

Whether returned log-probs are modified by the sampling parameters or not.

request_metadata_types: Optional[List[Tuple[str, torch.dtype]]]

None

A list of the per-request metadata types to track. Each entry is a tuple consisting of the string label and the target dtype.

use_synchronous_zmq_collectives: bool

False

Whether to use synchronous ZMQ collectives for inference. If True, the all_reduce_max operation will be performed synchronously, which can help reduce performance variability for MoEs.

disable_ep_consensus: bool

False

If True, the engine skips the EP-group consensus all-reduce in run_engine_with_coordinator and decides whether to step based on local state alone. The rank still calls controller.dummy_forward() whenever local_pending == 0, so EP collectives (NCCL all-to-all, etc.) stay in sync — without this, a peer running a real forward would deadlock waiting on this rank’s all-to-all participation. Trades off the consensus all-reduce CPU cost for unconditional dummy_forwards on idle ranks.

ep_consensus_interval: int

20

How many steps to skip between EP-consensus all-reduces when the engine has pending work. Consensus is always run immediately when there is no global work (to detect new arrivals quickly); this interval only applies to the busy case, where skipping avoids per-step all-reduce overhead. In the worst case, pausing is delayed by this many steps (~10–20 ms per step at typical decode throughput).

verbose: dataclasses.InitVar[bool]

False

Whether to log detailed context configuration at initialization. This is an InitVar and is not stored as a field on the config.

allow_stale_multimodal_embeddings: bool

False

Allow projected-media embeddings to survive weight-change boundaries.

By default, suspend/resume and generation-epoch changes invalidate both the shared vision-embedding cache and request-local vision state. Enable this only when model weights are guaranteed not to change across those boundaries.

__post_init__(verbose: bool)