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).
Serve your custom model
You can serve a model with a custom architecture using the same max serve command that runs MAX-supported
models. Once max serve loads your architecture package, your
model gets batching, KV cache management, tokenization, and the rest of the MAX
serving stack on top of it. This page teaches you how to make your custom
architecture package compatible with max serve.
How MAX registers custom architectures
A custom model architecture package is a Python package (a directory containing
an __init__.py) that MAX can import. To serve the model, the __init__.py
must expose an ARCHITECTURES list of
SupportedArchitecture
instances. Each SupportedArchitecture
instance bundles the model class, config, tokenizer, encodings, and weight
adapters under a name.
A minimal __init__.py looks like this:
from .arch import my_arch
ARCHITECTURES = [my_arch]A typical arch.py builds a SupportedArchitecture from
max.pipelines.lib with types from
max.pipelines.modeling.types:
from max.graph.weights import WeightsFormat
from max.pipelines.architectures.llama3_modulev3.batch_processor import (
Llama3ModuleV3BatchProcessor,
)
from max.pipelines.context import TextContext
from max.pipelines.kv_cache.memory_planner import PagedMemoryPlanner
from max.pipelines.lib import SupportedArchitecture, TextTokenizer
from max.pipelines.modeling.types import PipelineTask
from . import weight_adapters
from .model import MyModel
from .model_config import MyModelConfig
my_arch = SupportedArchitecture(
name="MyModelForCausalLM",
task=PipelineTask.TEXT_GENERATION,
example_repo_ids=["my-org/my-model"],
default_encoding="bfloat16",
supported_encodings={"bfloat16"},
default_weights_format=WeightsFormat.safetensors,
weight_adapters={
WeightsFormat.safetensors: weight_adapters.convert_safetensor_state_dict,
},
pipeline_model=MyModel,
tokenizer=TextTokenizer,
context_type=TextContext,
config=MyModelConfig,
batching=Llama3ModuleV3BatchProcessor,
memory_planner=PagedMemoryPlanner,
supports_overlap_scheduler=False,
supports_device_graph_capture=False,
)MyModel subclasses
ModuleV3PipelineModelWithKVCache,
which builds your max.experimental.nn module
and compiles it with the checkpoint weights. Set
model_config_cls = MyModelConfig on MyModel: the pipeline model uses it to
size the KV cache and build the batch processor, and raises an error at startup
when it's unset. MAX reads the config field in arch.py to plan memory, so
both name the same class. The batching field names the batch processor that
prepares each batch's inputs, and memory_planner sizes the KV cache.
When you pass your package to max serve with --custom-architectures, MAX
imports it, reads the ARCHITECTURES list, and registers every entry. When the
server starts, MAX matches the checkpoint's architectures[0] field against the
registered names to pick the right implementation. If you expose a
SupportedArchitecture whose name matches a built-in architecture's name,
your custom architecture takes precedence.
Run max serve with your custom architecture
From the directory that contains your architecture package, run:
max serve \
--model my-org/my-model \
--custom-architectures my_arch_package-
--modelspecifies the Hugging Face model ID or local path for the checkpoint to load. -
--custom-architecturesspecifies your custom architecture package. If you run the command from outside your architecture's root directory, you can pass a directory path followed by a colon and the Python package name, such asfolder/path/to/import:my_arch_package. -
--device-memory-utilization(default0.9) sets the fraction of free device memory the server uses. The KV cache gets what remains after the weights and activations.
If startup reports that the model and KV cache exceed available memory, reduce
--max-length or --max-batch-size. If other processes share the device and
allocation fails, lower the fraction, for example
--device-memory-utilization 0.6.
An AI coding agent can run this for you with the serve-model skill. The
skill sets up the environment and picks the flags your model needs, including
--custom-architectures, --devices, --quantization-encoding, and
--trust-remote-code. See Modular skills to
install it.
Send a request
When the server is ready to accept requests, you'll see:
Server ready on http://0.0.0.0:8000 (Press CTRL+C to quit)The server exposes an OpenAI-compatible API. Send a chat completion request with cURL or the OpenAI Python client:
- cURL
- Python
curl -X POST http://localhost:8000/v1/chat/completions \
-H "Content-Type: application/json" \
-d '{
"model": "my-org/my-model",
"messages": [
{"role": "user", "content": "Hello! Can you help me with a simple task?"}
],
"max_completion_tokens": 100
}'from openai import OpenAI
client = OpenAI(
base_url="http://localhost:8000/v1",
api_key="EMPTY", # Required by the API but not used by MAX
)
response = client.chat.completions.create(
model="my-org/my-model",
messages=[
{"role": "user", "content": "Hello! Can you help me with a simple task?"}
],
max_completion_tokens=100,
)
print(response.choices[0].message.content)The model field in the request body must match the value you passed to
--model when you started the server.
Troubleshoot serve-time errors
Serve-time errors such as an unsupported quantization encoding, a missing
weight adapter, or an unmatched architecture name usually indicate a mismatch
between your SupportedArchitecture and the checkpoint. If you see one of
these errors, check that:
-
namematches thearchitectures[0]field in your checkpoint'sconfig.json, character for character. -
supported_encodingsincludes every encoding your checkpoint declares. -
weight_adaptershas an entry for each weight format you want to load. For example, a.safetensorscheckpoint needs aWeightsFormat.safetensorsentry, and a.ggufcheckpoint needs aWeightsFormat.ggufentry.
If the server fails while it compiles the model, compare your weight
adapter's output with your module's parameter names. A parameter with no
matching weight raises a KeyError, and a weight whose shape or dtype differs
from its parameter raises a ValueError. See Map checkpoint
weights.
To learn more about these fields, see Model pipeline.
Next steps
Serving features and performance optimizations work on any architecture that MAX loads.
- Prefix caching: Reuse the KV cache across requests that share a prompt prefix to cut time-to-first-token on repeated workloads.
- LoRA adapters: Serve multiple fine-tuned adapters on top of your base architecture without loading separate model copies.
- Benchmark MAX on NVIDIA or AMD GPUs: Measure the throughput and latency of the endpoint you started.