Model pipeline

Neural network modules​

MAX provides two neural network module packages:

• max.nn provides reusable layers that build MAX graphs. Most existing pipeline architectures use this package.
• max.experimental.nn provides the module API used by newer ModuleV3 architectures.

The max.nn package includes common components like:

• Core modules: Module, Linear, Embedding, Sequential, LayerList.
• Normalization: RMSNorm and LayerNorm.
• Positional encodings: RotaryEmbedding.

For graph-based components like attention mechanisms and transformers, import from max.nn:

from max.nn import Module, Linear
from max.nn.attention import AttentionWithRope

The Module base class in max.experimental.nn standardizes how modules manage weights and devices. Here's an example of building a simple multi-layer perceptron with the ModuleV3 API:

from max.driver import Accelerator
from max.experimental.nn import Linear, Module, module_dataclass

@module_dataclass
class MLP(Module):
    fc1: Linear
    fc2: Linear

    def forward(self, x):
        return self.fc2(self.fc1(x))

# Create a model with two linear layers
model = MLP(fc1=Linear(10, 20), fc2=Linear(20, 5))

# Weights are tracked automatically through the module hierarchy
for name, param in model.parameters:
    print(f"{name}: {param.shape}")
# fc1.weight: [20, 10]
# fc1.bias: [20]
# fc2.weight: [5, 20]
# fc2.bias: [5]

# Move all parameters to an accelerator (GPU)
model.to(Accelerator())

In this example, the Module base class automatically tracks all parameters through the module hierarchy, letting you iterate over them or inspect them. The to() method provides a simple way to move the entire model and all its parameters to a different device with a single call.

KV cache module​

The max.pipelines.kv_cache package provides cache management for transformer inference. The main component is PagedKVCacheManager, which handles memory allocation for key-value pairs across generation steps.

For most use cases, you don't interact with the cache manager directly. The pipeline handles cache management automatically using paged attention based on the supported_encodings in your architecture config.

How modules and pipelines work together​

Understanding the relationship between modules and pipelines is key to working with MAX. When you build a model for the pipeline system, define your architecture as a module, using either the graph API (max.nn.Module) or the ModuleV3 API (max.experimental.nn.Module). MAX then compiles your module into an optimized executable for the target device, and the pipeline orchestrates execution.

Here's the workflow:

1. Define your model: Create a module that defines your model architecture. The module's computation method defines what happens when processing inputs: __call__() for the graph API, or forward() for the ModuleV3 API.

2. MAX compiles your model: MAX compiles the module into an optimized executable for the target device. This compilation happens once, and you can reuse the result for many inference calls.

3. Pipelines orchestrate execution: The pipeline receives pre-tokenized context objects (which the tokenizer creates), manages the KV cache, calls your model's PipelineModel.execute() method, samples output tokens, and returns results.

This separation lets you work at the right level of abstraction: use max.nn or max.experimental.nn to define model architectures, let MAX handle compilation, and rely on pipelines for production serving.

Request execution example​

To see how these components work together in practice, here's a complete example of generating text with a pipeline.

The PIPELINE_REGISTRY is the central system that maps model architectures to their compiled pipelines and tokenizers. When you call retrieve() with a PipelineConfig, the registry looks up the model's architecture from its Hugging Face config, validates compatibility, downloads weights if needed, compiles the model, and returns both a tokenizer and a ready-to-use pipeline instance.

import asyncio
from max.pipelines.modeling.types import RequestID, TextGenerationInputs, TextGenerationRequest
from max.pipelines import PIPELINE_REGISTRY, PipelineArgs, PipelineConfig
from max.pipelines.context import TextContext

# 1. Configure and retrieve the pipeline and tokenizer
config = PipelineConfig.from_args(
    PipelineArgs.from_flat_kwargs(
        model_path="meta-llama/Llama-3.1-8B-Instruct",
        max_length=512,
    )
)
tokenizer, pipeline = PIPELINE_REGISTRY.retrieve(config)

# 2. Get the KV cache manager from the pipeline
kv_cache_manager = pipeline.kv_manager

# 3. Create a request with the text prompt
request = TextGenerationRequest(
    request_id=RequestID(),
    model_name=config.model.model_path,
    prompt="Explain how neural networks work",
)

# 4. Create a context object (tokenization happens here)
context = asyncio.run(tokenizer.new_context(request))

# 5. Allocate space in the KV cache for this request
kv_cache_manager.claim(context)

