core.inference.engine_factory

Construction helpers for dynamic inference engines.

Module Contents

Functions

build_dynamic_inference_engine

Build and configure a dynamic inference engine.

API

core.inference.engine_factory.build_dynamic_inference_engine(

*,

model: Any,

tokenizer: Any,

inference_config: megatron.core.inference.config.InferenceConfig | None = None,

inference_wrapper_cls: type[megatron.core.inference.model_inference_wrappers.abstract_model_inference_wrapper.AbstractModelInferenceWrapper] = GPTInferenceWrapper,

engine_cls: type[megatron.core.inference.engines.dynamic_engine.DynamicInferenceEngine] | None = None,

) → megatron.core.inference.engines.dynamic_engine.DynamicInferenceEngine

Build and configure a dynamic inference engine.

This is the shared construction path for the high-level inference API and framework integrations. When disaggregation_shards is configured, it selects the disaggregated engine, reserves any recurrent-state dummy slot, and applies the prefill/decode coordinator configuration.

Parameters:

• model – Distributed Megatron model to serve.

• tokenizer – Tokenizer consumed by the text generation controller.

• inference_config –

  Runtime inference configuration. Defaults to

  class:

  InferenceConfig.

• inference_wrapper_cls – Model inference wrapper implementation.

• engine_cls – Optional engine implementation override. A disaggregated config requires a subclass of :class:DisaggDynamicInferenceEngine.

Returns:

A fully configured dynamic inference engine.

Raises:

ValueError – If a disaggregated config disables prefix caching or is paired with an incompatible explicit engine class.