For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Python class

PipelineModel

PipelineModel

class max.pipelines.PipelineModel(pipeline_config, session, devices, kv_cache_config, weights, adapter, return_logits, *, memory_plan, return_hidden_states=ReturnHiddenStates.NONE, max_batch_size=1)

source

Bases: ABC, Generic[BaseContextType]

A pipeline model with setup, input preparation and execution methods.

Parameters:

• pipeline_config (PipelineConfig)
• session (InferenceSession)
• devices (list[Device])
• kv_cache_config (KVCacheConfig)
• weights (Weights)
• adapter (WeightsAdapter | None)
• return_logits (ReturnLogits)
• memory_plan (MemoryPlan)
• return_hidden_states (ReturnHiddenStates)
• max_batch_size (int)

arch_config_as()

arch_config_as(cls)

source

Returns the built-once arch config, narrowed to cls.

Parameters:

cls (type[ArchConfigT])

Return type:

ArchConfigT

batch_processor

property batch_processor: BatchProcessor[Any, Any] | None

source

Returns the batch processor when configured.

batch_processor_cls

batch_processor_cls: ClassVar[type[BatchProcessor[Any, Any]] | None] = None

source

Optional batch processor class for input/output handling.

compute_log_probabilities()

compute_log_probabilities(session, model_inputs, model_outputs, next_tokens, batch_top_n, batch_echo)

source

Optional method that can be overridden to compute log probabilities.

Parameters:

• session (InferenceSession) – Inference session to compute log probabilities within.
• model_inputs (ModelInputs) – Inputs to the model returned by prepare_initial_token_inputs().
• model_outputs (ModelOutputs) – Outputs returned by execute().
• next_tokens (Buffer) – Sampled tokens. Should have shape=[batch size]
• batch_top_n (list[int]) – Number of top log probabilities to return per input in the batch. For any element where top_n == 0, the LogProbabilities is skipped.
• batch_echo (list[bool]) – Whether to include input tokens in the returned log probabilities.

Returns:

List of log probabilities.

Return type:

list[LogProbabilities | None]

dtype

property dtype: DType

source

Returns the model data type.

execute()

abstract execute(model_inputs)

source

Executes the graph with the given inputs.

Parameters:

model_inputs (ModelInputs) – The model inputs to execute, containing tensors and any other required data for model execution.

Returns:

ModelOutputs containing the pipeline’s output tensors.

Return type:

ModelOutputs

This is an abstract method that must be implemented by concrete PipelineModels to define their specific execution logic.

huggingface_config

property huggingface_config: AutoConfig

source

Returns the HuggingFace config from pipeline config.

For multimodal models (e.g., Pixtral, Gemma3 multimodal), this returns the top-level config which contains both text_config and vision_config. Models should explicitly access .text_config or .vision_config as needed.

Returns:

The HuggingFace AutoConfig for this model.

Raises:

ValueError – If HuggingFace config could not be loaded.

lora_manager

property lora_manager: LoRAManagerV3 | None

source

Returns the LoRA manager if LoRA is enabled, otherwise None.

lora_modulev3

lora_modulev3: ClassVar[bool] = False

source

Whether this arch serves LoRA via the ModuleV3 adapters-as-inputs path (LoRAManagerV3). Non-ModuleV3 archs cannot serve LoRA.

lora_targets

lora_targets: ClassVar[tuple[LoRATargetModule, ...]] = ()

source

The ModuleV3 LoRA target projections this arch wraps. Read by the base to construct LoRAManagerV3; empty for non-ModuleV3-LoRA archs.

max_seq_len

property max_seq_len: int

source

The effective maximum sequence length, read from the memory plan.

A view of the plan’s planned_max_length — the model stores no copy of it.

model_config_cls

model_config_cls: ClassVar[type[Any] | None] = None

source

Config class used to build the arch config and delegate KV params.

planned_max_batch_total_tokens

property planned_max_batch_total_tokens: int | None

source

The plan’s batch token budget.

prepare_initial_token_inputs()

prepare_initial_token_inputs(replica_batches, kv_cache_inputs=None, return_n_logits=1)

source

Prepares the initial inputs to be passed to execute().

The inputs and functionality can vary per model. For example, model inputs could include encoded tensors, unique IDs per tensor when using a KV cache manager, and kv_cache_inputs (or None if the model does not use KV cache). This method typically batches encoded tensors, claims a KV cache slot if needed, and returns the inputs and caches.

When batch_processor_cls is set, delegates to the batch processor.

Parameters:

• replica_batches (Sequence[Sequence[BaseContextType]])
• kv_cache_inputs (KVCacheInputsInterface[Buffer, Buffer] | None)
• return_n_logits (int)

Return type:

ModelInputs

sampler_custom_extensions

property sampler_custom_extensions: Sequence[Path]

source

Custom-op extension paths to compile the sampler graph with.

signal_buffers

property signal_buffers: list[Buffer]

source

Lazily initialize signal buffers for multi-GPU communication collectives.

Signal buffers are only needed during model execution, not during compilation. By deferring their allocation, we avoid memory allocation in compile-only mode.

Returns:

List of signal buffer tensors, one per device for multi-device setups, or an empty list for single-device setups or compile-only mode.