Is this page helpful?

core.inference.apis.llm#

Sync high-level inference API for Megatron (MegatronLLM).

Module Contents#

Classes#

MegatronLLM

Sync high-level inference API for Megatron.

API#

class core.inference.apis.llm.MegatronLLM(
*,
model,
tokenizer,
inference_config: Optional[megatron.core.inference.config.InferenceConfig] = None,
use_coordinator: bool = True,
coordinator_host: Optional[str] = None,
coordinator_port: Optional[int] = None,
)#

Bases: megatron.core.inference.apis._llm_base._MegatronLLMBase

Sync high-level inference API for Megatron.

See :class:_MegatronLLMBase for execution modes (direct vs coordinator), caller responsibilities, and the model.eval() contract.

On top of the base, this class provides:

  • meth:

    generate accepting one prompt or a batch; always returns a list[DynamicInferenceRequest] (single-prompt input returns a one-element list – deliberate asymmetry vs the async API).

  • Sync lifecycle controls: :meth:pause / :meth:unpause /

    meth:

    suspend / :meth:resume / :meth:shutdown /

    meth:

    wait_for_shutdown.

  • meth:

    serve for OpenAI-compatible HTTP serving on the primary rank.

  • Context-manager protocol: with MegatronLLM(...) as llm:; exit calls :meth:shutdown.

Initialization

generate(
prompts: Union[str, List[int], List[str], List[List[int]]],
sampling_params: Optional[megatron.core.inference.sampling_params.SamplingParams] = None,
) → List[megatron.core.inference.inference_request.DynamicInferenceRequest]#

Run inference for one prompt or a batch.

Returns list[DynamicInferenceRequest] in input order. Single-prompt input returns a one-element list – the always-list shape is the deliberate sync-vs-async asymmetry.

No concurrency guard: sync is single-caller by Python’s GIL. If you need to call generate concurrently from multiple threads, callers must serialize externally.

Raises:

RuntimeError – if called on a non-primary rank in coordinator mode.

pause() → None#

Transition the engine to PAUSED. Coordinator mode only.

Raises:

RuntimeError – in direct mode (use_coordinator=False).

unpause() → None#

Transition the engine from PAUSED back to RUNNING.

Raises:

RuntimeError – in direct mode (use_coordinator=False).

suspend() → None#

Transition the engine to SUSPENDED (offloads GPU buffers).

The caller must pause() first; this method does not enforce that.

Raises:

RuntimeError – in direct mode (use_coordinator=False).

resume() → None#

Transition the engine from SUSPENDED to RESUMED.

Raises:

RuntimeError – in direct mode (use_coordinator=False).

shutdown() → None#

Tear down the engine and runtime. Idempotent. Direct mode is a no-op.

serve(
serve_config: megatron.core.inference.apis.serve_config.ServeConfig,
*,
blocking: bool = True,
) → None#

Start the OpenAI-compatible HTTP frontend.

Coordinator mode only. The HTTP frontend runs only on the primary rank (global rank 0); other ranks no-op the HTTP setup but still respect blocking (so all ranks return together).

With blocking=True (default), this blocks the calling thread until the engine loop terminates via :meth:shutdown – suitable for standalone serving scripts. With blocking=False, this returns once the HTTP frontend is up (primary) or immediately (workers); the engine loop continues in the background runtime, and the user can call

Meth:

generate / :meth:shutdown afterward.

Raises:

ValueError – if use_coordinator=False (HTTP serving requires the coordinator path).

wait_for_shutdown() → None#

Block until the engine loop terminates. Direct mode no-op.

__enter__() → core.inference.apis.llm.MegatronLLM#
__exit__(exc_type, exc, tb) → None#