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

MemoryPlanner

MemoryPlanner

class max.pipelines.kv_cache.MemoryPlanner(config)

source

Bases: object

Base class for pipeline model memory planning.

Provides default implementations for all estimation methods. Subclasses override the methods that require architecture-specific logic:

• Estimating KV cache memory requirements.
• Estimating activation, weight, and signal-buffer memory overheads specific to the model.

A MemoryPlanner is constructed from a ModelConfig alone (not from a full PipelineConfig) so that it can be used independently of the pipeline stack.

Initializes the memory planner with the model config.

Parameters:

config (Any) – Model configuration.

estimate_activation_memory()

estimate_activation_memory(pipeline_config, huggingface_config)

source

Estimates activation memory beyond model weights.

The default implementation returns 0. Override in subclasses that require temporary buffers for large intermediate tensors (e.g. MLA up-projection during prefill, expert-parallel routing buffers).

Parameters:

• pipeline_config (Any) – Pipeline configuration.
• huggingface_config (Any) – HuggingFace model configuration.

Returns:

Estimated activation memory in bytes.

Return type:

int

estimate_signal_buffer_memory()

estimate_signal_buffer_memory(pipeline_config, arch_config=None)

source

Estimates signal-buffer memory in bytes across all devices.

Signal buffers are fixed-size per-GPU allocations used by P2P collectives. The default returns 0 for single-device pipelines and delegates to pipeline_config.estimate_signal_buffer_memory for multi-device.

Models that perform allreduce unconditionally (e.g. via VocabParallelEmbedding) need signal buffers even on a single device. Set always_signal_buffers=True on the planner class to enable this.

Parameters:

• pipeline_config (Any) – Pipeline configuration.
• arch_config (Any | None) – Unused; kept for interface parity with PipelineConfig.estimate_signal_buffer_memory().

Returns:

Estimated signal-buffer memory in bytes across all devices.

Return type:

int

estimate_weights_size()

estimate_weights_size(pipeline_config)

source

Estimates the memory consumed by model weights in bytes.

The default implementation delegates to pipeline_config.model.weights_size(). Override in subclasses that need architecture-specific weight accounting (e.g. expert-parallel sharding adjustments).

Parameters:

pipeline_config (Any) – Pipeline configuration providing the model config.

Returns:

Estimated weight memory in bytes.

Return type:

int

infer_max_batch_size()

infer_max_batch_size(pipeline_config, devices, weights_size)

source

Infers an architecture-specific default max_batch_size.

Memory planning calls this when the user did not set max_batch_size, before estimate_activation_memory() runs. The default returns None, deferring to the framework-wide inference in memory estimation. Override in planners for architectures with per-request device memory beyond the KV cache (e.g. recurrent-state pools) that need a tighter default.

Parameters:

• pipeline_config (Any) – Pipeline configuration.
• devices (list[Device]) – Loaded devices the model will run on.
• weights_size (int) – Estimated model weights size in bytes.

Returns:

The inferred max_batch_size, or None to use the framework default.

Return type:

int | None