recurrent¶
RecurrentMixer: how a recurrent mixer's values are reached, whatever the values are.
A recurrent mixer (a gated DeltaNet, a state-space layer) runs its sequence
through a kernel function the modeling module calls: one kernel for a prompt,
another for a one-token step over a cached state (a decode step of
generate), picked by what the forward binds before it branches. Its values are the arguments and results of that
call. What differs between mixers is which arguments mean what; what they
share is how the call is found, which kernel names must be transformers'
pure-torch ones for there to be a call to read inside, how the family's
kernels are rebound (route_kernels), and, when the token-by-token kernel
materializes it, the state after every token. That shared part is
RecurrentMixer; a subclass (LinearAttention, SelectiveScan, StateSpace)
names its kernels in class constants and declares its values.
module-attribute
¶
States = Float[Tensor, f'{BATCH} {SEQ} {MIXER_HEADS} {KEY_DIM} {VALUE_DIM}']
Bases: Standard
A sequence mixer with a recurrent state, read at the kernel call its forward makes.
The mechanism every recurrent mixer shares, apart from what its values
are. A subclass names its kernels in class constants and declares
its values at kernel("inputs") / kernel("output"):
BRANCH: the binding the forward makes before it picks a kernel,Truewhen the call continues from a cached state.CHUNK_KERNEL: the call a prompt runs through.RECURRENT_KERNEL: the call each decode step ofgenerateruns through.STATE_OP: inside the token-by-token kernel, the binding of the state after each token's update;Nonewhen the mixer's kernels do not materialize it.STEP_STATE_OP: set when the decode kernel is a single-step update rather than the token loop (Mamba-1): the binding of the new state inside it. The prompt's kernel is then the token loop.
KERNEL names the call that fires on this call of the mixer, decided
once per call by the forward's own test: RECURRENT_KERNEL for one token
over a cached state (BRANCH true and a sequence axis of 1 at SEQ_OP),
CHUNK_KERNEL otherwise, a prompt or several tokens over a cached state.
So the same values work in a trace and at every step of
tracer.iter. The kernels have to be transformers' pure-torch ones
for there to be a call to read inside (needs_torch_kernels).
The state after every token of a prompt exists only in the
token-by-token kernel: a chunked kernel carries the state between chunks.
Like eager attention, that is the user's choice:
nnterp.route_kernels(model.family, "torch") routes the family's
prompts through the token-by-token kernel, and then state, states,
state_after and set_state_after read and write the state at any
position; without it, or on a mixer with no STATE_OP, reading one
raises Unavailable with the reason and support reports it. A subclass
whose kernels keep the state another way overrides them (StateSpace
reads states off the chunk scan's boundaries).
class-attribute
instance-attribute
¶
states = DerivedEProperty(_states, description="The recurrent state after every token of this call; needs route_kernels(family, 'torch')", unavailable=needs_recurrent_routing)
staticmethod
¶
The kernel call that fires on this call of the mixer, decided once per call (per_call).
The forward's own test, use_precomputed_states and seq_len == 1:
the token-by-token update for one token over a cached state, the
prompt's kernel for everything else. The two bindings are read in the
order the forward makes them, which differs between mixers.
state(value: Tensor) -> State
The state after a token of the prompt: one occurrence per token.
A location inside the token-by-token kernel, so it takes nnsight's
own iteration: for t in tracer.iter[:]: mix.state reads the state
after every token, tracer.iter[4] the one after token 4, and an
assignment there is a write the following tokens continue from.
Outside any tracer.iter a read is the state after token 0. It
lives on the prompt's kernel call: under generate walk it in an
inner tracer.iter on step 0, and take a decode step's state from
the call's output, one token per step. Needs route_kernels(family, "torch").
The state after token t of this call: for _ in tracer.iter[t]: mix.state on a prompt, as a call that also counts from the call's own first token on a decode step.
Assign state at token t: the tokens after it continue from value.
A write inside the prompt. Reads follow the forward: in the same
trace, read positions before t before the write and positions
after it afterwards; states (every position) only before it.
Bind a family's recurrent kernels, process-wide.
"torch" binds the prompt's and the decode step's kernel names in the
family's modeling module to transformers' pure-torch kernels, the
functions the dispatcher falls back to: the slower path, and the only one
with Python source to read values inside. On a mixer whose token-by-token
kernel materializes the state after every token (STATE_OP set), the
prompt's name is bound to that kernel, so a prompt runs through it and
RecurrentMixer.state, states, state_after and set_state_after
read and write the state at any position. On a gated DeltaNet that is
the decode kernel, so both names are bound to it; on Mamba-1 the
pure-torch selective scan is itself the token loop, and each name keeps
its own (RecurrentMixer.STEP_STATE_OP). "default" restores what
the module bound at import. Like installing a kernel, it applies to every
model of the family; call it before tracing a layer, since a forward
.source has already instrumented keeps the binding it was compiled
with. family is the family module (model.family, or
nnterp.families.qwen3_5_text) or its modeling module.
route_kernels for a gated DeltaNet, in its own words.
"recurrent" is route_kernels(family, "torch"): the family's
prompts run through transformers' token-by-token gated delta rule, which
materializes the state after every token. "chunked" is
route_kernels(family, "default").
Why a recurrent mixer's kernel values are unavailable: the kernel has no Python source to read inside.
Why the per-token state is unavailable: only the token-by-token kernel materializes it.
compute() once per call of the envoy's module, under key.
For something several values in one call depend on and the model serves once: the kernel a forward branches to, a call's arguments, its sequence length. The record is kept until the worker is in another call of the module.
Which call that is comes from two facts nnsight keeps. A read pinned by
tracer.iter to step k is served at the k-th occurrence of its
location, and a mixer's kernel fires once per call, so a pinned read is
in call k. After a step's first read the pin relaxes, and outside
tracer.iter (or on step 0) there is none; then the call is how many
times the module's .output has been passed, which nnsight counts for
every location whether or not it was read: inside call c that many have
returned, and between calls it names the one about to start.
The records live on the worker, the greenlet running this intervention
code. Every run of every invoke has its own (a replayed model.edit
too), so two invokes reading one mixer never share a record, a record
never outlives its run, and the tensors in it are freed with the worker.
Pin this worker's reads to occurrence n of their location, as tracer.iter[n] does; None relaxes the pin.
A pinned read is served at the n-th occurrence of its location and
relaxes the pin, so one read per with. The pin the worker had is
restored on the way out.
A key at whichever kernel fires on this call: source.<KERNEL>.<attribute>.