# 6. Run the generation loop
generated_text = ""
while True:
    # Allocate KV cache for the next token
    kv_cache_manager.alloc(context)

    # Execute the pipeline with the current context
    inputs = TextGenerationInputs[TextContext](
        batches=[[context]],
    )
    output = pipeline.execute(inputs)

    # Decode and accumulate generated tokens
    for token in output[context.request_id].tokens:
        generated_text += asyncio.run(
            tokenizer.decode(token, skip_special_tokens=True)
        )

    # Check if generation is complete
    if output[context.request_id].is_done:
        break

print(generated_text)

In this example, you can see the key phases of pipeline execution:

1. The registry maps the model path to the appropriate tokenizer and compiled pipeline.
2. The cache manager tracks memory allocation for the request's key-value pairs across generation steps.
3. The new_context() method handles tokenization internally and creates a context object that tracks the request's state throughout generation.
4. The pipeline processes tokens, the model executes, and the sampler selects new tokens until completion.

Notice how the pipeline itself is stateless, all request-specific state lives in the context object and the KV cache manager. The pipeline orchestrates execution based on the inputs it receives.

For more information on the stateless nature of the pipeline system, see Stateless orchestration below.

Top-level interfaces​

The max.pipelines.modeling.types package defines the contracts that all pipeline components must implement. These abstractions enable MAX to work uniformly across different model architectures and tasks.

The key interfaces are:

• Pipeline: Abstract base class for all pipelines. Defines execute() and release() methods that all pipeline implementations must provide.

• PipelineInputs: Base class for inputs to a pipeline, such as text generation requests or embeddings requests.

• PipelineOutput: Protocol for pipeline outputs. Must implement is_done to signal when generation is complete.

• PipelineTokenizer: Interface for tokenizers that convert between text and token IDs, and create context objects for requests.

• PipelineModel: Abstract base class for model implementations. Defines methods like execute() and input preparation methods that all architectures must implement.

These interfaces are task-agnostic. Specialized variants like TextGenerationInputs and TextGenerationOutput extend them for specific use cases.

Pipeline registry​

The PIPELINE_REGISTRY is a singleton that tracks all available model architectures. When you run the max serve command with a model, the registry:

1. Looks up the model's architecture from its Hugging Face config.
2. Validates that it supports the requested encoding and settings.
3. Returns the appropriate tokenizer and pipeline for that architecture.

You can interact with the registry directly to retrieve a model's tokenizer and compiled pipeline:

from max.pipelines import PIPELINE_REGISTRY, PipelineArgs, PipelineConfig

# Create configuration for a model
config = PipelineConfig.from_args(
    PipelineArgs.from_flat_kwargs(
        model_path="meta-llama/Llama-3.1-8B-Instruct",
    )
)

# Retrieve tokenizer and compiled pipeline
tokenizer, pipeline = PIPELINE_REGISTRY.retrieve(config)

# Or get a factory for deferred compilation
retrieved = PIPELINE_REGISTRY.retrieve_factory(config)
pipeline = retrieved.factory()  # Compile when ready

In this example, retrieve() returns a ready-to-use pipeline, while retrieve_factory() returns a RetrievedPipeline with the tokenizer, a factory that performs compilation when invoked, and the memory plan the pipeline was sized against. The factory pattern is useful when you need to pass the pipeline across process boundaries, since it avoids serializing the compiled model.

Supported architecture​

A SupportedArchitecture configuration defines each model architecture. This bridges the gap between Hugging Face model conventions and MAX's execution system.

When you point MAX at a Hugging Face model (like meta-llama/Llama-3.1-8B-Instruct), MAX downloads and reads the model's config.json file. Inside that file is an architectures field listing the model class name (like "LlamaForCausalLM"). The registry uses this name to look up the corresponding SupportedArchitecture, which tells MAX which model (which subclass of PipelineModel) to use, what quantization formats it supports, how to load weights, and which tokenizer to instantiate.

Here's how you define an architecture:

from max.graph.weights import WeightsFormat
from max.pipelines.modeling.types import PipelineTask
from max.pipelines.lib import (
    SupportedArchitecture,
    TextTokenizer,
)

