nnterp · encyclopedia bamba · nnterp.families.bamba

family bamba

Bamba

Llama's pre-norm block whose mixer is a Mamba-2 state-space scan on all but three blocks and, on those three, attention that rotates half of each head; the block returns a tuple.

HybridState space (Mamba-2)Partial rotaryGrouped-query attention

The block

01 · the residual stream

One block of the stack. Hover any part for its nnterp name; click to pin the card. The slider picks a block, each tick's color gives its block's shape, and the diagram redraws for the block picked.

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

identity layers[i].input + linear_attn.attention_output + mlp.mlp_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 BambaForCausalLMlayoutwhat 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] BambaDecoderLayerlayoutwhat 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 Residual [batch seq hidden_size] The residual stream after the attention add

the child module post_attention_layernorm's input
mlp_norm_output Residual [batch seq hidden_size] The normed stream entering the MLP

the child module post_attention_layernorm's output
mlp_input Residual [batch seq hidden_size] What the MLP is called with

the child module mlp's input
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].mlp BambaMLPlayoutwhat it isread at
mlp_activation Neurons [batch seq intermediate_size] The activation function's output

the child module act_fn'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

the child module down_proj's input
mlp_output Residual [batch seq hidden_size] What the MLP adds to the residual stream

the module's own output
model.layers[i].linear_attn BambaMixerlayoutwhat 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
model.layers[i].self_attn BambaAttentionlayoutwhat it isread at
⚠︎attention_queries Queries [batch num_heads seq qk_head_dim] The queries entering attention

inside the forward: the arguments of attention_interface_1, element 1
⚠︎attention_keys Keys [batch num_kv_heads seq qk_head_dim] The keys entering attention

inside the forward: the arguments of attention_interface_1, element 2
⚠︎attention_values Values [batch num_kv_heads seq head_dim] The values entering attention

inside the forward: the arguments of attention_interface_1, element 3
⚠︎attention_scores Pattern [batch num_heads query key] The attention scores entering the softmax, masked

inside the forward: the input of nn_functional_softmax_0 inside attention_interface_1
attention_premix HeadOutputs [batch seq num_heads head_dim] Every head's output side by side, as the output projection takes it

the child module o_proj's input
attention_output Residual [batch seq hidden_size] What the attention adds to the residual stream

the module's own output
⚠︎attention_probabilities Pattern [batch num_heads query key] The attention pattern the values are mixed with

inside the forward: the output of nn_functional_dropout_0 inside attention_interface_1
⚠︎attention_head_outputs HeadOutputs [batch seq num_heads head_dim] The per-head outputs before the output projection

inside the forward: the output of attention_interface_1, element 0

Printout

What print(model) shows on ibm-ai-platform/Bamba-9B-v2: the native modules, standard names and every standard value with its layout.

