nnterp · encyclopedia falcon_mamba · nnterp.families.falcon_mamba

family falcon_mamba

Falcon-MambaFalcon3-Mamba

Mamba's block, one norm and a selective-scan mixer, with no attention and no MLP; inside the mixer, weightless RMS norms sit on B, C and the low-rank step input before the scan.

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× 64
→
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 FalconMambaForCausalLMlayoutwhat 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] FalconMambaBlocklayoutwhat 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 FalconMambaMixerlayoutwhat 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 tiiuae/falcon-mamba-7b: the native modules, standard names and every standard value with its layout.

FalconMambaForCausalLM(
  (lm_head): Linear(in_features=4096, out_features=65024, bias=False)
  (generator): Generator(
    (streamer): Streamer()
  )
  (embed_tokens): Embedding(65024, 4096)
  (layers): ModuleList(
    (0-63): 64 x FalconMambaBlock(
      (input_layernorm/norm): FalconMambaRMSNorm((4096,), eps=1e-05)
      (linear_attn/mixer): FalconMambaMixer(
        (conv1d): Conv1d(8192, 8192, kernel_size=(4,), stride=(1,), padding=(3,), groups=8192)
        (act): SiLUActivation()
        (in_proj): Linear(in_features=4096, out_features=16384, bias=False)
        (x_proj): Linear(in_features=8192, out_features=288, bias=False)
        (dt_proj): Linear(in_features=256, out_features=8192, bias=True)
        (out_proj): Linear(in_features=8192, out_features=4096, bias=False)
        (dt_layernorm): FalconMambaWeightlessRMSNorm(eps=1e-06)
        (b_layernorm): FalconMambaWeightlessRMSNorm(eps=1e-06)
        (c_layernorm): FalconMambaWeightlessRMSNorm(eps=1e-06)
        (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): FalconMambaRMSNorm((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_layers64
model.hidden_size4096
model.vocab_size65024
model.intermediate_size8192
architecturesFalconMambaForCausalLM
num_hidden_layers64
hidden_size4096
intermediate_size8192
vocab_size65024
hidden_actsilu
tie_word_embeddingsFalse
dtypebfloat16
layer_norm_epsilon1e-05
state_size16
expand16
conv_kernel4
residual_in_fp32True
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): u, z = in_proj(h)
          u = silu(conv1d(u))              # attention_values
          dt, B, C = x_proj(u)
          dt, B, C = rms(dt), rms(B), rms(C)   # weightless
          y = scan(u, dt_proj(dt), A, B, C) + D * u
          return out_proj(y * silu(z))     # attention_output

The tree, the block and the values are Mamba's (the mamba page): the block's norm is input_layernorm (native norm), the mixer is linear_attn, and model.norm is the final norm_f. What is this family's own sits inside the mixer: dt_layernorm, b_layernorm and c_layernorm, RMS norms with no gain (their weight is a buffer of ones the torch path never reads), with mixer_rms_eps (1e-6 on falcon-mamba-7b, 1e-5 on Falcon3-Mamba).

Where the norms sit

x_proj maps attention_values (the convolved x) to a low-rank step, B and C; the three norms apply to those three outputs, and the scan reads what they return.

  • attention_keys and attention_queries are the normed B and C: each token's vector of state_size (16) entries has a root-mean-square of one.
  • betas is softplus(dt_proj(rms(dt))), dt_proj's bias included: the norm sits on the low-rank step (time_step_rank, 256) before dt_proj, so betas itself carries no fixed scale, and decays is betas * A as on Mamba.
  • The un-normed B and C are in x_proj's output, columns time_step_rank: of it. Read it in a trace of its own. The kernel's values first look up which kernel the call runs, at the top of the mixer's forward, before x_proj runs, and are then read at the kernel call, after it, so a trace that reads x_proj.output and a kernel value fails with OutOfOrderError in either order.
mix = model.layers[1].linear_attn
with model.trace(prompt):
    raw = mix.x_proj.output.save()          # [batch, seq, rank + 2 * state_size]
with model.trace(prompt):
    B = mix.attention_keys.save()           # [batch, seq, 1, state_size]

r, n = model.config.time_step_rank, model.config.state_size
b = raw[..., r:r + n].float()
normed = b * torch.rsqrt(b.pow(2).mean(-1, keepdim=True) + model.config.mixer_rms_eps)
torch.testing.assert_close(normed.to(B.dtype), B[:, :, 0])

A write to attention_keys or attention_queries replaces the normed vector the scan reads, so a written vector need not have unit scale; a write to x_proj.output passes through the norms first.

The contribution is the whole block

attention_output is out_proj's output and all the block adds: layers[i].input + attention_output == layer_output, exactly, in 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)

Loading: route the kernels

With mamba_ssm installed every trace on CPU fails inside the scan kernel (RuntimeError: Expected u.is_cuda() to be true), and on GPU every linear_attn value but attention_output reports the kernel in support(). Routing binds the scan and the decode step to transformers' torch functions in Falcon-Mamba's own modeling module; after it support() reports every value, state and states included.

import nnterp
from nnterp import StandardizedTransformer, route_kernels

route_kernels(nnterp.families.falcon_mamba, "torch")
model = StandardizedTransformer("tiiuae/falcon-mamba-7b")

The mixer's values

On falcon-mamba-7b the mixer is 8192 channels wide, each with a state of 16: attention_values, betas and attention_head_outputs are [batch, seq, 8192], attention_keys and attention_queries [batch, seq, 1, 16], decays and states [batch, seq, 8192, 16]. config.expand reads 16 on the falcon-mamba-7b checkpoints and 2 on Falcon3-Mamba; the mixer's width is intermediate_size, 8192 on both. The scan is Mamba's, so the hidden attention on the mamba page rebuilds attention_head_outputs here unchanged, with the normed B and C.

The state

state_output is [batch, 8192, 16] and equals states[:, -1]. On a decode step read state_output before attention_head_outputs: the step updates the state before it reads it.

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 readout and the embeddings

lm_head has its own weight. project_on_vocab casts the normed stream to the head's dtype and the logits to float32, as the model does, so at the last block's layer_output it equals logits. token_embeddings equals layers[0].input.

The checkpoints and the tokenizer

falcon-mamba-7b and its -instruct and -pre-decay checkpoints, and Falcon3-Mamba-7B -Base and -Instruct, all 64 blocks of width 4096. They tokenize with Falcon's 65024-token BPE and prepend nothing. The tokenizer's end-of-text token is <|endoftext|> (id 11), and it sets no padding token; the -pre-decay and Falcon3-Mamba configs set eos_token_id to 0.

From the family module

04 · nnterp/families/falcon_mamba.py

Falcon-Mamba (FalconMambaForCausalLM): Mamba-1 with normed scan inputs.

Mamba's tree and block: backbone.{embeddings, layers[i].{norm, mixer}, norm_f} plus lm_head, a pre-norm and a selective-scan mixer per block, no attention and no MLP (see nnterp.families.mamba). What differs is inside the mixer: weightless RMS norms (dt_layernorm, b_layernorm, c_layernorm) on the step-size projection, B and C before the scan. The values are read at the kernel call, after those norms, so attention_keys and attention_queries are the normed B and C the scan uses, and the base holds. The kernels are the same functions under the same names, in Falcon-Mamba's own modeling module: nnterp.route_kernels(model.family, "torch") routes them.