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

PipelinedModel

PipelinedModel

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

source

Bases: PipelineModel[BaseContextType, Any, Any]

Mixin for a pipeline model whose decoder can run as pipeline stages.

Every stage’s graph takes the whole model’s graph inputs, so the engine that serves the stages prepares each stage’s inputs as it would the whole model’s. A stage after the first takes what the previous stage’s decoder returned after them.

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)

stage_graph()

stage_graph(name, input_types, **graph_kwargs)

source

Builds this model’s graph, or its stage’s part of it.

Yields the graph, the values of input_types and the boundary inputs. A stage after the first takes what the previous stage’s decoder returned as its last graph inputs, typed by its stage’s boundary_types. The boundary inputs are those values in the tree the decoder returned them, or None on the first stage. The model calls its decoder with the boundary inputs and outputs what stage_outputs() returns. Without a pipeline stage this is the plain graph.

Parameters:

• name (str) – The graph’s name.
• input_types (Iterable[Type[Any]]) – The model’s graph input types, the same on every stage.
• graph_kwargs (Any) – Passed to Graph.

Raises:

ValueError – If a stage after the first has no boundary types.

Return type:

Iterator[tuple[Graph, Sequence[Value[Any]], Any]]

stage_outputs()

stage_outputs(outputs)

source

Returns what a stage graph outputs, from what its decoder returned.

On the last stage, and without a pipeline stage, that is outputs. A stage before the last outputs the leaves of what its decoder returned.

Parameters:

outputs (Any) – What the decoder returned.

Raises:

• StageTraced – On a stage that only traces, carrying the types of outputs.
• ValueError – If a tensor crossing the cut has an algebraic dim.

Return type:

Sequence[Value[Any]]