IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /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. /get-started.md).

Python module

max.pipelines.architectures.mimo_v2

MiMo-V2.6-Flash (MiMoV2ForCausalLM).

MiMoV2Config​

class max.pipelines.architectures.mimo_v2.MiMoV2Config(*, vocab_size, hidden_size, num_hidden_layers, rms_norm_eps, layer_types, num_attention_heads, head_dim, v_head_dim, qkv_layouts, rotary_dim, rope_thetas, sliding_window, attention_value_scale, sinks, intermediate_size, moe_layers, n_routed_experts, num_experts_per_tok, moe_intermediate_size, norm_topk_prob, quant, dtype, devices, kv_params, max_seq_len, use_subgraphs=True, return_logits=ReturnLogits.LAST_TOKEN, return_hidden_states=ReturnHiddenStates.NONE, target_layer_ids=None)

source

Bases: ArchConfigWithStoredKVParams, ArchConfigWithKVCache

Model configuration for MiMo-V2.6-Flash.

attention_chunk_size in the checkpoint config is a leftover key and is deliberately not read: the model has no chunked local attention.

Parameters:

DEFAULT_ENCODING​

DEFAULT_ENCODING: ClassVar[max.pipelines.lib.config.SupportedEncoding] = 'float4_e2m1fnx2'

source

SUPPORTED_ENCODINGS​

SUPPORTED_ENCODINGS: ClassVar[set[max.pipelines.lib.config.SupportedEncoding]] = {'float4_e2m1fnx2'}

source

attention_scale​

property attention_scale: float

source

The softmax scale, from the unpadded Q/K head dim.

attention_value_scale​

attention_value_scale: float

source

construct_kv_params()​

classmethod construct_kv_params(huggingface_config, pipeline_config, devices, kv_cache_config, cache_dtype, *, allow_kv_head_replication=False)

source

Builds the {sliding_attention, full_attention} KV tree.

Parameters:

Return type:

MultiKVCacheParams

devices​

devices: list[DeviceRef]

source

dtype​

dtype: DType

source

from_huggingface_config()​

classmethod from_huggingface_config(huggingface_config, *, devices, kv_params, max_seq_len)

source

Builds the config from the checkpoint’s HuggingFace config.

Parameters:

Returns:

The config.

Raises:

ValueError – If the config describes a variant this implementation does not handle, or a quantization other than Xiaomi’s.

Return type:

Self

head_dim​

head_dim: int

source

Q and K head dim.

hidden_size​

hidden_size: int

source

initialize()​

classmethod initialize(pipeline_config, model_config=None, *, max_seq_len)

source

Initializes the config from the pipeline config.

Parameters:

  • pipeline_config (PipelineConfig) – The pipeline configuration.
  • model_config (MAXModelConfig | None) – The model configuration to read. Defaults to pipeline_config.model.
  • max_seq_len (int) – The maximum sequence length to build for.

Returns:

The initialized config.

Raises:

ValueError – If a device is not an SM100 GPU, or the checkpoint is not one this implementation handles.

Return type:

Self

intermediate_size​

intermediate_size: int

source

Dense MLP width.

kv_params​

kv_params: KVCacheParamInterface

source

layer_types​

layer_types: list[str]

source

Each decoder layer’s KV group, sliding_attention or full_attention.

max_seq_len​

max_seq_len: int

source

moe_intermediate_size​

moe_intermediate_size: int

source

moe_layers​

moe_layers: set[int]

source

n_routed_experts​

n_routed_experts: int

source

norm_topk_prob​

norm_topk_prob: bool

source

num_attention_heads​

num_attention_heads: int

source

num_experts_per_tok​

num_experts_per_tok: int

source

num_hidden_layers​

num_hidden_layers: int

source

qkv_layouts​

qkv_layouts: dict[str, QkvChunkLayout]

source

The fused qkv_proj chunk layout of each KV group.

quant​

quant: MiMoV2QuantScheme

source

return_hidden_states​

return_hidden_states: ReturnHiddenStates = 'none'

source

return_logits​

return_logits: ReturnLogits = 'last_token'

source

rms_norm_eps​

rms_norm_eps: float

source

rope_thetas​

rope_thetas: dict[str, float]

source

rotary_dim​

rotary_dim: int

source

Leading Q/K dims that RoPE rotates, NeoX style.

sinks​

sinks: dict[str, bool]

source

Whether each KV group has a learned per-head attention sink.

sliding_window​

sliding_window: int

source

Keys a sliding layer attends to, the query’s own included.

target_layer_ids​

target_layer_ids: list[int] | None = None

source

For ReturnHiddenStates.SELECTED_LAYERS, the layers whose outputs (post-residual, before the final norm) are returned.

use_subgraphs​

use_subgraphs: bool = True

source

Whether the MoE layers of each attention type share one compiled subgraph.

v_head_dim​

v_head_dim: int

source

vocab_size​

vocab_size: int

source

MiMoV2Model​

class max.pipelines.architectures.mimo_v2.MiMoV2Model(pipeline_config, session, devices, kv_cache_config, weights, *, memory_plan, adapter=None, return_logits=ReturnLogits.LAST_TOKEN, max_batch_size=1)

source

Bases: AlwaysSignalBuffersMixin, GraphPipelineModelWithKVCache[TextContext, Any]

Pipeline model for MiMo-V2.6-Flash.

The graph takes GPT OSS’s inputs (tokens, per-device row offsets, signal buffers and a {sliding_attention, full_attention} KV tree), so GPT OSS’s inputs and batching are reused.

Parameters:

batch_processor_cls​

batch_processor_cls

source

alias of GptOssBatchProcessor

execute()​

execute(model_inputs)

source

Runs the model on inputs from GptOssBatchProcessor.

Parameters:

model_inputs (ModelInputs)

Return type:

ModelOutputs

graph_specs()​

graph_specs(model_config)

source

Returns the decoder graph.

Parameters:

model_config (MiMoV2Config)

Return type:

dict[str, GraphSpec]

model​

property model: Model

source

The compiled and initialized model.

model_config_cls​

model_config_cls

source

alias of MiMoV2Config

Models

Cookbook

Puzzles

Blog

Community