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 module

max.pipelines.lib.module_registration

Serves architectures that register their modules by role.

An architecture that sets modules writes no pipeline model. module_pipeline_model() builds one from the registration, and it owns what a per-architecture pipeline model used to: finalizing the config, constructing each module as module_cls(config), the graph signature, batching and execution.

A registered module class can also set:

• graph_name, the name its graph compiles under (graph_name());
• input_layout, the role field names in a graph’s existing input order. Naming signals places the allreduce signal buffers there, on one device too, as the graph-API graph that keeps the layout takes them;
• device_context_length_dtype, when its graph takes batch_context_lengths as a graph-API graph does. Setting it changes both their count and their dtype: one per device, each replica’s on every device of the replica, of that dtype, rather than one int32 per replica;
• lora_targets, the projections it serves LoRA adapters on (lora_targets());
• supports_data_parallel, when it runs one replica per device (supports_data_parallel());
• supports_multi_device_replicas, when each of its data-parallel replicas may also span several devices (supports_multi_device_replicas());
• consumes_vision, when it reads vision (consumes_vision());
• supports_pipeline_stages, when its config builds one pipeline stage’s layers (supports_pipeline_stages());
• scratch_types(config, *, max_batch_size, devices), a classmethod that returns the buffer types of the workspace its kernels write each step (see decoder_scratch). Its batch processor must be a DecoderBatchProcessor, which passes the buffers;
• state_store(config, *, max_batch_size, devices), a classmethod that returns the store of the per-request state it carries between steps (see request_state). Its batch processor must be a DecoderBatchProcessor, which feeds the state.

A registration can add a VISION module beside the DECODER. The pipeline model then compiles the vision graph with the decoder’s, and the pipeline encodes each batch’s new images through its vision cache and passes the embeddings to the decoder’s vision input. The registered batch processor packs the images and runs the vision graph, since pixel layouts differ by architecture: it implements VisionBatchProcessor. A registered VISION class constructs from the config like any role, and also has input_types(), the vision graph’s inputs, and embedding_dtype and embedding_dim, the type of each embedding row it produces. Registration checks the combination up front: a VISION module needs a decoder that consumes the embeddings and a batch processor that prepares them, and a decoder that consumes them needs a VISION module.

registration_runtime lists what a registered config can declare about its deployment.

Pipeline models

ModuleV3DecoderPipelineModel	Serves a registered text-generation DECODER, and its VISION.
RegisteredDecoderBuilder	Constructs a registered DECODER for a caller that composes modules.

module_pipeline_model	Returns the pipeline model that serves a module registration.

Batch processors

DecoderBatchProcessor	Batches ragged text for a registered DECODER.
VisionBatchProcessor	The batch processor of a registration with a VISION module.

Model inputs

DecoderModelInputs	The prepared inputs of a registered DECODER.

Module construction

construct_module	Constructs a registered module as module_cls(config) on mesh.

Registration accessors

device_mesh	Returns the mesh registered modules are constructed on.
graph_name	Returns the name a registered module's graph compiles under.
supports_pipeline_stages	Returns whether a registered decoder builds as pipeline stages.

Graph input types

decoder_input_types	Returns a DECODER graph's input types, keyed by field.
decoder_vision_input_types	Returns the decoder's vision input types, replicated on mesh.

Workspace and state binding

bind_decoder_scratch	Allocates a registered decoder's workspace and gives it to the batcher.
bind_state_store	Builds a registered decoder's state store and gives it to the batcher.