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.architectures.mamba

Mamba state-space architecture for text generation.

MambaConfig

class max.pipelines.architectures.mamba.MambaConfig(*, hidden_size, intermediate_size, num_hidden_layers, vocab_size, max_seq_len, dtype, devices, d_state, dt_rank=None, conv_kernel=4, x_proj_dim=None, rms_norm_eps=None, use_bias=False, use_conv_bias=True, residual_in_fp32=True, tie_word_embeddings=True, return_logits=ReturnLogits.LAST_TOKEN, quant_config=None, use_subgraphs=True, data_parallel_degree=1, expand=2, quantization_encoding=None)

source

Bases: ArchConfigWithKVCache

Model configuration for Mamba graph construction/execution.

Parameters:

• hidden_size (int)
• intermediate_size (int)
• num_hidden_layers (int)
• vocab_size (int)
• max_seq_len (int)
• dtype (DType)
• devices (list[DeviceRef])
• d_state (int)
• dt_rank (int | str | None)
• conv_kernel (int)
• x_proj_dim (int | None)
• rms_norm_eps (float | None)
• use_bias (bool)
• use_conv_bias (bool)
• residual_in_fp32 (bool)
• tie_word_embeddings (bool)
• return_logits (ReturnLogits)
• quant_config (QuantConfig | None)
• use_subgraphs (bool)
• data_parallel_degree (int)
• expand (int)
• quantization_encoding (SupportedEncoding | None)

DEFAULT_ENCODING

DEFAULT_ENCODING: ClassVar[SupportedEncoding] = 'float32'

source

SUPPORTED_ENCODINGS

SUPPORTED_ENCODINGS: ClassVar[set[SupportedEncoding]] = {'bfloat16', 'float32'}

source

calculate_max_seq_len()

static calculate_max_seq_len(pipeline_config, huggingface_config, model_config=None)

source

Parameters:

• pipeline_config (PipelineConfig)
• huggingface_config (AutoConfig)
• model_config (MAXModelConfig | None)

Return type:

int

conv_kernel

conv_kernel: int = 4

source

d_state

d_state: int

source

data_parallel_degree

data_parallel_degree: int = 1

source

devices

devices: list[DeviceRef]

source

dt_rank

dt_rank: int | str | None = None

source

dtype

dtype: DType

source

expand

expand: int = 2

source

finalize()

finalize(huggingface_config, state_dict, return_logits)

source

Set parameters that require introspecting the state dict.

Parameters:

• huggingface_config (AutoConfig)
• state_dict (dict[str, WeightData])
• return_logits (ReturnLogits)

Return type:

None

get_kv_params()

get_kv_params()

source

Return minimal dummy KV cache params.

Mamba uses SSM state caching, not attention KV cache, but the pipeline interfaces require KVCacheParams for memory estimation and PagedKVCacheManager initialization. The tiny dimensions (1 head, 1 dim, 1 layer) make the allocation negligible.

Return type:

KVCacheParams

get_max_seq_len()

get_max_seq_len()

source

Returns the default maximum sequence length for the model.

Subclasses should determine whether this value can be overridden by setting the --max-length (pipeline_config.model.max_length) flag.

Return type:

int

get_num_layers()

static get_num_layers(huggingface_config)

source

Parameters:

huggingface_config (AutoConfig)

Return type:

int

get_ssm_cache_params()

get_ssm_cache_params()

source

Return type:

SSMStateCacheParams

help()

static help()

source

Return type:

dict[str, str]

hidden_size

hidden_size: int

source

initialize()

classmethod initialize(pipeline_config, model_config=None)

source

Initialize the config from a PipelineConfig.

Parameters:

• pipeline_config (PipelineConfig) – The pipeline configuration.
• model_config (MAXModelConfig | None) – The model configuration to read from. When None (the default), pipeline_config.model is used. Pass an explicit config (e.g. pipeline_config.draft_model) to initialize the arch config for a different model.

Return type:

Self

initialize_from_config()

classmethod initialize_from_config(pipeline_config, huggingface_config, model_config=None)

source

Parameters:

• pipeline_config (PipelineConfig)
• huggingface_config (AutoConfig)
• model_config (MAXModelConfig | None)

Return type:

Self

intermediate_size

intermediate_size: int

source

max_seq_len

max_seq_len: int

source

num_hidden_layers

num_hidden_layers: int

source

quant_config

quant_config: QuantConfig | None = None

source

quantization_encoding

quantization_encoding: SupportedEncoding | None = None

source

residual_in_fp32

residual_in_fp32: bool = True

source

return_logits

return_logits: ReturnLogits = 'last_token'

source

rms_norm_eps

rms_norm_eps: float | None = None

source

tie_word_embeddings

tie_word_embeddings: bool = True

source

use_bias

use_bias: bool = False

source

use_conv_bias

use_conv_bias: bool = True

source

use_subgraphs

use_subgraphs: bool = True

source

vocab_size

vocab_size: int

source

x_proj_dim

x_proj_dim: int | None = None

source

MambaModel

