View as Markdown

Recurrent mixers¶

A recurrent mixer's values are the arguments and results of a kernel call its forward makes, and RecurrentMixer in nnterp/components/recurrent.py holds everything about reaching that call that does not depend on which mixer it is: which of the two kernels fires, decided once per call by the forward's own test; the requirement that the kernels be transformers' pure-torch ones, so .source has a body to read inside; route_kernels, which rebinds a family's kernels process-wide; and the state after every token, read by occurrence arithmetic. LinearAttention (gated DeltaNet), SelectiveScan (Mamba-1) and StateSpace (Mamba-2) name their kernels in class constants and declare their values.

import torch
import nnterp
from nnterp import StandardizedTransformer, route_kernels

model = StandardizedTransformer("Qwen/Qwen3.5-0.8B", dispatch=True, device="cpu")
mix = model.layers[0].linear_attn
cls = type(mix)
print(cls.__name__, cls.CHUNK_KERNEL, cls.RECURRENT_KERNEL, cls.STATE_OP)   # LinearAttention torch_chunk_gated_delta_rule_0 torch_recurrent_gated_delta_rule_0 last_recurrent_state_3
print(cls.attention_queries.key, cls.state.key, cls.states.key)            # <kernel.inputs> <_token_state_op> <_states>
print(model.support()["linear_attn.states"][0])
# the state after each token is materialized only by the token-by-token kernel; the chunked kernel a prompt runs through carries it between chunks. Call nnterp.route_kernels(model.family, 'torch') before tracing this layer (slower, like attn_implementation='eager')

route_kernels(model.family, "torch")   # before the first trace that touches a mixer of this family
prompt = "Hello world there"
n = len(model.tokenizer(prompt).input_ids)
with model.trace(prompt):
    q = mix.attention_queries.save()
    prompt_states = mix.states.save()
    out = mix.state_output.save()
print(n, tuple(q.shape), tuple(prompt_states.shape), torch.equal(prompt_states[:, -1], out))  # 3 (1, 3, 16, 128) (1, 3, 16, 128, 128) True

firsts, seqs, stacks, outs = [], [], [], []          # bound outside: a list first bound in the block does not survive it
with model.generate(prompt, max_new_tokens=3, do_sample=False) as tracer:
    for step in tracer.iter[:]:
        op, first, seq = mix._token_op()              # (the state op, the occurrence its first token is, this call's length)
        seqs.append(seq)
        firsts.append(first)
        stacks.append(mix.states.save())
        outs.append(mix.state_output.save())
print(seqs, firsts)                                   # [3, 1, 1] [0, 0, 1]
print(torch.equal(stacks[0], prompt_states), torch.equal(stacks[2][:, 0], outs[2]), torch.equal(stacks[2][:, 0], outs[1]))  # True True False
route_kernels(model.family, "default")

Qwen3.5 is a hybrid, and block 0 holds a linear_attn. support() lists the value as a dict because every fourth block has a self_attn instead. The open tracer.iter[:] ends with nnsight's "never reached" warning, which is how an open slice ends. Without the first offset, step 2 would ask for occurrence 0 of the recurrent kernel's state op, which step 1 consumed.

Which kernel fires¶

The mixer's forward binds a flag before it branches and then calls the token-by-token kernel when it has a cached state and one token, the chunked kernel otherwise. A subclass names the pieces in class constants.

constant	meaning	`LinearAttention`
`BRANCH`	the binding made before the branch, `True` when the call continues a cached state	`"use_precomputed_states_0"`, the base's default
`SEQ_OP`	the forward's masking of its input, once per call, with the call's length on axis 1	`"apply_mask_to_padding_states_0"`, the base's default
`CHUNK_KERNEL`	the call a prompt runs through	`"torch_chunk_gated_delta_rule_0"`
`RECURRENT_KERNEL`	the call a decode step runs through	`"torch_recurrent_gated_delta_rule_0"`
`STATE_OP`	inside the token-by-token kernel, the binding of the state after each token's update, or `None`	`"last_recurrent_state_3"`
`STEP_STATE_OP`	set when the decode kernel updates the state once instead of looping: the binding of the new state inside it	`None`

