nnterp · encyclopedia mamba · nnterp.families.mamba

family mamba

Mamba

No attention and no MLP: each block is one norm and a selective-scan mixer with a state per channel, and reading inside the mixer needs mamba_ssm's CUDA kernels routed to transformers' torch functions.

Selective scan (Mamba-1)

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 MambaForCausalLMlayoutwhat 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] MambaBlocklayoutwhat 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 input_layernorm'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 MambaMixerlayoutwhat it isread at
attention_output Residual [batch seq hidden_size] What the mixer adds to the residual stream

the module's own output
⚠︎state ScanState [batch channels state_dim] The state after one token of the prompt; iterate it with tracer.iter; needs route_kernels(family, 'torch')

inside the forward, at the operation that _token_state_op finds on each call
⚠︎states ScanStates [batch seq channels state_dim] The state after every token of this call; needs route_kernels(family, 'torch')

inside the forward, at the operation that _scan_states finds on each call
⚠︎attention_queries ScanQK [batch seq groups state_dim] C, the vector each token reads the state with

inside the forward: the arguments 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
⚠︎attention_keys ScanQK [batch seq groups state_dim] B, the vector each token writes into the state with

inside the forward: the arguments 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
⚠︎attention_values ScanValues [batch seq channels] x, the input the scan writes 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), the element that call returns it in
⚠︎betas ScanSteps [batch seq channels] The step size dt = softplus(dt + dt_bias), per token and channel; derived, read-only

inside the forward, at the operation that _betas finds on each call
⚠︎decays ScanDecays [batch seq channels state_dim] The per-token log decay dt * A of the state; derived, read-only

