MiniMax-M2 / MiniMax-M2.1 / MiniMax-M2.5 / MiniMax-M2.7

MiniMaxM2ForCausalLM( (lm_head): Linear(in_features=3072, out_features=200064, bias=False) (generator): Generator( (streamer): Streamer() ) (embed_tokens): Embedding(200064, 3072) (layers): ModuleList( (0-61): 62 x MiniMaxM2DecoderLayer( (self_attn): MiniMaxM2Attention( (q_proj): Linear(in_features=3072, out_features=6144, bias=False) (k_proj): Linear(in_features=3072, out_features=1024, bias=False) (v_proj): Linear(in_features=3072, out_features=1024, bias=False) (o_proj): Linear(in_features=6144, out_features=3072, bias=False) (q_norm): MiniMaxM2RMSNorm((6144,), eps=1e-06) (k_norm): MiniMaxM2RMSNorm((1024,), eps=1e-06) (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention (attention_keys) -> Keys [batch num_kv_heads seq qk_head_dim]: The keys entering attention (attention_values) -> Values [batch num_kv_heads seq head_dim]: The values entering attention (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked (attention_premix) -> HeadOutputs [batch seq num_heads head_dim]: Every head's output side by side, as the output projection takes it (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection ) (mlp): MiniMaxM2SparseMoeBlock( (router/gate): MiniMaxM2TopKRouter() (experts): MiniMaxM2Experts( (act_fn): SiLUActivation() ) (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream ) (input_layernorm): MiniMaxM2RMSNorm((3072,), eps=1e-06) (post_attention_layernorm): MiniMaxM2RMSNorm((3072,), eps=1e-06) (layer_input) -> Residual [batch seq hidden_size]: The residual stream entering the block (attention_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the attention (layer_mid) -> Residual [batch seq hidden_size]: The residual stream after the attention add (mlp_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the MLP (mlp_input) -> Residual [batch seq hidden_size]: What the MLP is called with (layer_output) -> Residual [batch seq hidden_size]: The residual stream leaving the block, a tensor even when the block returns a tuple ) ) (norm): MiniMaxM2RMSNorm((3072,), eps=1e-06) (logits) -> Logits [batch seq vocab_size]: The logits the model returns, after anything it does past lm_head (a softcap, a scale) (token_embeddings) -> Residual [batch seq hidden_size]: The embedding module's output; layers[0].input is what enters the first block (final_norm_output) -> Residual [batch seq hidden_size]: The residual stream after the final norm, what lm_head reads (lm_head_output) -> Logits [batch seq vocab_size]: The head's logits before anything the model does past lm_head (a softcap, a scale) (next_token_probs) -> NextTokenProbs [batch vocab_size]: The next-token distribution at the last position; derived, read-only (input_ids) -> Tokens [batch seq]: The token ids the model was called with (attention_mask) -> Tokens [batch seq]: The attention mask the model was called with; zeros are padding (input_size): [batch, seq] of the current call; read-only )

The block, in order

h   = x + self_attn(input_layernorm(x))       # q_norm, k_norm over the whole projection
out = h + mlp(post_attention_layernorm(h))    # 8 of 256 experts per token, on every block

Llama's pre-norm block and Llama's names. Every block's MLP is a mixture: there is no dense MLP and no shared expert, so intermediate_size (1536) is one expert's width. Nothing norms a sublayer's output, so attention_output is o_proj's output and mlp_output the mixture's, and the identity is the plain sum:

with model.trace(prompt):
    x = model.layers[1].input.save()
    attn = model.layers[1].self_attn.attention_output.save()
    mlp = model.layers[1].mlp.mlp_output.save()
    out = model.layers[1].layer_output.save()

torch.testing.assert_close(x + attn + mlp, out)

Queries and keys are normed across all heads

q_norm is one RMSNorm over q_proj's whole output (48 heads × 128 = 6144 wide) and k_norm one over k_proj's (8 × 128 = 1024), applied before the heads are split and before the rotary embedding; attention_queries and attention_keys are read after both. One RMS covers every head, so an edit to one head's slice of q_proj.output changes the other heads' queries too. Edit a single head at attention_queries or attention_keys, which no norm follows:

with model.trace(prompt):
    model.layers[1].self_attn.attention_queries[:, 0] = 0   # head 0 only

48 query heads share 8 key/value heads, so attention_keys and attention_values are [batch, 8, seq, 128] and query head h reads key/value head h // 6: an edit to one key/value head reaches 6 query heads. head_dim is the config's 128, so num_heads * head_dim (6144) is wider than the 3072-wide stream.

Rotary rotates half of each head

The Hub configs set rotary_dim: 64 of each head's 128 dimensions, which MiniMaxM2Config reads as rope_parameters["partial_rotary_factor"] 0.5. The first 64 dimensions of every query and key are rotated, as two halves of 32 (rotate_half). Dimensions 64 to 127 of attention_queries and attention_keys are the normed projections unrotated, and carry no position.

Load with eager for the attention interior

The attention interior needs attn_implementation="eager"; the default load runs sdpa. There is no softcap, sink or window, so the two loads compute the same attention up to rounding (2e-7 in the logits on the pinned tiny checkpoint, in float32). The checkpoints' configs carry an FP8 quantization_config (128 × 128 weight blocks) that leaves the router, e_score_correction_bias and lm_head unquantized.

The router: sigmoid scores, a selection bias, renormalized weights

router_logits is the router's projection, [batch, seq, 256], in the model's dtype. The router takes their sigmoid in float32 and adds the mixture's selection bias, mlp.e_score_correction_bias (a buffer on the mixture module, not on the router), to choose 8 experts. expert_weights are the chosen experts' sigmoids without the bias, renormalized to sum to one, with no scaling factor; they stay float32 under a bfloat16 load. There is no shared expert, so routed_output is mlp_output.

moe = model.layers[1].mlp
with model.trace(prompt):
    logits = moe.router_logits.save()           # [batch, seq, num_experts]
    w = moe.expert_weights.save()               # [batch, seq, top_k]
    idx = moe.expert_indices.save()
    each = moe.expert_outputs.save()            # [batch, seq, top_k, hidden_size]
    out = moe.mlp_output.save()

scores = logits.float().sigmoid()
choice = (scores + moe._module.e_score_correction_bias).topk(moe.top_k, dim=-1).indices
assert torch.equal(choice.sort(-1).values, idx.sort(-1).values)     # chosen with the bias
chosen = scores.gather(-1, idx)
torch.testing.assert_close(w, chosen / chosen.sum(-1, keepdim=True))  # weighted without it
torch.testing.assert_close(each.sum(2), out)                         # no shared expert

The router asks topk for unsorted slots (sorted=False), so slot 0 need not hold the largest weight: compare sets of experts, not slots. A written router_logits moves both the choice and the weights; a written bias moves only the choice.

Ablating an expert touches only the tokens routed to it

Zeroing expert e's weight where expert_indices == e removes that slot's term and leaves the token's other weights as they were, so they then sum to less than one. mlp_output changes on exactly the tokens that chose e:

moe, e = model.layers[1].mlp, 4
with model.trace(prompt):
    idx = moe.expert_indices.save()
    clean = moe.mlp_output.save()
with model.trace(prompt):
    moe.expert_weights = moe.expert_weights.masked_fill(moe.expert_indices == e, 0)
    ablated = moe.mlp_output.save()

changed = (ablated != clean).any(-1)                       # [batch, seq]
assert torch.equal(changed, (idx == e).any(-1))            # only the tokens that chose e

The readout and the tokenizer

logits is lm_head.output, with no cap or scale, and project_on_vocab on the last block's layer_output equals logits. The final norm is a plain RMSNorm whose gain is norm.weight. lm_head and embed_tokens are separate weights. The tokenizer prepends nothing, so model.input_ids is the prompt's tokens alone; its begin and end tokens are ]~!b[ (id 200034) and [e~[ (id 200020).

MiniMax-M2MiniMax-M2.1MiniMax-M2.5MiniMax-M2.7

The block

The API

Values, by host

Printout

Config