KERNEL is one staticmethod on the base. It reads the BRANCH and SEQ_OP outputs in the order the forward binds them, sorted by each operation's source line, and returns RECURRENT_KERNEL when the call continues a cached state with one token, CHUNK_KERNEL otherwise, which means that several tokens over a cached state read the prompt's kernel. kernel(attribute) is the key function the values are declared on and returns source.<KERNEL(envoy)>.<attribute>: attention_queries is kernel("inputs") with select=0, state_output is kernel("output") with select=1. attention_output is the module's own output, unwrapped from a tuple and rewrapped on write, as on Attention.

Once per call¶

The two bindings are served once per module call, and every value of the call depends on the kernel they pick, so KERNEL makes its choice once, inside per_call(envoy, "kernel", choose). per_call keeps one record per (envoy.path, key) on the worker greenlet running the intervention code and recomputes when the record's call is not the current one. The current call is mediator.iteration or mediator.occurrence(f"{envoy.path}.output"). 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 the read is in call k. Relaxed, on step 0, or outside tracer.iter, the call is how many times the module's output has been passed, which nnsight counts whether or not anything read it (interleaver). Every run of every invoke has its own worker, so two invokes reading one mixer never share a record and nothing stays on the envoy. The records are the kernel choice, _token_op's (op, first, seq) under "call", and StateSpace's kernel arguments under "arguments".

Pure-torch kernels and routing¶

transformers decorates each kernel with its hub dispatcher, which binds the module-level name to a closure holding torch_function, the pure-torch body, and implementation, the optimized kernel when its package is installed. _dispatch reads that closure with inspect.getclosurevars, and needs_torch_kernels, the unavailable predicate of every kernel value, refuses when implementation is not torch_function, naming the package: A compiled kernel has no Python source to drill into.

route_kernels(family, "torch") rebinds the family's two kernel names in its modeling module to their pure-torch functions, stashing the originals in the module's _nnterp_kernels, from which "default" restores them. On a mixer with a STATE_OP the prompt's name is bound to the loop kernel too, which on a DeltaNet is the decode kernel: after routing both names run the token-by-token rule, the operation keeps its label torch_chunk_gated_delta_rule_0, and STATE_OP under it fires once per token. The routing is process-wide and must come before the first trace that touches a mixer of that family, because nnsight builds an instrumented forward over a copy of the module's globals on the first drill. A mixer traced before routing keeps calling the kernel it was compiled with, and a read of states on it fails as nnsight's "no attribute states" while an untouched mixer of the same model serves it. route_delta_rule(family, "recurrent" | "chunked") is the DeltaNet spelling.

mamba = StandardizedTransformer("state-spaces/mamba-130m-hf", dispatch=True, device="cpu")
scan = mamba.layers[0].linear_attn
print(type(scan).__name__, type(scan).STATE_OP, type(scan).STEP_STATE_OP)   # SelectiveScan ssm_state_3 ssm_state_0
print(mamba.support()["linear_attn.attention_queries"][0])
# read inside transformers' pure-torch mamba_selective_scan, but this process dispatches it to an optimized kernel (mamba_ssm) with no Python source; uninstall it, or call nnterp.route_kernels(model.family, 'torch'), to read these
route_kernels(mamba.family, "torch")
with mamba.trace(prompt):
    states = scan.states.save()
    out = mamba.layers[0].layer_output.save()
print(mamba.support()["linear_attn.states"], tuple(states.shape), tuple(out.shape))   # None (1, 3, 1536, 16) (1, 3, 768)
route_kernels(mamba.family, "default")

With mamba_ssm installed, a Mamba-1 or Mamba-2 model needs the routing before any trace on a CPU, including one that reads only layer_output: The compiled kernel refuses a CPU tensor (Expected u.is_cuda() to be true).

The state after every token¶

state is an EProperty with the key function _token_state_op, which returns source.<kernel>.source.<STATE_OP>.output, a location with one occurrence per token, so nnsight's own tracer.iter walks it: for t in tracer.iter[:n]: mix.state is the state after every prompt token, and an assignment there is a write the following tokens continue from (iterator). Under tracer.iter the pin at the read is a token index while the bindings KERNEL reads and the kernel call's .fn fire once per forward. _token_state_op decides the kernel and drills into its call inside with pinned(None), which relaxes the pin and restores it on exit, and the pinned read that follows asks for that token's occurrence. pinned(n) sets the worker mediator's iteration the way tracer.iter[n] does, and one read relaxes it: A with holds one read.