inside the forward, at the operation that _decays finds on each call
⚠︎state_input ScanState | None [batch channels state_dim] The state entering the layer on a decode step (a copy of the cache's buffer), None on a prompt

inside the forward: the arguments 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
⚠︎attention_head_outputs ScanValues [batch seq channels] y = C.h + D.x, the scan's read of the state before the gate and out_proj

inside the forward, at the operation that inside.READ_OPS finds on each call
⚠︎state_output ScanState [batch channels state_dim] The state leaving the layer

inside the forward, at the operation that _state_output_key finds on each call

Printout

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

MambaForCausalLM(
  (lm_head): Linear(in_features=768, out_features=50280, bias=False)
  (generator): Generator(
    (streamer): Streamer()
  )
  (embed_tokens): Embedding(50280, 768)
  (layers): ModuleList(
    (0-23): 24 x MambaBlock(
      (input_layernorm/norm): MambaRMSNorm(768, eps=1e-05)
      (linear_attn/mixer): MambaMixer(
        (conv1d): Conv1d(1536, 1536, kernel_size=(4,), stride=(1,), padding=(3,), groups=1536)
        (act): SiLUActivation()
        (in_proj): Linear(in_features=768, out_features=3072, bias=False)
        (x_proj): Linear(in_features=1536, out_features=80, bias=False)
        (dt_proj): Linear(in_features=48, out_features=1536, bias=True)
        (out_proj): Linear(in_features=1536, out_features=768, bias=False)
        (attention_queries) -> ScanQK [batch seq groups state_dim]: C, the vector each token reads the state with
        (attention_keys) -> ScanQK [batch seq groups state_dim]: B, the vector each token writes into the state with
        (attention_values) -> ScanValues [batch seq channels]: x, the input the scan writes into the state
        (betas) -> ScanSteps [batch seq channels]: The step size dt = softplus(dt + dt_bias), per token and channel; derived, read-only
        (decays) -> ScanDecays [batch seq channels state_dim]: The per-token log decay dt * A of the state; derived, read-only
        (state_input) -> ScanState | None [batch channels state_dim]: The state entering the layer on a decode step (a copy of the cache's buffer), None on a prompt
        (attention_head_outputs) -> ScanValues [batch seq channels]: y = C.h + D.x, the scan's read of the state before the gate and out_proj
        (state_output) -> ScanState [batch channels state_dim]: The state leaving the layer
        (state) -> ScanState [batch channels state_dim]: The state after one token of the prompt; iterate it with tracer.iter; needs route_kernels(family, 'torch')
        (states) -> ScanStates [batch seq channels state_dim]: The state after every token of this call; needs route_kernels(family, 'torch')
        (attention_output) -> Residual [batch seq hidden_size]: What the mixer adds to the residual stream
      )
      (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
      (attention_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the attention
      (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): MambaRMSNorm(768, eps=1e-05)
  (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.vocab_size50280
model.intermediate_size1536
architecturesMambaForCausalLM
num_hidden_layers24
hidden_size768
intermediate_size1536
vocab_size50280
hidden_actsilu
tie_word_embeddingsTrue
dtypefloat32
n_layer24
layer_norm_epsilon1e-05
state_size16
expand2
conv_kernel4
residual_in_fp32True
d_model768
d_inner1536
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 (residual_in_fp32)
mixer(h): u, z = in_proj(h);  u = silu(conv1d(u))  # causal, width 4; u is attention_values
          dt, B, C = x_proj(u);  dt = dt_proj(dt)
          y = scan(u, dt, A, B, C) + D * u         # y is attention_head_outputs
          return out_proj(y * silu(z))             # attention_output

The block's pre-norm is natively norm, and nnterp also serves it as input_layernorm; the model's final norm is natively norm_f, and nnterp serves it as model.norm. So model.layers[i].norm is the block's pre-norm, while model.norm is the final norm.

The contribution is the whole block

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

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 its norm at the last token of a ten-token prompt is 7.4 entering block 0 and 188 entering block 23, and block 23's contribution (262) is larger than the stream it joins. A steering vector sized for one depth is small or large at another.

Loading: route the kernels

With mamba_ssm installed transformers dispatches the scan and the decode step to its CUDA kernels. On CPU every trace then fails inside the kernel, even one that reads only layer_output (RuntimeError: Expected u.is_cuda() to be true). On GPU, unrouted, layer_output, attention_output and logits read normally and every other linear_attn value raises nnterp.Unavailable. Route before the trace that reads the mixer's values. Routing still takes effect after an unrouted trace (one that read only layer_output, one that failed on CPU, or one whose read was refused):

import nnterp
from nnterp import StandardizedTransformer, route_kernels

route_kernels(nnterp.families.mamba, "torch")
model = StandardizedTransformer("state-spaces/mamba-130m-hf", device="cpu")

The two paths compute the same model. In float32 on GPU the kernel and the routed torch scan give logits within 2e-4 of each other (130m and 370m, logits up to 98), with the same top token everywhere. In bfloat16 they differ by up to 2.25 on 130m, about what the torch scan on CPU and on GPU differ by (2.0), so the gap is bfloat16 rounding. The torch scan is a token loop: on a 22-token prompt it took 83 ms against the kernel's 46 ms on 130m, 286 ms against 46 ms on 370m.

model.train() sends the forward down a third path, the fused mamba_inner_fn, which route_kernels does not reroute: with mamba_ssm installed and causal-conv1d not, every trace in training mode fails with AssertionError: causal_conv1d_cuda is not available, on CPU and GPU, routed or not. Gradients do not need it; backward() runs through the routed scan in eval mode.

The mixer's values

Each channel of the mixer's width (1536 on 130m) keeps its own state of state_size (16) numbers. attention_queries and attention_keys are C and B, [batch, seq, 1, 16]: one vector per token, shared by every channel. attention_values is x, [batch, seq, 1536], after the convolution and the silu. betas is the step size softplus(dt_proj(dt)) (dt_proj's bias included), one per token and channel; decays is betas[..., None] * A, [batch, seq, 1536, 16], the log of how much of each state entry a token keeps. Both are read-only.

The rates are far apart on real weights. A = -exp(A_log) runs from -2e-6 to -3e8 on 130m (median -12.6), so over every block and a ten-token prompt exp(decays) has a median of 0.48, 28% of entries keep less than 1% and 12% keep more than 99% (on 370m: 0.55, 15%, 9%). The pinned tiny checkpoint has the initialization's A, -1 to -16 on every channel.

The hidden attention

The scan is a causal linear map over tokens, one per channel: y_t = sum_s alpha[t, s] x_s + D x_t with alpha[t, s] = C_t . (exp(decays_{s+1} + ... + decays_t) * betas_s B_s). The served values rebuild it exactly (float32, as loaded):

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

D = mix._module.D                                 # after the trace: the weights load on the first one
cum = decay.cumsum(0)
causal = torch.ones(len(x), len(x), dtype=torch.bool, device=x.device).tril()[..., None, None]
kept = (cum[:, None] - cum[None, :]).masked_fill(~causal, -torch.inf).exp()
alpha = torch.einsum("tn,tscn,sn->tsc", C, kept, B) * beta     # [query t, key s, channel]
torch.testing.assert_close(torch.einsum("tsc,sc->tc", alpha, x) + D * x, y)

kept is [seq, seq, channels, state_dim], 10 MB for ten tokens on 130m and growing with the square of the prompt: slice the channels on a long one.

The state

state_output is the state after the call's last token, [batch, 1536, 16], and equals states[:, -1]; states holds every token's and needs the routing. A write to the state at token t is not all the block carries forward: the width-4 convolution mixes each token's input with the three before it. On 130m, zeroing every block's state after token 4 of a ten-token prompt and running the last five tokens alone give logits up to 87 apart and different top tokens at three of the five positions.

On a decode step the update comes before the read, the reverse of the prompt's scan, so under generate read state_output before attention_head_outputs on every step after the first. The wrong order raises nothing: step 1's state_output comes back as step 2's.

heads, leaving = [], []
with model.generate(prompt, max_new_tokens=3, do_sample=False) as tracer:
    for step in tracer.iter[:3]:
        if step == 0:                        # the prompt's scan: y, then the state
            heads.append(mix.attention_head_outputs.save())
            leaving.append(mix.state_output.save())
        else:                                # a decode step: the state, then y
            leaving.append(mix.state_output.save())
            heads.append(mix.attention_head_outputs.save())

Edit C and B by assignment

attention_queries and attention_keys are views of torch.split's outputs, and an in-place edit under autograd raises RuntimeError: Output 0 of Select is a view and is being modified inplace. Assign a copy, or edit in place under torch.no_grad(); both reach the scan and give the same logits. attention_values takes in-place edits.

with model.trace(prompt):
    q = mix.attention_queries.clone()
    q[:, -1] = 0                  # the last token reads nothing from the state
    mix.attention_queries = q

The readout and the embeddings

lm_head is embed_tokens' weight; the checkpoint stores no lm_head, so an edit to one is an edit to the other. project_on_vocab at the last block's layer_output equals logits exactly. The final norm's gain is its weight (mean 1.05, up to 5.2 on 130m). Nothing is added after embed_tokens, so token_embeddings equals layers[0].input.

The tokenizer is GPT-NeoX's

The checkpoints use GPT-NeoX-20B's BPE (GPTNeoXTokenizer, 50277 tokens; the embedding has 50280 rows). It prepends no BOS, and <|endoftext|> (id 0) is BOS, EOS and padding. The mixer zeroes padded positions before and after the convolution, so the state stays zero across left padding: a padded prompt's logits match the prompt alone to 2e-4. On 370m config.d_inner reads 160; the mixer's width is intermediate_size, 2048.

The family's checkpoints

This family is the -hf conversions of the original Mamba, 130m to 2.8b. Falcon-Mamba is its own family, falcon_mamba, with RMS norms on B, C and dt inside the mixer; Mamba-2 is mamba2, a different mixer (SSD, one state per head); Jamba's Mamba blocks carry this mixer.

Published work

  • The Hidden Attention of Mamba Models (Ali, Zimerman and Wolf, 2024) reads Mamba's scan as the attention matrix above, built from attention_queries, attention_keys, betas and decays.
  • Locating and Editing Factual Associations in Mamba (Sen Sharma, Atkinson and Bau, 2024) localizes factual recall with causal tracing over Mamba's hidden states, layer_output here, and inserts facts with rank-one edits.

From the family module

04 · nnterp/families/mamba.py

Mamba (MambaForCausalLM): Mamba-1, a pure state-space model.

Its tree is backbone.{embeddings, layers[i].{norm, mixer}, norm_f} plus lm_head. A block is a pre-norm and a selective-scan mixer, nothing else: layer_output = input + mixer(norm(input)), so there is no self_attn and no mlp, and the identity is input + linear_attn.attention_output == layer_output. The mixer answers to linear_attn, the standard name of a recurrent mixer, and carries the SelectiveScan values; the block's norm answers to input_layernorm.

The kernels: with mamba_ssm installed the forward dispatches the scan and the decode step to its compiled CUDA kernels, which have no source to read inside and do not run on CPU; nnterp.route_kernels(model.family, "torch") before the first trace binds them to transformers' pure-torch functions. With residual_in_fp32 the block adds in float32, so layer_output is float32 whatever the load dtype; the model casts the stream to lm_head's dtype and returns float32 logits, and so does the family's project_on_vocab. The config's intermediate_size is the mixer's inner width (expand * hidden_size, the channels of the state); there are no attention heads, so the head sizes have nothing to read.