state_space¶
StateSpace: a Mamba-2 (SSD) mixer's values, at the scan kernel call.
module-attribute
¶
SSDHeadOutputs = Float[Tensor, f'{BATCH} {SEQ} {MIXER_HEADS} {HEAD_DIM}']
module-attribute
¶
NO_STATE_OCCURRENCES = "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)"
module-attribute
¶
NO_STATE_WRITES = "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"
Bases: RecurrentMixer
A Mamba-2 (SSD) mixer (Mamba-2, Nemotron-H, Bamba, Falcon-H1, Zamba2, GraniteMoeHybrid): a selective state space.
SSD, the state-space dual, is linear attention with a scalar decay per
head. Per head, with the state h a head_dim by state_dim
matrix, each token does::
h = exp(dt * A) * h + dt * x B^T
y = h C + D * x
so the shared names mean what they mean on a gated DeltaNet
(LinearAttention):
attention_queriesisC, what reads the state;attention_keysisB, where a token writes into it; both[batch, seq, groups, state_dim], one per group of heads (n_groups), shared by the heads in the group.attention_valuesisx(hidden_statesat the call), what is written,[batch, seq, heads, head_dim].betasisdt = softplus(dt + dt_bias), the write strength,[batch, seq, heads];decaysisA * dt, the log of how much of the state each token keeps,[batch, seq, heads], float32 and non-positive. Both are the kernel'sdtargument seen throughdt_bias, the softplus andA; an assignment is carried back intodt(_dt_for), so the kernel computes the assigned gate.state_input/state_output: the state entering and leaving the call, served as[batch, heads, state_dim, head_dim](key side first, the sharedStatelayout), the transpose of the cache's[batch, heads, head_dim, state_dim]; an assignment is transposed back.state_inputisNoneon a fresh prompt.attention_head_outputsisy, the scan's output with theDskip, before the gated norm andout_proj,[batch, seq, heads, head_dim].
Everything but attention_output is read at the scan call:
transformers' mamba2_chunk_scan on a prompt, its
mamba2_selective_state_update on a decode step (RecurrentMixer.KERNEL). The two
take their arguments in different places and the update has no sequence
axis, so each value selects by the kernel that fires (argument) and a
decode step's tensors are served with a sequence axis of 1, removed again
on assignment. The update writes the new state into the cache's buffer
and returns only y, so a decode step's state_output is read inside
it, at UPDATED_STATE.
The chunk scan materializes the state at every chunk boundary, in one
tensor (new_states, CHUNK_STATES); with a chunk size of 1
(chunk_per_token) every token is a boundary, and states and
state_after(t) read the state after every token of the call (a decode
step's is its state_output). There is no per-token occurrence to walk
with tracer.iter (state is unavailable, STATE_OP is None),
and no per-token write: the scan computes every boundary state in one
cumulative step from the initial state, so set_state_after is
unavailable too. With mamba_ssm installed the kernels have no Python
source; route_kernels(model.family, "torch") binds transformers'
pure-torch ones (RecurrentMixer).
A value that needs another argument of the call (betas needs
dt_bias, a prompt's state_output whether the scan returns a
state) reads the call's arguments once per call (_arguments), at the
call's first need, and finds them there after the model has moved into
the kernel. A value read after the model has moved past its location is
an out-of-order read.
class-attribute
instance-attribute
¶
class-attribute
instance-attribute
¶
CHUNK_ARGUMENTS = {'hidden_states': 0, 'dt': 1, 'A': 2, 'B': 3, 'C': 4, 'dt_bias': 'dt_bias', 'state': 'initial_states'}
class-attribute
instance-attribute
¶
RECURRENT_ARGUMENTS = {'state': 0, 'hidden_states': 1, 'dt': 2, 'A': 3, 'B': 4, 'C': 5, 'dt_bias': 'dt_bias'}
property
¶
Width of one head's values and outputs: the mixer's head_dim, which on a hybrid is not its softmax attention's.
class-attribute
instance-attribute
¶
states = DerivedEProperty(_states, description='The state after every token of this call; needs nnterp.chunk_per_token(model)', unavailable=needs_per_token_chunks)
The state after token t of this call, [batch, heads, state_dim, head_dim]: states[:, t].
A select function: where the scan call that fires takes name (the two kernels' call sites differ).
Run every Mamba-2 mixer of model with a chunk size of 1, so the chunk scan keeps the state after every token.
The chunk scan materializes the state at each chunk boundary
(new_states, StateSpace.states); with chunk_size == 1 every
token is a boundary. The outputs are the same; the cost is not: the
recurrence between chunks is quadratic in the number of chunks, so a long
prompt runs markedly slower. Unlike route_kernels, which binds a
family's kernels process-wide, this sets an attribute of each mixer module
of this one model (chunk_size, which the forward reads on every
call), so it applies from the next trace and to no other model. Call it
on a loaded model: a lazily built model's modules are replaced when its
weights arrive. enabled=False restores the chunk size each mixer was
built with (config.chunk_size, mamba_chunk_size on Bamba, Falcon-H1
and GraniteMoeHybrid).
Why StateSpace.states is unavailable: the kernels have no source to read, or the chunk scan's chunks are longer than a token.