states, state_after(t) and set_state_after(t, value) address the same location by index and need _token_op's three numbers, computed once per call: the state op under the drilled kernel, its sequence length from _seq() (the queries' axis 1 on a DeltaNet), and first, the op's occurrence count at the moment the kernel's arguments are read. Reading an argument parks the worker at the call's start, the one moment that count is the number of tokens earlier calls put through the op. Occurrences are counted per location since the op was drilled this run, so on step k of a generate the recurrent op's count is k minus 1, and the prompt's tokens, under the chunk kernel's label, do not offset it. _states then loops with pinned(first + t): op.output and stacks on axis 1. states is a DerivedEProperty, read-only.

needs_recurrent_routing guards state and states, and answers in this order: a mixer with no STATE_OP says its kernels do not materialize the state per token, an optimized kernel gives needs_torch_kernels' reason, and a CHUNK_KERNEL name still bound to the chunked function says to route. The predicate reads the live binding, so support() follows the routing, as the first block shows.

Mamba-1 and Mamba-2¶

SelectiveScan, in nnterp/components/selective_scan.py, differs from the DeltaNet at the points the base leaves open to a subclass. Its forward binds the branch flag before the mask, which KERNEL's line ordering handles. Its pure-torch prompt kernel is itself the token loop (STATE_OP = "ssm_state_3") while the decode kernel updates once (STEP_STATE_OP = "ssm_state_0"): route_kernels binds each name to its own function, and _state_op(kernel) picks the binding for the kernel that fires. Its arguments sit at different positions in the two kernels, and select is a function looking the name up in ARGUMENTS by whether the call decodes. Its tensors are channel-first, and each preprocess returns a tokens-first view that the postprocess lays back. betas and decays are DerivedEProperty values computed from dt, dt_bias and A the way the kernel computes them.

StateSpace, in nnterp/components/state_space.py, has STATE_OP = None, since neither Mamba-2 kernel binds the state once per token. state and set_state_after are unavailable(...) with reasons, and states is a DerivedEProperty over the chunk scan's new_states_0 binding, the state at every chunk boundary. chunk_per_token(model) sets chunk_size to 1 on each StateSpace module of one model, which makes every token a boundary. It is per model, not process-wide, and needs_per_token_chunks reports the chunk size until then. The two kernels take their arguments in different places (CHUNK_ARGUMENTS, RECURRENT_ARGUMENTS, through argument(name)). The update has no sequence axis, and a decode step's values gain one on read and lose it on write. betas and decays are the kernel's dt seen through dt_bias, the softplus and A, with an inverse postprocess, and an assignment reaches the kernel. Values that need another argument read the call's (args, kwargs) once per call through _arguments.

mamba2 = StandardizedTransformer("AntonV/mamba2-130m-hf", dispatch=True, device="cpu")
ssd = mamba2.layers[0].linear_attn
route_kernels(mamba2.family, "torch")
print(mamba2.support()["linear_attn.states"][0])
# 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)
nnterp.chunk_per_token(mamba2)
with mamba2.trace(prompt):
    states = ssd.states.save()
    out = ssd.state_output.save()
print(mamba2.support()["linear_attn.states"], tuple(states.shape), torch.equal(states[:, -1], out))   # None (1, 3, 24, 128, 64) True
nnterp.chunk_per_token(mamba2, enabled=False)
route_kernels(mamba2.family, "default")

Adding a mixer¶

A new mixer subclasses RecurrentMixer, sets CHUNK_KERNEL and RECURRENT_KERNEL (and BRANCH, SEQ_OP or KERNEL itself, as a staticmethod, when its forward spells the test otherwise), sets STATE_OP only when a kernel binds the state once per token, and declares its values at kernel("inputs") and kernel("output") with unavailable=needs_torch_kernels. A kernel with tensors not laid out [batch, seq, ...] overrides _seq. tests/test_base.py holds the base's own tests, and a family's file mixes in SelectiveScanSuite or the tests in tests/families/ssd.py (Testing).