llama_arch = SupportedArchitecture(
    # Must match the HuggingFace model class name
    name="LlamaForCausalLM",

    # The type of task this architecture supports
    task=PipelineTask.TEXT_GENERATION,

    # Example models that use this architecture
    example_repo_ids=[
        "meta-llama/Llama-3.1-8B-Instruct",
        "meta-llama/Llama-3.2-3B-Instruct",
    ],

    # Quantization support
    default_encoding="q4_k",
    supported_encodings={
        "bfloat16",
        "q4_k",
    },

    # Implementation classes
    pipeline_model=Llama3Model,
    tokenizer=TextTokenizer,

    # Weight handling
    default_weights_format=WeightsFormat.safetensors,
    weight_adapters={
        WeightsFormat.safetensors: convert_safetensor_state_dict,
        WeightsFormat.gguf: convert_gguf_state_dict,
    },

    # Architecture-specific settings
    multi_gpu_supported=True,
)

The name field must match the architectures field in the model's Hugging Face config.json. Common architecture names include LlamaForCausalLM, DeepseekV3ForCausalLM, Qwen3VLMoeForConditionalGeneration, and more. However, if you are using a custom architecture, you will need to use a custom name that is specific to your architecture.

The task field determines which Pipeline subclass MAX uses to orchestrate execution. Different types of models serve different purposes, and each task type has its own execution strategy:

• TEXT_GENERATION: Autoregressive text generation for chat and completion use cases. MAX uses TextGenerationPipeline to handle the prefill and decode loop, KV cache management, and token sampling.
• EMBEDDINGS_GENERATION: Vector embeddings from text input. MAX uses EmbeddingsPipeline to produce dense vector representations suitable for semantic search and retrieval.

For example, if you are building a model for text generation, you will set task=PipelineTask.TEXT_GENERATION.

Set the default_encoding to the quantization encoding you want to use by default, and supported_encodings to the set of encodings your architecture supports. For example, if your model supports q4_k quantization, set supported_encodings={"q4_k"}. For more information on quantization, see the quantization guide.

The pipeline_model field is the class that builds and executes the model. The tokenizer field is the class that handles text encoding and decoding.

Finally, the weight_adapters field is a dictionary of functions that convert weights from different formats. You can use the weight_adapters field to convert weights from different formats to the default format.

You use these fields to configure the pipeline for your model. In most cases, if you are implementing a new model architecture, you will have intimate knowledge of the model's architecture and you will be able to set these fields accordingly.

Pipeline model​

The PipelineModel abstract class defines the interface for model implementations. Every model architecture implements execute(), and either implements prepare_initial_token_inputs() directly or delegates input preparation to a batch processor (see Delegating input preparation to a batch processor):

from max.pipelines.lib import PipelineModel, ModelInputs, ModelOutputs

class MyModel(PipelineModel):
    # Link to your architecture config class for max_seq_len and KV params
    model_config_cls = MyModelConfig

    def execute(self, model_inputs: ModelInputs) -> ModelOutputs:
        """Run inference on the compiled model."""
        ...

    def prepare_initial_token_inputs(
        self, context_batch, kv_cache_inputs=None, return_n_logits=1
    ) -> ModelInputs:
        """Prepare inputs for the first forward pass (prefill)."""
        ...

The execute() method runs your model's forward pass on the compiled executable. prepare_initial_token_inputs() builds the tensors for each scheduled batch (prefill and decode steps share this path in the serving pipeline):

1. Prefill: Process the entire prompt at once, building up the KV cache.
2. Decode: Generate tokens one at a time, reusing the cached keys/values.

When to customize ModelInputs​

For most text-to-text transformer models, you can use the default ModelInputs implementation provided by MAX. Start with the default ModelInputs and only create a custom subclass if your model's forward pass requires additional tensors beyond the standard transformer inputs. Custom ModelInputs are only necessary for:

• Models that process both text and image inputs, require custom input structures to handle image tensors, pixel values, or image embeddings alongside text tokens.

• Models that have unique input patterns, for example: mixture-of-experts with routing tensors or retrieval augmented models with document embeddings.

If you're implementing a standard decoder-only language model (like Llama, Mistral, or similar architectures), you likely don't need to subclass ModelInputs. The default implementation handles token IDs, position IDs, attention masks, and KV cache inputs, which covers most use cases.

Stateless orchestration​

A core design principle of the MAX pipeline system is that pipelines are stateless orchestrators. The pipeline itself does not own or maintain per-request state. Instead, it operates on the state passed to it through inputs:

• Context objects track all request-specific information (tokens, sampling parameters, generation status). You pass these into execute(), and the pipeline updates them but doesn't store them internally.

• KV cache manager owns the allocation and lifecycle of cached key-value pairs across all requests. The pipeline uses the cache manager but doesn't own it.