Mamba¶
On Mamba and Falcon-Mamba every block is a norm and a Mamba-1 mixer, layers[i].linear_attn, a SelectiveScan. Jamba puts one on every block except the attention blocks that attn_layer_period places. Per channel of its inner width the mixer keeps a state of state_dim numbers and runs, for each token, h = exp(dt * A) * h + dt * B * x and y = C . h + D * x. nnterp serves the scan's tensors under the recurrent mixer's shared names: C reads the state as a query does, B says where a token writes as a key does, x is what it writes, dt is the write strength (betas) and dt * A the log decay (decays). state-spaces/mamba-130m-hf is the smallest public checkpoint.
import torch
import nnterp
from nnterp import StandardizedTransformer, route_kernels
route_kernels(nnterp.families.mamba, "torch") # before any trace: mamba_ssm's kernels have no source and need CUDA
model = StandardizedTransformer("state-spaces/mamba-130m-hf", device="cpu")
mix = model.layers[0].linear_attn
with model.trace("The Eiffel Tower is in the city of"):
C = mix.attention_queries.save()
B = mix.attention_keys.save()
x = mix.attention_values.save()
dt = mix.betas.save()
decays = mix.decays.save()
y = mix.attention_head_outputs.save()
state = mix.state_output.save()
out = mix.attention_output.save()
print(tuple(C.shape), tuple(B.shape), tuple(x.shape), tuple(dt.shape), tuple(decays.shape))
print(tuple(y.shape), tuple(state.shape), tuple(out.shape))
# (1, 10, 1, 16) (1, 10, 1, 16) (1, 10, 1536) (1, 10, 1536) (1, 10, 1536, 16)
# (1, 10, 1536) (1, 1536, 16) (1, 10, 768)
route_kernels comes first because transformers calls mamba_ssm's compiled kernels when that package is installed, and they have no Python source to read inside and do not run on CPU: unrouted, a CPU trace fails inside the kernel with Expected u.is_cuda() to be true, even when it reads only layer_output. What the call does, and why a recurrent mixer's values need it, is on the DeltaNet page. On Mamba-1 the pure-torch scan is itself a token loop, and each kernel name is bound to its own pure-torch function. A Mamba block has no self_attn and no mlp: support() lists neither, and layers[i].input + linear_attn.attention_output == layer_output. Jamba's blocks are Llama's, with one mixer each.
The values¶
| value | shape | what it is |
|---|---|---|
attention_output |
[batch, seq, hidden_size] |
what the mixer adds to the residual stream |
attention_queries |
[batch, seq, 1, state_dim] |
C, what each token reads the state with; one group, shared by every channel |
attention_keys |
[batch, seq, 1, state_dim] |
B, what each token writes into the state along |
attention_values |
[batch, seq, channels] |
x, the input each channel writes, after the convolution and the activation |
betas |
[batch, seq, channels] |
softplus(dt + dt_bias), the step size; derived, read-only |
decays |
[batch, seq, channels, state_dim] |
dt * A, the log decay per channel and state dimension; derived, read-only |
state_input |
[batch, channels, state_dim] |
a decode step's cached state, a copy; None on a prompt, whose scan starts from zeros |
attention_head_outputs |
[batch, seq, channels] |
y = C . h + D * x, before the silu(z) gate and out_proj |
state_output |
[batch, channels, state_dim] |
the state after the call's last token, as the cache receives it |
channels is the mixer's intermediate_size (1536 here) and state_dim its ssm_state_size (16). The kernel takes these tensors channel-first; each value is a tokens-first view of its argument, laid back out on assignment. Unlike a DeltaNet's, decays has a state axis, because A is per channel and state dimension. betas and decays are derived from the kernel's arguments and read-only. On Falcon-Mamba and Jamba B, C and the step size pass through RMS norms before the scan, and the values are the normed ones. Signatures are on the selective_scan page of the Documentation tab.
The state¶
The pure-torch scan loops over the tokens, and once routed the state after every token of a prompt exists with no second kernel: states stacks them, state is one occurrence per token for tracer.iter to read or write, and state_after(t) and set_state_after(t, value) are the same two operations as calls:
prompt = "The Eiffel Tower is in the city of"
with model.trace(prompt):
states = mix.states.save() # [batch, seq, channels, state_dim]
final = mix.state_output.save()
with model.trace(prompt) as tracer:
for t in tracer.iter[2]:
mix.state = torch.zeros_like(final) # tokens 3.. continue from an empty state
for t in tracer.iter[3]:
s3 = mix.state.save()
print(tuple(states.shape), torch.equal(states[:, -1], final))
print(round(s3.norm().item(), 2), round(states[:, 3].norm().item(), 2))
# (1, 10, 1536, 16) True
# 3.66 10.6
A write at token t changes token t's own output too, since y_t reads h_t, and the convolution in front of the scan (width 4) still carries the three tokens before t into the three after it, as on a DeltaNet. Under generate the prompt runs the scan and each decode step runs mamba_selective_state_update over one token, the state handed from state_output to the next step's state_input:
with model.generate(prompt, max_new_tokens=2, do_sample=False) as tracer:
for step in tracer.iter[1]: # a decode step: the state is updated first, then read
entering = mix.state_input.save()
leaving = mix.state_output.save()
y1 = mix.attention_head_outputs.save()
print(tuple(entering.shape), tuple(y1.shape))
# (1, 1536, 16) (1, 1, 1536)
Read order on a decode step
The scan computes y before it returns the state, but the decode kernel updates the state first and reads y from it. On a decode step read state_output before attention_head_outputs, as above. The other order does not raise: state_output binds the next step's state, and only on the last step is the loop cut short, with a was never reached warning that blames the loop.
Assign, do not edit in place¶
C and B are views that torch refuses to edit in place while autograd is on. Assign them, or edit in place under torch.no_grad(). attention_values and attention_head_outputs take in-place edits either way, and a write to state_output changes what the next decode step starts from, not this call's output.
with model.trace(prompt):
mix.attention_queries = mix.attention_queries * 0 # assign: C is a view of the kernel's argument
y_noC = mix.attention_head_outputs.save()
with model.trace(prompt):
with torch.no_grad():
mix.attention_queries[:, -1] = 0 # or edit in place without autograd
y_noC_last = mix.attention_head_outputs.save()
print(torch.allclose(y_noC[:, -1], y_noC_last[:, -1]))
try:
with model.trace(prompt):
mix.attention_queries[:, -1] = 0
except RuntimeError as e:
print(str(e)[:60])
try:
with model.trace(prompt):
mix.betas = mix.betas * 0
except AttributeError as e:
print(e)
# True
# Output 0 of Select is a view and is being modified inplace.
# betas is derived and read-only
Mamba-2's mixer has the same names and a different state: Mamba-2.