nnterp · encyclopedia falcon_h1 · nnterp.families.falcon_h1

family falcon_h1

Falcon-H1

Every block runs a Mamba-2 mixer and attention side by side on one normed input and adds both, each times its own multiplier, then a gated MLP; the embeddings and the logits carry multipliers too.

Parallel mixersState space (Mamba-2)Scaled residual addsEmbedding multiplierScaled logitsGrouped-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.

embed_tokenstoken_embeddings
→
layers× 36
→
normfinal_norm_output
→
lm_headlm_head_output
→
logitsscaled
i = 0

identity layers[i].input + (linear_attn.attention_output + self_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 FalconH1ForCausalLMlayoutwhat 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] FalconH1DecoderLayerlayoutwhat 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 FalconH1MLPlayoutwhat 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 FalconH1Mixerlayoutwhat it isread at
attention_output Residual [batch seq hidden_size] What the Mamba-2 mixer adds to the residual stream: its output times ssm_out_multiplier

the sibling module source.mamba_hidden_states_1's 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 FalconH1Attentionlayoutwhat 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: its output times attention_out_multiplier

the sibling module source.attention_hidden_states_0's 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 tiiuae/Falcon-H1-0.5B-Base: the native modules, standard names and every standard value with its layout.

FalconH1ForCausalLM(
  (lm_head): Linear(in_features=1024, out_features=32784, bias=False)
  (generator): Generator(
    (streamer): Streamer()
  )
  (embed_tokens): Embedding(32784, 1024, padding_idx=0)
  (layers): ModuleList(
    (0-35): 36 x FalconH1DecoderLayer(
      (mlp/feed_forward): FalconH1MLP(
        (gate_proj): Linear(in_features=1024, out_features=2048, bias=False)
        (up_proj): Linear(in_features=1024, out_features=2048, bias=False)
        (down_proj): Linear(in_features=2048, out_features=1024, 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
      )
      (linear_attn/mamba): FalconH1Mixer(
        (act): SiLUActivation()
        (conv1d): Conv1d(1792, 1792, kernel_size=(4,), stride=(1,), padding=(3,), groups=1792)
        (in_proj): Linear(in_features=1024, out_features=3352, bias=False)
        (norm): Identity()
        (out_proj): Linear(in_features=1536, out_features=1024, bias=False)
        (attention_output) -> Residual [batch seq hidden_size]: What the Mamba-2 mixer adds to the residual stream: its output times ssm_out_multiplier
        (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
      )
      (self_attn): FalconH1Attention(
        (q_proj): Linear(in_features=1024, out_features=512, bias=False)
        (k_proj): Linear(in_features=1024, out_features=128, bias=False)
        (v_proj): Linear(in_features=1024, out_features=128, bias=False)
        (o_proj): Linear(in_features=512, out_features=1024, bias=False)
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream: its output times attention_out_multiplier
        (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_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
      )
      (input_layernorm): FalconH1RMSNorm((1024,), eps=1e-05)
      (post_attention_layernorm/pre_ff_layernorm): FalconH1RMSNorm((1024,), eps=1e-05)
      (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): FalconH1RMSNorm((1024,), 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_layers36
model.hidden_size1024
model.num_heads8
model.num_kv_heads2
model.head_dim64
model.qk_head_dim64
model.vocab_size32784
model.intermediate_size2048
architecturesFalconH1ForCausalLM
num_hidden_layers36
hidden_size1024
intermediate_size2048
num_attention_heads8
num_key_value_heads2
head_dim64
vocab_size32784
max_position_embeddings16384
hidden_actsilu
rms_norm_eps1e-05
tie_word_embeddingsFalse
dtypebfloat16
embedding_multiplier5.656854249492381
mamba_n_heads24
mamba_d_head64
mamba_d_state128
mamba_expand2
mamba_chunk_size128
mamba_n_groups1
mamba_rms_normFalse
ssm_in_multiplier1.25
ssm_out_multiplier0.23570226039551587
attention_in_multiplier1.0
attention_out_multiplier0.9375
key_multiplier0.39062499999999994
lm_head_multiplier0.0390625
attention_biasFalse

The Notes

03 · what to know before you run

The block, in order

n   = input_layernorm(x)
h   = x + (mamba(n) * ssm_out_multiplier
           + self_attn(n * attention_in_multiplier) * attention_out_multiplier)
out = h + feed_forward(pre_ff_layernorm(h))

Every block has both mixers: linear_attn (native mamba, a Mamba-2 mixer) and self_attn read the same normed input, the Mamba-2 mixer runs first, and the block adds the two at once. The MLP (native feed_forward, nnterp's mlp) then reads the stream after that add through pre_ff_layernorm, which nnterp names post_attention_layernorm. The block returns a one-element tuple; layer_output is the tensor. config.layer_types is hybrid on every block.

The contributions are the scaled outputs

Each mixer's attention_output is the block's own product, the module's output times ssm_out_multiplier or attention_out_multiplier: the tensor the block adds, so a read, an assignment or an in-place edit needs no arithmetic. linear_attn.output and self_attn.output[0] are the unscaled module outputs. The identity has four terms:

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

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

On 0.5B, over the positions after the first, the Mamba-2 mixer adds a larger vector than the attention in 33 of 36 blocks (a median norm ratio of 2.2), and the MLP's is about 3.4 times the attention's.

Every multiplier, and where the values read it

The configs carry µP (maximal update parametrization) multipliers; nnterp reads each value where the forward has already applied them.

  • ssm_out_multiplier, attention_out_multiplier: inside each mixer's attention_output (0.236 and 0.9375 on 0.5B; 0.088 and 0.0375 on 34B).
  • attention_in_multiplier: the attention reads input_layernorm's output times it, 1.0 on every released checkpoint, so self_attn.input is the norm's output.
  • key_multiplier: k_proj's output is multiplied by it before the rotary, so attention_keys carry it (0.39 on 0.5B).
  • ssm_in_multiplier and ssm_multipliers: the mixer scales its input, then each part of in_proj's output (z, x, B, C, dt), before the convolution; attention_queries, attention_keys, attention_values and betas are read after them.
  • mlp_multipliers: the MLP scales the gate's pre-activation by the first and down_proj's output by the second, so mlp_output (the MLP's return) carries both, while mlp.down_proj.output is the product before the second.
  • embedding_multiplier and lm_head_multiplier: on the strip above.

Loading

The attention interior needs attn_implementation="eager". The default sdpa load computes the same attention: on 0.5B the eager and sdpa logits differ by at most 2.3e-5 in float32 and 0.5 in bfloat16, with the same top token at every position.

With mamba_ssm installed transformers sends the Mamba-2 chunk scan and decode update to its Triton kernels: every linear_attn value but attention_output reports the kernel in support(), and on a CPU no trace runs. Route before the first trace. The routed and the default kernels give logits within 0.019 of each other in float32 (logits up to 46) and 0.5 in bfloat16 on 0.5B, with the same top tokens.

import nnterp
from nnterp import StandardizedTransformer, route_kernels

route_kernels(nnterp.families.falcon_h1, "torch")
model = StandardizedTransformer("tiiuae/Falcon-H1-0.5B-Base", attn_implementation="eager")

The Mamba-2 mixer

On 0.5B the mixer has 24 heads of 64 channels and a state of 128 per head, in one group: attention_queries and attention_keys are [batch, seq, 1, 128], attention_values and attention_head_outputs [batch, seq, 24, 64], and the state [batch, 24, 128, 64]. states needs nnterp.chunk_per_token(model): the scan keeps the state every mamba_chunk_size tokens (128). The family module's docstring and docs/usage/state-space.md give the rest.

mamba_rms_norm decides what follows the scan. It is false on 0.5B and the Tiny checkpoints: linear_attn.norm is an identity, y is gated by silu(z) and goes to out_proj, and on a decode step the gate is applied inside the update kernel, so a decode step's attention_head_outputs is gated where a prompt's is not. It is true on 1.5B, 1.5B-Deep, 3B, 7B, 34B and H1R-7B: a gated RMSNorm follows the scan, and attention_head_outputs is ungated on a prompt and on a decode step alike.

Position 0 carries a very large norm

On 0.5B the stream at the first position has a norm of 570 to 920 after block 9, where the other positions' median is about 60; the gap closes in the last blocks (214 to 545 at position 0 after block 27, about 250 elsewhere). Block 18's attention puts more than half of its mass on position 0, averaged over heads and queries (0.58 to 0.85 on five test texts). The tokenizer prepends nothing, so position 0 is the prompt's own first token. Leave it out before averaging activations for steering vectors or probes.

with model.trace(prompt):
    resid = model.layers[9].layer_output.save()

resid[0].norm(dim=-1)        # 0.5B: 570 to 920 at position 0, about 60 elsewhere

Attention

Every block's attention applies rotary embeddings with rope_theta 1e11: of the 32 frequency pairs of a 64-wide head, about ten turn once within 0.5B's 16384-token context, and the rest barely move. 0.5B has 8 query heads over 2 key/value heads of 64: query head h reads key/value head h // 4. attention_keys are scaled by key_multiplier before the rotary.

The readout and the embeddings

logits is lm_head.output * lm_head_multiplier (0.039 on 0.5B), and project_on_vocab applies the multiplier, so at the last block's layer_output it equals logits. The final norm's gain is large (mean 16.2, up to 106 on 0.5B) and the head's raw output reaches 1179 where logits stops at 46. layers[0].input is token_embeddings * embedding_multiplier (5.66 on 0.5B to 34B; 0.08 to 0.11 on the Tiny checkpoints).

The checkpoints and the tokenizer

0.5B, 1.5B, 1.5B-Deep (66 blocks of width 1280), 3B, 7B and 34B, each Base and Instruct; Falcon-H1R-7B, a reasoning model on 7B's shape; and the Falcon-H1-Tiny 90M checkpoints, whose mixers' and attention's multipliers are all 1. The vocabulary is 32768 tokens on 0.5B and grows with size (261120 on 34B). The tokenizer prepends no BOS on a plain prompt; <|end_of_text|> (id 11) ends a sequence. The Instruct chat template opens with <|begin_of_text|> and uses <|im_start|> and <|im_end|> turns.

From the family module

04 · nnterp/families/falcon_h1.py

Falcon-H1 (FalconH1ForCausalLM).

Llama's containers over a parallel hybrid block: every block runs a Mamba-2 mixer and attention side by side on the same normed input, then an MLP::

model.layers[i]          FalconH1DecoderLayer
    .input_layernorm
    .mamba               FalconH1Mixer, an SSD mixer
    .self_attn           FalconH1Attention
    .pre_ff_layernorm
    .feed_forward        FalconH1MLP

h = x + mamba(n) * ssm_out_multiplier + self_attn(n * attention_in_multiplier) * attention_out_multiplier
out = h + feed_forward(pre_ff_layernorm(h))          with n = input_layernorm(x)

mamba is linear_attn (nnterp.StateSpace), feed_forward is mlp and pre_ff_layernorm is post_attention_layernorm. Every block has both mixers, so the contribution identity has four terms, input + self_attn.attention_output + linear_attn.attention_output + mlp_output == layer_output.

The config's µP multipliers, and where each is handled:

  • ssm_out_multiplier / attention_out_multiplier: the block scales each mixer's output before adding it, so attention_output on each is that product, the block's own binding of it (mamba_hidden_states_1, attention_hidden_states_0): the tensor the block adds, so reads, writes and in-place edits need no arithmetic.
  • lm_head_multiplier: the model multiplies the head's output by it, so the family defines project_on_vocab with that step.
  • embedding_multiplier: the model multiplies the embedding module's output by it before the first block, so token_embeddings (the module's output) times the multiplier is layers[0].input.
  • attention_in_multiplier: the attention's input is the normed stream times it (1.0 on the released checkpoints), so self_attn.input is the norm's output only then.
  • mlp_multipliers, key_multiplier, ssm_in_multiplier, ssm_multipliers: applied inside the MLP, the attention and the mixer; the values read where the forward already applied them.

With mamba_rms_norm false the decode step passes the gate into the update kernel, so a decode step's attention_head_outputs is gated by silu(z) where a prompt's is not; with it true (Falcon-H1 1.5B and up, H1R-7B) both are ungated. Falcon-H1-0.5B and the Falcon-H1-Tiny checkpoints set it false, and their linear_attn.norm is an identity. The block returns a one-element tuple.