class max.pipelines.architectures.mamba.MambaModel(pipeline_config, session, devices, kv_cache_config, weights, adapter=None, return_logits=ReturnLogits.LAST_TOKEN, return_hidden_states=ReturnHiddenStates.NONE, max_batch_size=1)

source

Bases: PipelineModelWithKVCache[TextContext]

Mamba pipeline model with incremental SSM state caching.

Uses separate compiled prefill and step models. The prefill model processes the full prompt and extracts per-layer conv/ssm states. The step model processes single new tokens using cached states.

Parameters:

• pipeline_config (PipelineConfig)
• session (InferenceSession)
• devices (list[Device])
• kv_cache_config (KVCacheConfig)
• weights (Weights)
• adapter (WeightsAdapter | None)
• return_logits (ReturnLogits)
• return_hidden_states (ReturnHiddenStates)
• max_batch_size (int)

batch_processor_cls

batch_processor_cls

source

alias of MambaBatchProcessor

compute_log_probabilities()

compute_log_probabilities(session, model_inputs, model_outputs, next_tokens, batch_top_n, batch_echo)

source

Optional method that can be overridden to compute log probabilities.

Parameters:

• session (InferenceSession) – Inference session to compute log probabilities within.
• model_inputs (ModelInputs) – Inputs to the model returned by prepare_initial_token_inputs().
• model_outputs (ModelOutputs) – Outputs returned by execute().
• next_tokens (Buffer) – Sampled tokens. Should have shape=[batch size]
• batch_top_n (list[int]) – Number of top log probabilities to return per input in the batch. For any element where top_n == 0, the LogProbabilities is skipped.
• batch_echo (list[bool]) – Whether to include input tokens in the returned log probabilities.

Returns:

List of log probabilities.

Return type:

list[LogProbabilities | None]

execute()

execute(model_inputs)

source

Executes the graph with the given inputs.

Parameters:

model_inputs (ModelInputs) – The model inputs to execute, containing tensors and any other required data for model execution.

Returns:

ModelOutputs containing the pipeline’s output tensors.

Return type:

ModelOutputs

This is an abstract method that must be implemented by concrete PipelineModels to define their specific execution logic.

get_kv_params()

classmethod get_kv_params(huggingface_config, pipeline_config, devices, kv_cache_config, cache_dtype)

source

Return minimal dummy KV cache params.

Mamba uses SSM state caching internally, not attention KV cache. These dummy params satisfy the PipelineModelWithKVCache interface with negligible memory overhead.

Parameters:

• huggingface_config (AutoConfig)
• pipeline_config (PipelineConfig)
• devices (list[DeviceRef])
• kv_cache_config (KVCacheConfig)
• cache_dtype (DType)

Return type:

KVCacheParamInterface

get_num_layers()

classmethod get_num_layers(huggingface_config)

source

Parameters:

huggingface_config (AutoConfig)

Return type:

int

model_config_cls

model_config_cls

source

alias of MambaConfig

release()

release(request_id)

source

Release SSM cache slot when a request completes.

Parameters:

request_id (RequestID)

Return type:

None

MambaModelInputs

class max.pipelines.architectures.mamba.MambaModelInputs(tokens, input_row_offsets, return_n_logits, is_prefill=True, layer_states=None, request_ids=None)

source

Bases: ModelInputs

Inputs for the Mamba pipeline model.

Parameters:

• tokens (Buffer)
• input_row_offsets (Buffer)
• return_n_logits (Buffer)
• is_prefill (bool)
• layer_states (list[Buffer])
• request_ids (list[RequestID])

input_row_offsets

input_row_offsets: Buffer

source

is_prefill

is_prefill: bool

source

layer_states

layer_states: list[Buffer]

source

request_ids

request_ids: list[RequestID]

source

return_n_logits

return_n_logits: Buffer

source

tokens

tokens: Buffer

source

SSMStateCacheParams

class max.pipelines.architectures.mamba.SSMStateCacheParams(num_layers, intermediate_size, d_state, conv_kernel, dtype)

source

Bases: object

Parameters for SSM state cache memory estimation.

SSM cache is fixed-size per batch element (no sequence-length scaling):

conv_state: num_layers * intermediate_size * conv_kernel * dtype_bytes ssm_state: num_layers * intermediate_size * d_state * dtype_bytes

Parameters:

• num_layers (int)
• intermediate_size (int)
• d_state (int)
• conv_kernel (int)
• dtype (DType)

compute_max_seq_len_fitting_in_cache()

compute_max_seq_len_fitting_in_cache(cache_memory)

source

SSM cache doesn’t scale with seq length — no constraint.

Parameters:

cache_memory (int)

Return type:

int | None

conv_kernel

conv_kernel: int

source

d_state

d_state: int

source

dtype

dtype: DType

source

estimated_memory_size()

estimated_memory_size(available_cache_memory, max_batch_size, max_seq_len)

source

Total SSM cache bytes. Independent of sequence length.

Parameters:

• available_cache_memory (int)
• max_batch_size (int)
• max_seq_len (int)

Return type:

int

intermediate_size

intermediate_size: int

source

num_layers

num_layers: int

source

per_element_bytes

property per_element_bytes: int

source