BambaForCausalLM(
  (lm_head): Linear(in_features=4096, out_features=128256, bias=False)
  (generator): Generator(
    (streamer): Streamer()
  )
  (embed_tokens): Embedding(128256, 4096, padding_idx=0)
  (layers): ModuleList(
    (0-8): 9 x BambaDecoderLayer(
      (mlp/feed_forward): BambaMLP(
        (gate_proj): Linear(in_features=4096, out_features=14336, bias=False)
        (up_proj): Linear(in_features=4096, out_features=14336, bias=False)
        (down_proj): Linear(in_features=14336, out_features=4096, bias=False)
        (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
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): BambaRMSNorm((4096,), eps=1e-05)
      (post_attention_layernorm/pre_ff_layernorm): BambaRMSNorm((4096,), eps=1e-05)
      (linear_attn/mamba): BambaMixer(
        (act): SiLUActivation()
        (conv1d): Conv1d(8448, 8448, kernel_size=(4,), stride=(1,), padding=(3,), groups=8448)
        (in_proj): Linear(in_features=4096, out_features=16768, bias=False)
        (norm): BambaRMSNormGated()
        (out_proj): Linear(in_features=8192, out_features=4096, 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
      )
      (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
    )
    (9): BambaDecoderLayer(
      (mlp/feed_forward): BambaMLP(
        (gate_proj): Linear(in_features=4096, out_features=14336, bias=False)
        (up_proj): Linear(in_features=4096, out_features=14336, bias=False)
        (down_proj): Linear(in_features=14336, out_features=4096, bias=False)
        (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
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): BambaRMSNorm((4096,), eps=1e-05)
      (post_attention_layernorm/pre_ff_layernorm): BambaRMSNorm((4096,), eps=1e-05)
      (self_attn): BambaAttention(
        (q_proj): Linear(in_features=4096, out_features=4096, bias=False)
        (k_proj): Linear(in_features=4096, out_features=1024, bias=False)
        (v_proj): Linear(in_features=4096, out_features=1024, bias=False)
        (o_proj): Linear(in_features=4096, out_features=4096, bias=False)
        (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
      )
      (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
    )
    (10-17): 8 x BambaDecoderLayer(
      (mlp/feed_forward): BambaMLP(
        (gate_proj): Linear(in_features=4096, out_features=14336, bias=False)
        (up_proj): Linear(in_features=4096, out_features=14336, bias=False)
        (down_proj): Linear(in_features=14336, out_features=4096, bias=False)
        (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
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): BambaRMSNorm((4096,), eps=1e-05)
      (post_attention_layernorm/pre_ff_layernorm): BambaRMSNorm((4096,), eps=1e-05)
      (linear_attn/mamba): BambaMixer(
        (act): SiLUActivation()
        (conv1d): Conv1d(8448, 8448, kernel_size=(4,), stride=(1,), padding=(3,), groups=8448)
        (in_proj): Linear(in_features=4096, out_features=16768, bias=False)
        (norm): BambaRMSNormGated()
        (out_proj): Linear(in_features=8192, out_features=4096, 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
      )
      (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
    )
    (18): BambaDecoderLayer(
      (mlp/feed_forward): BambaMLP(
        (gate_proj): Linear(in_features=4096, out_features=14336, bias=False)
        (up_proj): Linear(in_features=4096, out_features=14336, bias=False)
        (down_proj): Linear(in_features=14336, out_features=4096, bias=False)
        (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
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): BambaRMSNorm((4096,), eps=1e-05)
      (post_attention_layernorm/pre_ff_layernorm): BambaRMSNorm((4096,), eps=1e-05)
      (self_attn): BambaAttention(
        (q_proj): Linear(in_features=4096, out_features=4096, bias=False)
        (k_proj): Linear(in_features=4096, out_features=1024, bias=False)
        (v_proj): Linear(in_features=4096, out_features=1024, bias=False)
        (o_proj): Linear(in_features=4096, out_features=4096, bias=False)
        (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
      )
      (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
    )
    (19-26): 8 x BambaDecoderLayer(
      (mlp/feed_forward): BambaMLP(
        (gate_proj): Linear(in_features=4096, out_features=14336, bias=False)
        (up_proj): Linear(in_features=4096, out_features=14336, bias=False)
        (down_proj): Linear(in_features=14336, out_features=4096, bias=False)
        (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
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): BambaRMSNorm((4096,), eps=1e-05)
      (post_attention_layernorm/pre_ff_layernorm): BambaRMSNorm((4096,), eps=1e-05)
      (linear_attn/mamba): BambaMixer(
        (act): SiLUActivation()
        (conv1d): Conv1d(8448, 8448, kernel_size=(4,), stride=(1,), padding=(3,), groups=8448)
        (in_proj): Linear(in_features=4096, out_features=16768, bias=False)
        (norm): BambaRMSNormGated()
        (out_proj): Linear(in_features=8192, out_features=4096, 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
      )
      (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
    )
    (27): BambaDecoderLayer(
      (mlp/feed_forward): BambaMLP(
        (gate_proj): Linear(in_features=4096, out_features=14336, bias=False)
        (up_proj): Linear(in_features=4096, out_features=14336, bias=False)
        (down_proj): Linear(in_features=14336, out_features=4096, bias=False)
        (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
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): BambaRMSNorm((4096,), eps=1e-05)
      (post_attention_layernorm/pre_ff_layernorm): BambaRMSNorm((4096,), eps=1e-05)
      (self_attn): BambaAttention(
        (q_proj): Linear(in_features=4096, out_features=4096, bias=False)
        (k_proj): Linear(in_features=4096, out_features=1024, bias=False)
        (v_proj): Linear(in_features=4096, out_features=1024, bias=False)
        (o_proj): Linear(in_features=4096, out_features=4096, bias=False)
        (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
      )
      (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
    )
    (28-31): 4 x BambaDecoderLayer(
      (mlp/feed_forward): BambaMLP(
        (gate_proj): Linear(in_features=4096, out_features=14336, bias=False)
        (up_proj): Linear(in_features=4096, out_features=14336, bias=False)
        (down_proj): Linear(in_features=14336, out_features=4096, bias=False)
        (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
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): BambaRMSNorm((4096,), eps=1e-05)
      (post_attention_layernorm/pre_ff_layernorm): BambaRMSNorm((4096,), eps=1e-05)
      (linear_attn/mamba): BambaMixer(
        (act): SiLUActivation()
        (conv1d): Conv1d(8448, 8448, kernel_size=(4,), stride=(1,), padding=(3,), groups=8448)
        (in_proj): Linear(in_features=4096, out_features=16768, bias=False)
        (norm): BambaRMSNormGated()
        (out_proj): Linear(in_features=8192, out_features=4096, 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
      )
      (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): BambaRMSNorm((4096,), 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_layers32
model.hidden_size4096
model.num_heads32
model.num_kv_heads8
model.head_dim128
model.qk_head_dim128
model.vocab_size128256
model.intermediate_size14336
architecturesBambaForCausalLM
num_hidden_layers32
hidden_size4096
intermediate_size14336
num_attention_heads32
num_key_value_heads8
vocab_size128256
max_position_embeddings262144
hidden_actsilu
rms_norm_eps1e-05
tie_word_embeddingsFalse
mamba_n_heads128
mamba_d_head64
mamba_d_state128
mamba_expand2
mamba_chunk_size256
mamba_n_groups1
attention_biasFalse

The Notes

03 · what to know before you run

The block, in order

h   = x + mixer(input_layernorm(x))     # linear_attn (the Mamba-2 mixer) or self_attn
out = h + mlp(pre_ff_layernorm(h))      # SwiGLU, the same class on every block
return out, attention_weights

One block class holds either mixer: config.attn_layer_indices names the attention blocks, 9, 18 and 27 of 32 on every checkpoint, and the other 29 hold a Mamba-2 mixer. config.layer_types says the same per block (linear_attention, full_attention). A block has linear_attn or self_attn, never both; mamba is linear_attn, feed_forward is mlp and pre_ff_layernorm is post_attention_layernorm. Pick blocks outside the trace:

kinds = model.config.layer_types
ssm = [i for i, t in enumerate(kinds) if t == "linear_attention"]
attn = [i for i, t in enumerate(kinds) if t == "full_attention"]     # 9, 18, 27

The contributions are the sublayers' outputs

The block returns a tuple, (hidden_states, attention_weights), and layer_output is its first element. Nothing scales or norms a sublayer's output on its way to the stream, so the identity is the plain sum, exact in float32 on both kinds of block:

with model.trace(prompt):
    x = model.layers[0].input.save()
    mix = model.layers[0].linear_attn.attention_output.save()
    mlp = model.layers[0].mlp.mlp_output.save()
    out = model.layers[0].layer_output.save()

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

On an attention block the middle term is model.layers[9].self_attn.attention_output.

Loading: route the kernels

With mamba_ssm installed, transformers runs the Mamba-2 scans in its CUDA kernels: every linear_attn value but attention_output is unavailable, and on a CPU no trace runs at all (RuntimeError: invalid argument to exchangeDevice, even for a read of layer_output). nnterp.route_kernels(model.family, "torch") before the first trace binds transformers' pure-torch scans. states and state_after also need nnterp.chunk_per_token(model): the scan keeps the state every mamba_chunk_size tokens, 256 on every checkpoint.

import nnterp

nnterp.route_kernels(nnterp.families.bamba, "torch")       # before the first trace
model = StandardizedTransformer("ibm-ai-platform/Bamba-9B-v2", attn_implementation="eager")

The attention interior on the three attention blocks needs attn_implementation="eager"; the Mamba-2 values are read at the scan's call and need no eager load.

The Mamba-2 mixer

Every checkpoint has 128 heads of 64 channels in one group, and a state of 128 per head: attention_queries (C) and attention_keys (B) are [batch, seq, 1, 128], one for all 128 heads, attention_values (x) and attention_head_outputs (y) are [batch, seq, 128, 64], betas (dt) and decays (A * dt) are [batch, seq, 128], and the state is [batch, 128, 128, 64], key side first. A width-4 convolution runs over x, B and C before the scan. After it, y is multiplied by silu(z) and RMS-normed over all 8192 channels at once (BambaRMSNormGated), then out_proj maps it to the stream; attention_head_outputs is read before both.

mix = model.layers[0].linear_attn
with model.trace(prompt):
    C = mix.attention_queries.save()          # [batch, seq, groups, state_dim]
    x = mix.attention_values.save()           # [batch, seq, heads, head_dim]
    dt = mix.betas.save()                     # [batch, seq, heads]
    state = mix.state_output.save()           # [batch, heads, state_dim, head_dim]

A zero step skips a token exactly

time_step_limit is (0, inf) on every checkpoint, so the scan's clamp on dt changes nothing, and a written betas of zero runs as zero. betas[:, t] = 0 is then the exact "skip this token": no write and no decay, so the state after token t equals the state after t - 1 bit for bit. Assign betas; an in-place edit does not reach the scan.

nnterp.chunk_per_token(model)
with model.trace(prompt):
    betas = mix.betas.clone()
    betas[:, 3] = 0
    mix.betas = betas
    states = mix.states.save()                # [batch, seq, heads, state_dim, head_dim]

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

Attention rotates half of each head

The three attention blocks apply rotary embeddings to the first 64 of each head's 128 dimensions (partial_rotary_factor 0.5, which the config class sets for every checkpoint; attn_rotary_emb is 64), with rope_theta 10000. The other 64 carry no position: attention_queries[..., 64:] is q_proj's output split into heads, unchanged. There are 32 query heads over 8 key/value heads, so an edit to attention_values[:, j] reaches query heads 4j to 4j + 3; the score scale is 128 ** -0.5.

attn = model.layers[9].self_attn
with model.trace(prompt):
    q_raw = attn.source.self_q_proj_0.output.save()
    q = attn.attention_queries.save()              # [batch, num_heads, seq, head_dim]

b, s, _ = q_raw.shape
q_raw = q_raw.view(b, s, -1, model.head_dim).transpose(1, 2)
half = model.head_dim // 2
assert torch.equal(q_raw[..., half:], q[..., half:])

The readout, the embeddings 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 exactly. lm_head and embed_tokens are separate weights, and nothing is added after embed_tokens, so token_embeddings equals layers[0].input. The tokenizer is Llama 3's (128256 tokens): it prepends <|begin_of_text|> (id 128000), so position 0 is that token, and <|end_of_text|> (128001) ends a text. Bamba-9B-v2's config names ids 1 and 2 as bos_token_id and eos_token_id; the tokenizer's are 128000 and 128001. The checkpoints are base models and ship no chat template.

The checkpoints

Bamba-9B-2T is the end of the first pretraining stage (2T tokens) and Bamba-9B-1.8T a checkpoint inside it; Bamba-9B-v1 (Bamba-9B on the Hub is the same repository) is the end of the second stage, and Bamba-9B-v2 is trained on 1T tokens more, as their model cards say. All four share one shape: 32 blocks of width 4096, an MLP of 14336, and the attention blocks at 9, 18 and 27. max_position_embeddings is 4096 on v1 and 262144 on the others (1.8T and 2T leave it to the config class's default). GraniteMoE-Hybrid (Granite 4.0-H) is Bamba's block with Granite's multipliers, and is the granitemoehybrid family.

From the family module

04 · nnterp/families/bamba.py

Bamba (BambaForCausalLM).

Llama's containers over one block class that holds either a Mamba-2 mixer or attention, by config.layers_block_type (attn_layer_indices name the attention blocks)::

model.layers[i]          BambaDecoderLayer
    .input_layernorm
    .mamba               BambaMixer, an SSD mixer   (or .self_attn, BambaAttention)
    .pre_ff_layernorm
    .feed_forward        BambaMLP

A pre-norm block, as Llama: the mixer's output and then the MLP's are added to the stream as the modules return them. The names are Llama's but for three: mamba is linear_attn (nnterp.StateSpace), feed_forward is mlp and pre_ff_layernorm is post_attention_layernorm. Each block has self_attn or linear_attn, never both. The block returns (hidden_states, attention_weights).