IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /max/get-started.md). For the complete documentation index, see llms.txt.
Skip to main content
For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/get-started.md).

Model pipeline

Implementing a model architecture in MAX involves more than just defining the layers and the forward pass. To serve your model, you must also integrate it into a pipeline that manages the entire inference lifecycle.

A pipeline connects a model implementation to MAX's serving system. It's responsible for tasks such as preparing inputs, executing the model, managing inference state such as the KV cache, performing post-processing, and returning results.

This page explains how the pipelines API works, how pipeline components work together, and when to customize pipelines.

The pipeline system​

The pipeline system uses a registry pattern where model architectures register their capabilities, and the infrastructure handles the execution details. When you point MAX at a model, the registry looks up the architecture, validates compatibility, downloads weights, compiles the model, and returns a ready-to-use pipeline.

This separates concerns cleanly:

  • Modules define model architectures and hold weights.
  • Pipelines orchestrate the inference loop and manage state.
  • Registry maps model identifiers to implementations.
  • Compilation transforms your model into optimized executables for the target device.

Pipelines let you focus on model architecture while MAX handles the production infrastructure, including batching, caching, compilation, and serving.

Building blocks​

Before diving into pipeline components, it's helpful to understand the two foundational packages that pipelines build on: max.nn and max.pipelines.kv_cache.

Neural network module​

The max.nn (neural network) package provides reusable neural network layers that serve as the bridge between the MAX Graph API and model implementations.

The max.nn package includes common components like:

These components are core to building model architectures.

The Module base class standardizes how layers manage weights and devices. Here's an example of building a simple multi-layer perceptron:

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

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, you define your architecture using the Module class. 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: You create a Module that defines your model architecture. The module's forward method defines what computations happen when processing inputs.

  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 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, PipelineConfig
from max.pipelines.context import TextContext

# 1. Configure and retrieve the pipeline and tokenizer
config = PipelineConfig.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.request_id, replica_idx=0)

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

    # 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.

Core components​

Now that you understand how modules and pipelines work together, you can explore the specific components that make up the pipeline system.

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, PipelineConfig

# Create configuration for a model
config = PipelineConfig.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
tokenizer, pipeline_factory = PIPELINE_REGISTRY.retrieve_factory(config)
pipeline = pipeline_factory()  # Compile when ready

In this example, retrieve() returns a ready-to-use pipeline, while retrieve_factory() returns a callable that performs compilation when invoked. 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
    rope_type="normal",
    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 format you want to use by default. supported_encodings maps quantization formats to compatible KV cache strategies. For example, if you are using a model that supports q4_k quantization, you will 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.

Delegating input preparation to a batch processor​

You can move the batching logic into a BatchProcessor instead of overriding prepare_initial_token_inputs() on the model, and register it through the batching field of SupportedArchitecture. At registration, MAX sets the processor as the model's batch_processor_cls, and the base PipelineModel.prepare_initial_token_inputs() delegates to it. For ragged batches, subclass RaggedBatchProcessor, which receives a BatchProcessorRuntime describing the active devices and KV cache.

Use a batch processor when several architectures share the same batching logic, or when you want that logic isolated from the model's graph construction. For example, one Llama 3 batch processor backs Llama 3, OLMo, Granite, Qwen 2, and Phi-3; see the llama3 architecture for a reference implementation. Override prepare_initial_token_inputs() on the model instead when the batching is specific to one architecture and tightly coupled to how it builds inputs.

Linking to the architecture config​

The model_config_cls class variable links your pipeline model to the architecture configuration class. This connection enables the base PipelineModel to delegate key behaviors to the config.

The base class delegates the following through model_config_cls:

  • Maximum sequence length: The calculate_max_seq_len() method determines the model's sequence length limit. The base class calls model_config_cls.calculate_max_seq_len() or initializes the config and calls get_max_seq_len().

  • KV cache parameters: The get_kv_params() method constructs the KV cache configuration. The base class calls model_config_cls.construct_kv_params() to build the appropriate KVCacheParams for your architecture. Architectures with multiple or hybrid caches (for example, separate sliding-window and full-attention caches) return a MultiKVCacheParams tree instead, which the cache manager consumes through the shared KVCacheParamInterface.

You can override calculate_max_seq_len() or get_kv_params() if needed.

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.

Pipeline execution​

Now let's look at how pipelines coordinate execution at runtime. When a pipeline runs, it orchestrates three main components: the KV cache manager (which tracks key-value pairs across generation steps), the model (which executes the forward pass), and the sampler (which selects the next token based on the model's output logits). The execute() method ties these together in a generation loop.

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.

Configuration flow​

Now let's see how configuration flows from user input to a running pipeline. When you run a model, configuration flows through several layers:

  1. Start with user arguments: MAX collects CLI or API arguments into a PipelineConfig object that specifies the model path, quantization settings, and runtime parameters.

  2. Load model metadata: The registry fetches the model's Hugging Face config to perform architecture lookup and extract hyperparameters like hidden size, number of layers, and vocabulary size.

  3. Validate compatibility: The registry checks that the architecture supports the requested encoding and KV cache strategy.

  4. Instantiate pipeline: Finally, the registry constructs and returns the tokenizer and compiled pipeline ready for inference.

from max.pipelines import PipelineConfig

config = PipelineConfig.from_flat_kwargs(
    # Model specification (Hugging Face repo ID or local path)
    model_path="meta-llama/Llama-3.1-8B-Instruct",

    # Sequence limits
    max_length=4096,

    # Batching
    max_batch_size=32,
)

The PipelineConfig consolidates all settings and provides defaults based on the model and hardware. See the PipelineConfig reference for all available options.

Next steps​

Now that you understand the pipeline architecture, continue learning:

Edit this page

Was this page helpful?

Models

Cookbook

Puzzles

Blog

Community

Contact