View as Markdown

Mamba-2¶

On Mamba-2, Nemotron-H, Bamba and Falcon-H1 the recurrent mixer layers[i].linear_attn is a StateSpace, a Mamba-2 (SSD) mixer: linear attention with one scalar decay per head. Per head, with the state a head_dim by state_dim matrix, each token does h = exp(dt * A) * h + dt * x B^T and y = h C + D * x. The shared names mean what they mean on a gated DeltaNet: C is attention_queries, B is attention_keys, x is attention_values, dt is betas, A * dt is decays and y is attention_head_outputs. AntonV/mamba2-130m-hf is the smallest public checkpoint. On it every block is the mixer alone. Nemotron-H and Bamba mix such blocks with attention and MLP blocks, and Falcon-H1 runs a mixer and attention side by side on every block.

import torch
import nnterp
from nnterp import StandardizedTransformer, route_kernels

route_kernels(nnterp.families.mamba2, "torch")   # before any trace: mamba_ssm's kernels have no source and need CUDA
model = StandardizedTransformer("AntonV/mamba2-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))
print(decays.dtype, decays.max().item() <= 0, dt.min().item() > 0)
# (1, 10, 1, 128) (1, 10, 1, 128) (1, 10, 24, 64) (1, 10, 24) (1, 10, 24)
# (1, 10, 24, 64) (1, 24, 128, 64) (1, 10, 768)
# torch.float32 True True

route_kernels comes first, as on the DeltaNet page: with mamba_ssm installed transformers calls its compiled kernels, which have no source to read and do not run on CPU. Unrouted, a CPU trace fails inside the kernel (Pointer argument cannot be accessed from Triton when a GPU is visible). On Mamba-2 only the two scan kernels are rerouted, and a prompt keeps the chunked scan.

The values¶

value	SSD	shape	what it is
`attention_output`		`[batch, seq, hidden_size]`	what the mixer adds to the residual stream
`attention_queries`	`C`	`[batch, seq, groups, state_dim]`	what reads the state; one per group of heads
`attention_keys`	`B`	`[batch, seq, groups, state_dim]`	where each token writes into the state; one per group
`attention_values`	`x`	`[batch, seq, heads, head_dim]`	what each token writes
`betas`	`dt`	`[batch, seq, heads]`	`softplus(dt + dt_bias)`, clamped to `time_step_limit` on a prompt: the write strength
`decays`	`A * dt`	`[batch, seq, heads]`, float32	the log of how much of the state each token keeps, non-positive
`state_input`	`h` in	`[batch, heads, state_dim, head_dim]`	the state the call starts from: `None` on a fresh prompt, a copy of the cached state on a decode step
`state_output`	`h` out	`[batch, heads, state_dim, head_dim]`	the state after the call's last token
`attention_head_outputs`	`y`	`[batch, seq, heads, head_dim]`	each head's read of the state plus the `D` skip, before the gated norm and `out_proj`

groups is the mixer's n_groups, 1 here (the heads of a group share B and C), heads its num_heads (24), head_dim 64 and state_dim its ssm_state_size (128). The state is served key side first, the layout every recurrent mixer's state shares and the transpose of the cache's buffer. An assignment is transposed back. Assign attention_queries, attention_keys and attention_values rather than editing them in place, as on the other recurrent mixers. Signatures are on the state_space page of the Documentation tab.

`betas` and `decays` are one argument¶

Both are the kernel's dt seen through dt_bias, the softplus and A. An assignment is carried back into dt, and the kernel computes the gate that was written. An in-place edit (mix.betas[:, t] = 0) does not reach dt.

prompt = "The Eiffel Tower is in the city of"
with model.trace(prompt):
    mix.decays = mix.decays * 0.5                 # half the log decay is half of dt: half-strength writes too
    halved = mix.betas.save()

with model.trace(prompt):
    betas = mix.betas.clone()
    betas[:, 3] = 0                               # token 3 writes nothing and decays nothing
    mix.betas = betas
    skipped = model.logits[0, -1].save()

with model.trace(prompt):
    mix.betas = mix.betas                         # the round trip through the softplus's inverse rounds dt
    written_back = model.logits[0, -1].save()

with model.trace(prompt):
    clean = model.logits[0, -1].save()

print(torch.allclose(halved, dt * 0.5, atol=1e-6))
print((skipped - clean).abs().max().item() > 0, (written_back - clean).abs().max().item())
# True
# True 0.000125885009765625

Writing decays rewrites betas

There is no way to change the decay alone. decays * 0.5 halves dt: each token keeps more of the state and writes at half strength. decays = 0 ("keep everything") sets dt to zero, which writes nothing, and the state stays at zero through the whole prompt.

betas[:, t] = 0 is the exact "skip this token": no write and exp(0) = 1, no decay, and the state after token t equals the state after t - 1. Nemotron-H's scan clamps dt to time_step_min first, and there the state still moves. A written betas must be positive or zero, and on a prompt the kernel clamps it to time_step_limit again. Writing the same values back is not a no-op: the round trip rounds dt, by 1.3e-4 in the last logit here. Compare an edit against an unedited run, not against a write-back.

The state after every token¶

The chunk scan computes the state at every chunk boundary in one tensor, and between boundaries there is nothing to read:

from nnterp import chunk_per_token

print(model.support(layer=0)["linear_attn.states"])
chunk_per_token(model)                            # this model only; undo with chunk_per_token(model, False)
with model.trace(prompt):
    states = mix.states.save()                    # [batch, seq, heads, state_dim, head_dim]
    final = mix.state_output.save()

with model.trace(prompt):
    betas = mix.betas.clone()
    betas[:, 3] = 0
    mix.betas = betas
    held = mix.states.save()

print(tuple(states.shape), torch.equal(states[:, -1], final), torch.equal(held[:, 3], held[:, 2]))
print(model.support(layer=0)["linear_attn.set_state_after"])
chunk_per_token(model, False)
# the chunk scan materializes the state only at chunk boundaries, every 256 tokens (chunk_size=256); call nnterp.chunk_per_token(model) to set every mixer's chunk_size to 1, so every token is a boundary (slower on long prompts)
# (1, 10, 24, 128, 64) True True
# 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

nnterp.chunk_per_token(model) sets every Mamba-2 mixer's chunk_size to 1. Every token is then a boundary, and states and state_after(t) read the state after every token of the call. The call touches this model only, because the forward reads the attribute on every call, where route_kernels rebinds a family's kernels for the process. chunk_per_token(model, False) restores the chunk size each mixer was built with. The recurrence between chunks is quadratic in their number, and a long prompt runs markedly slower. The same scan in another order also moves the logits by rounding. states is a read-only copy. state and set_state_after stay unavailable either way, with the reasons support() prints: the scan has no per-token occurrence for tracer.iter to walk, as a DeltaNet has, and it computes every boundary state in one cumulative step from the initial state, in which a state written at one token does not flow into the next. To change where a call starts, assign state_input.

The state between steps¶

A prompt runs mamba2_chunk_scan and each decode step runs mamba2_selective_state_update, which takes its arguments in other places and has no sequence axis. The values select from the kernel that fires, and a decode step's tensors are served with a sequence axis of 1.

entering, leaving = [], []
with model.generate(prompt, max_new_tokens=3, do_sample=False) as tracer:
    for step in tracer.iter[:3]:
        s = mix.state_input
        entering.append(s.save() if s is not None else None)
        leaving.append(mix.state_output.save())

print(entering[0] is None, torch.equal(leaving[0], entering[1]), torch.equal(leaving[1], entering[2]))
# True True True

The update writes the new state into the cache and returns only y, which is where a decode step's state_output is read. Read it before attention_head_outputs on a decode step, the trap the Mamba page describes. On Falcon-H1-0.5B and the Falcon-H1-Tiny checkpoints, which set mamba_rms_norm off, a decode step's attention_head_outputs is gated by silu(z) where a prompt's is not.

Mamba-2¶

The values¶

betas and decays are one argument¶

The state after every token¶

The state between steps¶

`betas` and `decays` are one argument¶