nnterp · encyclopedia mamba2 · nnterp.families.mamba2

family mamba2

Mamba-2

No attention and no MLP: each block is one norm and a Mamba-2 (SSD) mixer whose heads each keep a state matrix, with one scalar decay per head and token.

State space (Mamba-2)One sublayer per block

The block

01 · the residual stream

One block of the stack. Hover any part for its nnterp name; click to pin the card.

embed_tokenstoken_embeddings
→
layers× 24
→
normfinal_norm_output
→
lm_headlm_head_output
→
logitsthe output
i = 0

identity layers[i].input + linear_attn.attention_output == layer_output

The API

02 · support()

Values, by host

⚠︎ marks a value this checkpoint has only under some load: hover it for the condition.

model Mamba2ForCausalLMlayoutwhat it isread at
logits Logits [batch seq vocab_size] The logits the model returns, after anything it does past lm_head (a softcap, a scale)

the module's own output
token_embeddings Residual [batch seq hidden_size] The embedding module's output; layers[0].input is what enters the first block

the child module embed_tokens's output
final_norm_output Residual [batch seq hidden_size] The residual stream after the final norm, what lm_head reads

the child module norm's output
lm_head_output Logits [batch seq vocab_size] The head's logits before anything the model does past lm_head (a softcap, a scale)

the child module lm_head's output
next_token_probs NextTokenProbs [batch vocab_size] The next-token distribution at the last position; derived, read-only

the module's own output
input_ids Tokens [batch seq] The token ids the model was called with

the module's own input
attention_mask Tokens [batch seq] The attention mask the model was called with; zeros are padding

the module's own input
input_size scalar [batch, seq] of the current call; read-only

the module's own input
model.layers[i] Mamba2Blocklayoutwhat it isread at
layer_input Residual [batch seq hidden_size] The residual stream entering the block

the module's own input
attention_norm_output Residual [batch seq hidden_size] The normed stream entering the attention

the child module norm's output
⚠︎layer_mid scalar Unavailable: a state-space block: one sublayer, the mixer, so there is no stream between sublayers and no MLP

nowhere: the family does not serve it
⚠︎mlp_norm_output scalar Unavailable: a state-space block: one sublayer, the mixer, so there is no stream between sublayers and no MLP

nowhere: the family does not serve it
⚠︎mlp_input scalar Unavailable: a state-space block: one sublayer, the mixer, so there is no stream between sublayers and no MLP

nowhere: the family does not serve it
layer_output Residual [batch seq hidden_size] The residual stream leaving the block, a tensor even when the block returns a tuple

the module's own output
model.layers[i].linear_attn Mamba2Mixerlayoutwhat it isread at
attention_output Residual [batch seq hidden_size] What the mixer adds to the residual stream

the module's own output
⚠︎state scalar Unavailable: the chunk scan computes every token's state in one tensor per call, not one occurrence per token to walk with tracer.iter; read states, or state_after(t), after nnterp.chunk_per_token(model)

nowhere: the family does not serve it
⚠︎states States [batch seq heads key_dim value_dim] The state after every token of this call; needs nnterp.chunk_per_token(model)

inside the forward, at the operation that _states finds on each call
⚠︎attention_queries SSDQueries [batch seq groups state_dim] C, the queries reading the state

inside the forward: the arguments of the kernel this call runs (a prompt's chunked one or a decode step's recurrent one), argument C
⚠︎attention_keys SSDKeys [batch seq groups state_dim] B, the keys writing into the state

inside the forward: the arguments of the kernel this call runs (a prompt's chunked one or a decode step's recurrent one), argument B
⚠︎attention_values SSDValues [batch seq heads head_dim] x, the values written into the state

inside the forward: the arguments of the kernel this call runs (a prompt's chunked one or a decode step's recurrent one), argument hidden_states
⚠︎betas Gates [batch seq heads] dt, the per-token write strength into the state

inside the forward: the arguments of the kernel this call runs (a prompt's chunked one or a decode step's recurrent one), argument dt
⚠︎decays Gates [batch seq heads] A * dt, the per-token log decay of the state

inside the forward: the arguments of the kernel this call runs (a prompt's chunked one or a decode step's recurrent one), argument dt
⚠︎state_input State | None [batch heads key_dim value_dim] The state entering the layer, or None at the start of a prompt (a copy of the cache's buffer)

inside the forward: the arguments of the kernel this call runs (a prompt's chunked one or a decode step's recurrent one), argument state
⚠︎attention_head_outputs SSDHeadOutputs [batch seq heads head_dim] y, the per-head outputs before the gated norm and the output projection

inside the forward: the output of the kernel this call runs (a prompt's chunked one or a decode step's recurrent one), the element that call returns it in
⚠︎state_output State [batch heads key_dim value_dim] The state leaving the layer

inside the forward, at the operation that _state_output finds on each call
⚠︎set_state_after scalar Unavailable: the chunk scan computes every boundary state in one cumulative step from the initial state, so a state written at token t does not flow into later tokens' states; assign state_input to change where a call starts

nowhere: the family does not serve it

Printout

What print(model) shows on AntonV/mamba2-130m-hf: the native modules, standard names and every standard value with its layout.

Mamba2ForCausalLM(
  (lm_head): Linear(in_features=768, out_features=50288, bias=False)
  (generator): Generator(
    (streamer): Streamer()
  )
  (embed_tokens): Embedding(50288, 768)
  (layers): ModuleList(
    (0-23): 24 x Mamba2Block(
      (norm): Mamba2RMSNorm()
      (linear_attn/mixer): Mamba2Mixer(
        (act): SiLUActivation()
        (conv1d): Conv1d(1792, 1792, kernel_size=(4,), stride=(1,), padding=(3,), groups=1792)
        (in_proj): Linear(in_features=768, out_features=3352, bias=False)
        (norm): MambaRMSNormGated()
        (out_proj): Linear(in_features=1536, out_features=768, bias=False)
        (attention_queries) -> SSDQueries [batch seq groups state_dim]: C, the queries reading the state
        (attention_keys) -> SSDKeys [batch seq groups state_dim]: B, the keys writing into the state
        (attention_values) -> SSDValues [batch seq heads head_dim]: x, the values written into the state
        (betas) -> Gates [batch seq heads]: dt, the per-token write strength into the state
        (decays) -> Gates [batch seq heads]: A * dt, the per-token log decay of the state
        (state_input) -> State | None [batch heads key_dim value_dim]: The state entering the layer, or None at the start of a prompt (a copy of the cache's buffer)
        (attention_head_outputs) -> SSDHeadOutputs [batch seq heads head_dim]: y, the per-head outputs before the gated norm and the output projection
        (state_output) -> State [batch heads key_dim value_dim]: The state leaving the layer
        (states) -> States [batch seq heads key_dim value_dim]: The state after every token of this call; needs nnterp.chunk_per_token(model)
        (state): Unavailable: the chunk scan computes every token's state in one tensor per call, not one occurrence per token to walk with tracer.iter; read states, or state_after(t), after nnterp.chunk_per_token(model)
        (set_state_after): Unavailable: the chunk scan computes every boundary state in one cumulative step from the initial state, so a state written at token t does not flow into later tokens' states; assign state_input to change where a call starts
        (attention_output) -> Residual [batch seq hidden_size]: What the mixer adds to the residual stream
      )
      (attention_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the attention
      (layer_mid): Unavailable: a state-space block: one sublayer, the mixer, so there is no stream between sublayers and no MLP
      (layer_input) -> Residual [batch seq hidden_size]: The residual stream entering the block
      (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): Mamba2RMSNorm()
  (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
)

Config

model.num_layers24
model.hidden_size768
model.num_heads24
model.num_kv_heads24
model.head_dim64
model.qk_head_dim64
model.vocab_size50288
model.intermediate_size1536
num_hidden_layers24
hidden_size768
head_dim64
vocab_size50288
hidden_actsilu
tie_word_embeddingsTrue
layer_norm_epsilon1e-05
state_size128
expand2
conv_kernel4
residual_in_fp32True
n_groups1
chunk_size256
time_step_min0.001

The Notes

03 · what to know before you run

The block, in order

out = x + mixer(norm(x))                  # added in float32
mixer(h): z, xBC, dt = in_proj(h)
          x, B, C = silu(conv1d(xBC))     # causal, width 4
          y = ssd(x, dt, A, B, C) + D * x # attention_head_outputs
          return out_proj(norm(y * silu(z)))

The block's norm keeps its native name, norm, and model.layers[i].norm is the norm before the mixer; model.layers[i].linear_attn.norm is the mixer's own gated norm and model.norm the final norm_f. The convolution runs over x, B and C together. The gated norm multiplies y by silu(z) and then normalizes over the mixer's whole width (1536 on 130m), whatever n_groups is; attention_head_outputs is read before the gate and the norm.

The contribution is the whole block

attention_output is the mixer's output (out_proj's), and it is all the block adds: layers[i].input + attention_output == layer_output exactly, so ablating attention_output ablates the block. The add is in float32: attention_output is in the load dtype and layer_output is float32.

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

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

The stream grows through the stack: on 130m, at the last token of a 14-token prompt, its norm is 7.0 entering block 0 and 1316 entering block 23, and block 23 adds a vector of norm 1376.

Loading: route the kernels

With mamba_ssm installed transformers sends the prompt's chunk scan and the decode step's update to its Triton kernels. On a CPU every trace then fails inside the chunk scan, even one that reads only layer_output (ValueError: Pointer argument cannot be accessed from Triton (cpu tensor?), or RuntimeError: invalid argument to exchangeDevice with no GPU visible). On a GPU, unrouted, layer_output, attention_output and logits read normally and every other linear_attn value reports the kernel in support(). Route before the first trace:

import nnterp
from nnterp import StandardizedTransformer, route_kernels

route_kernels(nnterp.families.mamba2, "torch")
model = StandardizedTransformer("AntonV/mamba2-130m-hf")

On 130m in float32 the kernel and the routed torch scan give logits within 0.06 of each other (logits up to 120), with the same top token at every position; in bfloat16 within 2.25. The torch scan took 146 ms on a 14-token prompt against the kernel's 23 ms.

The mixer's values

On 130m each mixer has 24 heads of 64 channels and a state of 128 per head, in one group: attention_queries (C) and attention_keys (B) are [batch, seq, 1, 128], shared by all 24 heads, and attention_values (x) and attention_head_outputs (y) are [batch, seq, 24, 64]. betas is the step softplus(dt + dt_bias) and decays is betas * A, both [batch, seq, 24]: one scalar decay per head and token. Mamba-Codestral has 128 heads in 8 groups, so C and B are [batch, seq, 8, 128] and each serves 16 heads.

The rates are far apart on real weights. A = -exp(A_log) runs from -4e-4 to -36316 over the 576 heads of 130m (median -0.21). Over every block of a 14-token prompt, exp(decays) has a median of 0.89; 5% of entries keep less than 1% of the state and 19% keep more than 99%.

A zero step skips a token exactly

Every Mamba-2 checkpoint here has time_step_limit (0.0, inf), so the chunk scan clamps nothing, and a written zero runs as zero: the token writes nothing and the state does not decay. On 130m, with chunk_per_token, the state after token 3 then equals the state after token 2 bit for bit, and a read of betas after the write returns the zero.

nnterp.chunk_per_token(model)
mix = model.layers[0].linear_attn
with model.trace(prompt):
    betas = mix.betas.clone()
    betas[:, 3] = 0
    mix.betas = betas
    states = mix.states.save()

assert torch.equal(states[:, 3], states[:, 2])

Writing every block's betas back unchanged moves 130m's float32 logits by up to 1.2e-4.

The hidden attention

Per head the scan is a causal linear map over tokens: y_t = sum_s alpha[t, s] x_s + D x_t with alpha[t, s] = (C_t . B_s) * exp(decays_{s+1} + ... + decays_t) * betas_s, one scalar per head and query-key pair, so the map is a [seq, seq] matrix per head. The served values rebuild y exactly (float32):

mix = model.layers[1].linear_attn
with model.trace(prompt):
    C = mix.attention_queries[0].save()         # [seq, groups, state_dim]
    B = mix.attention_keys[0].save()            # [seq, groups, state_dim]
    x = mix.attention_values[0].save()          # [seq, heads, head_dim]
    beta = mix.betas[0].save()                  # [seq, heads]
    decay = mix.decays[0].save()                # [seq, heads]
    y = mix.attention_head_outputs[0].save()    # [seq, heads, head_dim]

D = mix._module.D
per_group = x.shape[1] // C.shape[1]
C, B = C.repeat_interleave(per_group, 1), B.repeat_interleave(per_group, 1)
cum = decay.cumsum(0)
causal = torch.ones(len(x), len(x), dtype=torch.bool, device=x.device).tril()[..., None]
kept = (cum[:, None] - cum[None, :]).masked_fill(~causal, -torch.inf).exp()
alpha = torch.einsum("thn,shn->tsh", C, B) * kept * beta     # [query t, key s, head]
torch.testing.assert_close(torch.einsum("tsh,shd->thd", alpha, x) + D[:, None] * x, y)

The state

state_output is the state after the call's last token, [batch, 24, 128, 64] on 130m, key side first. states, the state after every token, needs nnterp.chunk_per_token(model): the chunk scan keeps the state only every chunk_size tokens (256). Then states[:, -1] equals state_output.

nnterp.chunk_per_token(model)
with model.trace(prompt):
    states = mix.states.save()           # [batch, seq, 24, 128, 64] on 130m
    final = mix.state_output.save()

assert torch.equal(states[:, -1], final)

On a decode step the update comes before the read, so under generate read state_output before attention_head_outputs on every step after the first:

heads, leaving = [], []
with model.generate(prompt, max_new_tokens=3, do_sample=False) as tracer:
    for step in tracer.iter[:3]:
        if step == 0:
            heads.append(mix.attention_head_outputs.save())
            leaving.append(mix.state_output.save())
        else:
            leaving.append(mix.state_output.save())
            heads.append(mix.attention_head_outputs.save())

The width-4 convolution runs over x, B and C, so a state written through state_input is not all a block carries into the next tokens.

The readout and the embeddings

On the 130m to 2.7b conversions lm_head is embed_tokens' weight, so an edit to one is an edit to the other; Mamba-Codestral's head has its own. project_on_vocab at the last block's layer_output equals logits exactly. The final norm's gain is its weight (mean 0.94, up to 4.9 on 130m). Nothing is added after embed_tokens, so token_embeddings equals layers[0].input.

The checkpoints and the tokenizer

state-spaces/mamba2-130m to -2.7b are the original release, stored for the mamba_ssm package with no model_type; transformers loads the AntonV/mamba2-*-hf conversions, which nnterp's docs use. They tokenize with GPT-NeoX-20B's BPE (GPTNeoXTokenizer, 50277 tokens; the embedding has 50288 rows), prepend no BOS, and use <|endoftext|> (id 0) as BOS, EOS and padding. Mamba-Codestral-7B is a code model with a 32768-token vocabulary.

Published work

  • Transformers are SSMs (Dao and Gu, 2024) defines SSD as the masked attention above: C, B and x in the roles of attention_queries, attention_keys and attention_values, with the decay mask built from decays.

From the family module

04 · nnterp/families/mamba2.py

Mamba-2 (Mamba2ForCausalLM).

A pure state-space model: no attention and no MLP. The tree::

backbone.embeddings
backbone.layers[i]       Mamba2Block
    .norm                RMSNorm, the block's only norm
    .mixer               Mamba2Mixer, an SSD mixer
backbone.norm_f
lm_head

Each block is hidden + mixer(norm(hidden)) (the residual in float32 when residual_in_fp32), so the block's one contribution is the mixer's output and layer.input + linear_attn.attention_output == layer_output. The mixer is nnterp.StateSpace under the standard name linear_attn (the recurrent mixer's name on every hybrid). The block's norm keeps its name: a "norm" key would also match the mixer's own gated norm, and its output is linear_attn.input. There is no self_attn and no mlp, so support lists neither.

The sizes are the mixer's: num_heads and head_dim are the SSD heads (config.num_heads, config.head_dim), and intermediate_size is the mixer's inner width, expand * hidden_size, what in_proj expands to and out_proj reads from (the model has no MLP). With residual_in_fp32 the residual stream is float32 whatever the weights are, and the model casts it to lm_head's dtype and its logits to float32, so project_on_vocab is the family's own.