components¶
The base envoys and the descriptor every family builds its toolkit from.
A family gives its modules standard values by wrapping them in Envoy
subclasses (nnsight's envoys=): Layer on the decoder block, Attention
on the softmax-attention module, Mlp on the feed-forward, and
Moe on a mixture of experts (an Mlp with routing values),
LinearAttention on a hybrid's gated DeltaNet mixer, SelectiveScan on a
Mamba-1 mixer and StateSpace on a Mamba-2 (SSD) mixer, all RecurrentMixers
(the mechanism every recurrent mixer shares). A multimodal wrapper's vision
tower is a Vision, its blocks VisionLayer / VisionAttention /
VisionMlp, whose stream values are Patches; the tower also serves
where the image meets the text model (ImageTokenMask, ImageFeatures). The families under
nnterp.families subclass these, so a family whose forward is spelled
differently overrides only what differs, and key them on its own module types
in its ENVOYS.
A value a checkpoint does not have says so: every descriptor here takes
unavailable= (a reason, or a function of the envoy returning one), reading
it raises Unavailable with the reason before the model runs, and
Standard.support / StandardizedTransformer.support report the reasons
without raising. A family without a value assigns unavailable("...").
The three boundary values mean the same thing on every family:
layer_output is the residual stream leaving the block, and
attention_output / mlp_output are the contributions the sublayers
add to it, so layer.input + attention_output + mlp_output == layer_output
whether the block is sequential or parallel. On a family whose sublayer adds
the residual inside the module (BLOOM, MPT, DBRX) or adds a post-sublayer
norm's output instead (Gemma-2/3, OLMo-2), the family's subclass points the
value at the right place; the identity is what the tests check.
One descriptor, EProperty, whose key is a path from the host envoy:
"output": the host's own output, a view over the same location as.output, so reading, writing and in-place edits all go through the interleaver the way.outputdoes (Layer.layer_output)."../post_attention_layernorm.output","embed_tokens.output": a value produced by a module named relative to this one (a sandwich block's post-sublayer norm, the root's embedding)."/inputs": a path walked from the model's root, by standard names, whatever the host (the vision tower'simage_token_mask)."source.attention_interface_1.inputs": an operation inside a forward, reached through.source, optionally one element of it (select).- Computed from several served values, a
DerivedEProperty(a DeltaNet layer's state after every token). An operation inside a called function only exists once someone has drilled into that call in the current run — the interleaver resolves the callee from the live value and clears what it built at the start of every run — so this descriptor drills in before every read or write, the step a plain eproperty has no place for (Attention.attention_probabilities).
module-attribute
¶
NOT_ON_INTERFACE = "The attention does its own arithmetic rather than transformers' shared attention interface; not mapped for this family yet"
module-attribute
¶
PARALLEL = 'a parallel block: the attention and the MLP read the same input, so there is no stream between them'
module-attribute
¶
POST_NORM_ATTENTION = 'a post-norm block: the attention reads the residual stream itself, layer_input'
module-attribute
¶
module-attribute
¶
ChannelGates = Float[Tensor, f'{BATCH} {SEQ} {MIXER_HEADS} {KEY_DIM}']
module-attribute
¶
States = Float[Tensor, f'{BATCH} {SEQ} {MIXER_HEADS} {KEY_DIM} {VALUE_DIM}']
module-attribute
¶
SSDHeadOutputs = Float[Tensor, f'{BATCH} {SEQ} {MIXER_HEADS} {HEAD_DIM}']
module-attribute
¶
PER_IMAGE = "the Qwen ViT's attention makes one interface call per image (attention_interface_2; one per window on Qwen2.5-VL's windowed blocks), so no one tensor is the block's pattern; read attention_queries and attention_keys, whole, and split them where the calls do: at the attention's cu_seqlens argument (self_attn.inputs[1]['cu_seqlens']), which per image is the processor's image_grid_thw"
module-attribute
¶
__all__ = ['BATCH', 'CHANNELS', 'EXPERTS', 'GROUPS', 'HEAD_DIM', 'HEADS', 'HIDDEN', 'IMAGE_TOKENS', 'IMAGES', 'INTERMEDIATE', 'KEY', 'KEY_DIM', 'KV_HEADS', 'TOKENS', 'MIXER_HEADS', 'PATCHES', 'QK_HEAD_DIM', 'QUERY', 'SEQ', 'STATE_DIM', 'STREAMS', 'TOP_K', 'VALUE_DIM', 'VISION_HIDDEN', 'VOCAB', 'Attention', 'ChannelGates', 'DISPATCH', 'DerivedEProperty', 'EProperty', 'ExpertIndices', 'ExpertOutputs', 'ExpertWeights', 'Gates', 'HeadOutputs', 'INTERFACE', 'ImageFeatures', 'ImageTokenMask', 'Keys', 'Layer', 'LinearAttention', 'Patches', 'LOGITS', 'LinearQK', 'LinearV', 'Mlp', 'Moe', 'PER_SLOT', 'Pattern', 'Queries', 'RecurrentMixer', 'Residual', 'RouterLogits', 'ScanDecays', 'ScanQK', 'ScanState', 'ScanStates', 'ScanSteps', 'ScanValues', 'SelectiveScan', 'State', 'States', 'StreamMixing', 'StreamWeights', 'Streams', 'TokenEProperty', 'Values', 'Vision', 'VisionAttention', 'VisionLayer', 'VisionMlp', 'axis_size', 'blocks_support', 'ImageScatter', 'PER_IMAGE', 'PixtralVision', 'QwenVision', 'QwenVisionAttention', 'NOT_ON_INTERFACE', 'SSDHeadOutputs', 'SSDKeys', 'SSDQueries', 'SSDValues', 'Standard', 'StateSpace', 'Unavailable', 'chunk_per_token', 'first_tensor', 'image_token_id', 'interface_reason', 'mixture_reason', 'needs_eager', 'needs_grouped_experts', 'needs_kernel_source', 'needs_per_token_chunks', 'needs_recurrent_routing', 'needs_token_loop', 'needs_torch_kernels', 'no_shared_expert', 'per_call', 'pinned', 'rewrap', 'route_delta_rule', 'route_kernels', 'seq_first', 'unavailable', 'Neurons', 'FlatResidual', 'PARALLEL', 'POST_NORM_ATTENTION', 'POST_NORM_MLP']
Bases: Standard
A softmax-attention module: its contribution, its pattern, and its interior.
Everything but attention_output is read inside transformers' shared
eager_attention_forward, reached through the module's
attention_interface call (INTERFACE): the queries, keys and values
it receives, the scores entering its softmax, the pattern leaving it, and
the per-head outputs it returns. They are unavailable unless the model was
loaded with attn_implementation="eager"; off_interface is the one
place that decides, so a family with another reason (GPT-2's
reorder_and_upcast_attn) overrides that method. A family whose
attention does its own arithmetic redefines the pattern on its own op and
marks the rest unavailable(NOT_ON_INTERFACE).
Its sizes (num_heads, num_kv_heads, head_dim, qk_head_dim) are
this block's, read off the module, so they hold on a model whose blocks
differ (Gemma-4, MiMo-V2-Flash) and outside a trace.
The pattern is the dropout after the softmax, not the softmax itself: that is the tensor the values are mixed with on every family, after the cast back to the model dtype and, on a model with an attention sink (GPT-OSS), after the sink column is dropped.
property
¶
Width of one head's values and outputs: the module's v_head_dim (latent attention, MiMo-V2-Flash), else its head_dim / head_size.
property
¶
Width of one head's queries and keys: the module's qk_head_dim (latent attention), else its head_dim / head_size.
property
¶
Query heads: the module's num_heads / num_attention_heads / n_heads / n_head, else the output projection's input width over head_dim.
property
¶
Key/value heads as projected: the module's num_key_value_heads / num_kv_heads / kv_heads, else num_heads over its num_key_value_groups, else the key projection's width over qk_head_dim, else num_heads.
Why the shared attention interface does not run on this module, or None.
attention_queries(value: Tensor) -> Queries
The queries the attention interface receives, [batch, num_heads, seq, head_dim].
After the query projection and, on a family with rotary embeddings,
after they are applied. Assign to replace them. In-place edits reach
the model where torch allows them: GPT-2's queries, keys and values are
split views of one c_attn tensor and torch refuses to edit those in
place, so assign there.
attention_keys(value: Tensor) -> Keys
The keys the attention interface receives, [batch, num_kv_heads, seq, qk_head_dim].
Before repeat_kv, so under grouped-query attention the head axis
is num_kv_heads wide. Assign to replace them; in-place edits reach
the model except on GPT-2 (see attention_queries).
attention_values(value: Tensor) -> Values
The values the attention interface receives, [batch, num_kv_heads, seq, head_dim].
Before repeat_kv, like the keys. Assign to replace them; in-place
edits reach the model except on GPT-2 (see attention_queries).
attention_scores(value: Tensor) -> Pattern
The scaled, masked scores entering the softmax, [batch, num_heads, query, key].
softmax(attention_scores) is attention_probabilities up to the
dtype cast. Assign to replace them; in-place edits reach the model.
attention_probabilities(value: Tensor) -> Pattern
The attention pattern, [batch, num_heads, query, key].
The post-softmax probabilities as the values are mixed with them: in the model's dtype, and with an attention sink's column already dropped. Rows sum to one (to less than one on a sink model). Read it, edit it in place, or assign a tensor of the same shape::
with model.trace(prompt):
pattern = model.layers[3].self_attn.attention_probabilities.save()
model.layers[3].self_attn.attention_probabilities[:, 5] = 0
attention_head_outputs(value: Tensor) -> HeadOutputs
Each head's output before they are concatenated and projected, [batch, seq, num_heads, head_dim].
What the attention interface returns; the module reshapes it to
[batch, seq, hidden_size] and applies the output projection to get
attention_output. Assign to replace it, or edit it in place: the
interface returns a contiguous tensor on most families and a plain
transposed view on GPT-2, and torch accepts in-place edits on both.
DerivedEProperty(compute: Callable[[Envoy], Any], description: str | None = None, unavailable: Any = None)
Bases: EProperty
An EProperty computed from other served values rather than read at one location.
compute(envoy) runs at read time inside the trace and may read any
number of values; the result is read-only (assign through whatever
method the host offers). Listed in the repr like any value.
EProperty(key: str | Callable[[Envoy], str] | None = None, description: str | None = None, unavailable: str | Callable[[Envoy], str | None] | None = None, select: int | str | Callable[[Envoy], int | str | None] | None = None)
Bases: eproperty
An eproperty whose key is a path from the host, and which can say when it is unavailable.
| PARAMETER | DESCRIPTION |
|---|---|
key
|
Where the value lives, relative to the host envoy: dotted segments
ending in
TYPE:
|
description
|
Shown in the model's repr, like any eproperty's.
TYPE:
|
unavailable
|
A reason string, or a function of the envoy returning a
reason or
TYPE:
|
select
|
One element of the served value: with
TYPE:
|
The location is served by nnsight the way any eproperty's is, whatever the
path: a module's output, a sibling norm's, or an operation's arguments,
with preprocess, postprocess and transform all available. An
operation inside a called function only exists once someone has drilled
into that call in the current run (the interleaver resolves the callee
from the live value and clears what it built at the start of every run),
so a path through source is walked before every read or write. An
operation that is not there raises SourceNotAvailable naming what is,
rather than the AttributeError a descriptor would otherwise swallow into
"no attribute".
property
¶
The value's shape as a jaxtyping type (Residual, Pattern, ...), or None.
Read off the return annotation of the function that defines the
value. isinstance(tensor, value.layout) checks rank and dtype;
dims names the axes, the same on every family.
property
¶
The axis names of layout: ("batch", "seq", "hidden_size").
Why this value is not available on obj, or None when it is.
The repr line: (name) -> Layout [axes]: description for a value with a layout, else nnsight's line.
The key for obj: the path itself, or what the key function returns for it.
Whether the value is an operation inside a forward (a key function, or a source segment on its path).
staticmethod
¶
The served attribute a path ends in: output, input, inputs, or another an engine serves (logits).
Bases: RuntimeError
A standard value this checkpoint does not have, and why.
Raised on access, before anything runs, so a wrong assumption fails at the
line that makes it. Standard.support reports the same reasons without
raising.
Bases: Standard
A decoder block. Its hidden states are layer_output, whatever the block returns.
| ATTRIBUTE | DESCRIPTION |
|---|---|
self_attn |
The softmax attention, an
TYPE:
|
linear_attn |
The recurrent mixer, a
TYPE:
|
mlp |
The feed-forward, an
TYPE:
|
input_layernorm, |
The block's norms under their aliased names, where the family has them (their meaning varies; see the README).
TYPE:
|
Between layer_input and layer_output the block has four more boundaries,
each the input or output of one of its modules: attention_norm_output,
layer_mid, mlp_norm_output and mlp_input. They are defined on the
pre-norm block (x + attn(norm1(x)), then + mlp(norm2(.))) and keyed on
the aliased norms, so on a block built otherwise the family's subclass points
them at its own modules or marks them unavailable with the reason; a block
without the module reports no <module> module on this block.
Skip this block, handing hidden on as its residual stream.
The block does not run; hidden takes the place of its
layer_output, packed the way the block would have returned it (a
tuple family gets (hidden, None): the second element is the
attention weights nothing downstream reads). Call it inside a trace,
before the block runs.
layer_input(value: Tensor) -> Residual
The residual stream entering this block, before any of its norms.
The block's first argument where the block is called with the stream
(every transformers family), and the stream wherever an engine passes
it otherwise (vLLM's fused blocks take it in two halves). With the
contributions it is the identity the values are defined by:
layer_input + attention_output + mlp_output == layer_output.
In-place edits and assignment reach the model.
attention_norm_output(value: Tensor) -> Residual
The attention's pre-norm output: what the attention is called with, [batch, seq, hidden_size].
input_layernorm.output on a pre-norm block. In-place edits and
assignment reach the model.
layer_mid(value: Tensor) -> Residual
The residual stream between the two sublayers, [batch, seq, hidden_size].
What the MLP's pre-norm is called with: layer_input +
attention_output on a pre-norm block, where it is read as
post_attention_layernorm.input. A parallel block has no such
stream, and marks this unavailable. In-place edits and assignment reach
the model.
mlp_norm_output(value: Tensor) -> Residual
The MLP's pre-norm output, [batch, seq, hidden_size]: post_attention_layernorm.output on a pre-norm block.
The same tensor as mlp_input where nothing sits between the norm
and the MLP. In-place edits and assignment reach the model.
Bases: RecurrentMixer
A gated DeltaNet mixer (Qwen3-Next, Qwen3.5/3.6, OLMo-Hybrid; Kimi-Linear's KDA subclasses it): linear attention with a recurrent state.
It projects queries, keys and values like attention, but mixes them
through a per-head recurrent state instead of a softmax over keys: for
each token the state decays by a learned gate (decays), takes up the
new key/value pair scaled by betas, and the query reads against it.
So the shared names mean what they mean on Attention — the contribution,
the queries/keys/values the mixer receives, the per-head outputs — and
there is no pattern and no scores; instead there are the gate, the beta,
and the state entering and leaving the layer, one key_dim by
value_dim matrix per head.
Everything but attention_output is read at the delta-rule kernel
call: transformers' chunked rule on a prompt, its token-by-token rule on
a decode step, where the sequence axis is 1 and state_input is the
cached state. How that call is found, the kernels it needs, and the
state after every token (route_kernels(model.family, "torch"), or
route_delta_rule(model.family, "recurrent")) are RecurrentMixer's.
class-attribute
instance-attribute
¶
attention_queries(value: Tensor) -> LinearQK
The queries the delta rule receives, [batch, seq, heads, key_dim]: after the conv, the activation and the repeat to num_v_heads.
attention_keys(value: Tensor) -> LinearQK
The keys the delta rule receives, [batch, seq, heads, key_dim].
attention_values(value: Tensor) -> LinearV
The values the delta rule receives, [batch, seq, heads, value_dim].
decays(value: Tensor) -> Gates
The gate: the log of how much of the state each token keeps, [batch, seq, heads], float32 and non-positive.
betas(value: Tensor) -> Gates
How strongly each token's key/value pair is written into the state, [batch, seq, heads], in (0, 1) ((0, 2) on OLMo-Hybrid with linear_allow_neg_eigval).
state_input(value: Any) -> State | None
The state this call starts from: None on a fresh prompt, the cached state on a decode step.
A copy: the cache hands the kernel its own buffer and overwrites it in place with the step's new state afterwards, so the live tensor would read as this step's output by the time the trace ends. Assign to replace what the step starts from.
Bases: Standard
A feed-forward module. Its contribution is mlp_output; its width, intermediate_size, is read off the module.
property
¶
Width of this MLP's hidden layer; on a mixture of experts, one routed expert's.
Read off the module: its experts' intermediate_dim / expert_dim /
intermediate_size / ffn_hidden_size, else its own
intermediate_size / ffn_dim, else the down projection's input
width. A family whose module spells it another way overrides this.
mlp_activation(value: Tensor) -> Neurons
The activation's output, [batch, seq, intermediate_size]: act(gate_proj(x)) on a gated MLP,
act(up_proj(x)) on a plain one. The activation is act_fn, the family aliasing its own
spelling (act, activation_fn, activation, gelu_impl) in RENAME. In-place edits
and assignment reach the model.
mlp_neurons(value: Tensor) -> Neurons
The neurons' values the down projection sums, [batch, seq, intermediate_size].
On a gated MLP act(gate_proj(x)) * up_proj(x); on a plain one (GPT-2) the
same tensor as mlp_activation. The projection is down_proj, the family
aliasing its own spelling (c_proj, dense_4h_to_h, fc2, fc_out,
w2, output_linear) in RENAME. In-place edits and assignment reach the model.
Bases: Mlp
A mixture of experts: its routing, each routed expert's output, and the shared expert's.
mlp_output is what the block adds, as on any Mlp. The rest, read
where the model consumes them (router, experts and
shared_experts are the standard child names; a family aliases its
own):
router_logits[batch, seq, num_experts]: before the scoring, at the operation in the router's forward that produces them (LOGITS), so a write changes the routing.expert_weights/expert_indices[batch, seq, top_k]: the experts module's arguments. Zero a weight to ablate that slot; write an index to reroute (the weight is not recomputed).expert_outputs[batch, seq, top_k, hidden_size]: each slot's weighted output, inside transformers'grouped_mm/batched_mmexperts forward; unavailable underexperts_implementation="eager".routed_output: the experts' sum,expert_outputs.sum(2).shared_expert_output: the shared expert's contribution, where there is one;routed_output + shared_expert_outputis the mixture's output.
Its sizes, num_experts and top_k, are read off the router or the
experts module. SCORING says what the logits mean. A family whose module
runs the mixture only on some checkpoints (Gemma-4's dense MLP, which hosts
the block's experts) says why not in no_mixture, the one place every
value's availability starts from.
property
¶
Routed experts in this mixture: the router's or experts module's num_experts / num_local_experts / n_routed_experts.
property
¶
Experts each token is routed to: the module's or router's top_k / num_experts_per_tok.
Why this module runs no mixture of experts, or None (every Moe module does, unless its family says otherwise).
router_logits(value: Tensor) -> RouterLogits
The router's logits, [batch, seq, num_experts], where the router's forward produces them.
Before the scoring (SCORING) and any selection bias, so
softmax(router_logits) is the routing distribution on a softmax
router. Assign or edit in place: the router scores what is written,
so the weights and indices follow.
expert_weights(value: Tensor) -> ExpertWeights
The weights the experts receive, [batch, seq, top_k], one per routing slot.
After the router's normalization and scaling: what multiplies each
slot's expert output. Zero [b, t, j] to remove slot j of token
t; w.masked_fill(expert_indices == e, 0) removes expert e.
expert_indices(value: Tensor) -> ExpertIndices
The experts each token is routed to, [batch, seq, top_k], int64.
Write an index to reroute that slot to another expert; its weight stays what the router gave the expert it chose.
expert_outputs(value: Tensor) -> ExpertOutputs
Each slot's expert output times its weight, [batch, seq, top_k, hidden_size], before the slots are summed.
expert_outputs.sum(2) == routed_output. The model's own tensor in
transformers' grouped_mm (the default) and batched_mm experts
forwards; unavailable under eager and hub kernels.
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.
Bases: RecurrentMixer
A Mamba-1 mixer (Mamba, Falcon-Mamba, Jamba's Mamba blocks): a selective state-space scan.
Each channel of the mixer's inner width carries a state of state_dim
numbers. For each token t the mixer computes a step size per channel,
dt = softplus(dt_proj(x) + dt_bias), and two input-dependent vectors
B and C of state_dim numbers shared by every channel, and
runs, per channel::
h_t = exp(dt_t * A) * h_{t-1} + dt_t * B_t * x_t
y_t = C_t . h_t + D * x_t
with A a learned [channels, state_dim] matrix of negative rates.
The output is y gated by silu(z) and projected back by
out_proj. In the shared vocabulary: C reads the state as a
query does, B writes the input into it as a key does, x is what
gets written (the values), dt is how strongly each token writes
(betas), and dt * A is the log of how much of the state each
token keeps (decays, per channel and state dimension, where a
DeltaNet's is one number per head).
The values are read at the kernel call the forward makes:
mamba_selective_scan(x, dt, A, B, C, D=, z=, delta_bias=, ...) on a
prompt, mamba_selective_state_update(state, x, dt, A, B, C, D, z=,
dt_bias=, ...) on a decode step (one token). The kernel's tensors are
channel-first (x and dt [batch, channels, seq], B and
C [batch, state_dim, seq]; a decode step drops the sequence
axis); each value is a view of the argument laid out tokens first, and
an assignment is laid back out, so reads, writes and in-place edits
reach the kernel.
attention_queries/attention_keys:C/B,[batch, seq, 1, state_dim](one group).attention_values:x,[batch, seq, channels].betas:softplus(dt + dt_bias), the step size, anddecays:A * betas,[batch, seq, channels, state_dim]; derived, read-only.state_input: a decode step's cached state,[batch, channels, state_dim](a copy; assign to replace it),Noneon a prompt, whose scan starts from zeros.attention_head_outputs:ybefore the gate andout_proj,[batch, seq, channels]: a binding inside the kernel.state_output: the state after the call's last token, as the cache receives it: the scan's returned state on a prompt, the updated state the decode kernel copies into the cache.state/states/state_after/set_state_after: the state after each token, from the scan's token loop (STATE_OP) on a prompt and the single update (STEP_STATE_OP) on a decode step.
The pure-torch selective scan is the token loop, so
route_kernels(model.family, "torch") binds each kernel name to its
own pure-torch function. With mamba_ssm installed the forward
dispatches both to its compiled kernels, which have no source to read
inside (and need CUDA): route before the first trace.
class-attribute
instance-attribute
¶
class-attribute
instance-attribute
¶
ARGUMENTS = {'x': (0, 1), 'dt': (1, 2), 'A': (2, 3), 'B': (3, 4), 'C': (4, 5), 'dt_bias': ('delta_bias', 'dt_bias')}
class-attribute
instance-attribute
¶
betas = DerivedEProperty(_betas, description='The step size dt = softplus(dt + dt_bias), per token and channel; derived, read-only', unavailable=needs_torch_kernels)
class-attribute
instance-attribute
¶
decays = DerivedEProperty(_decays, description='The per-token log decay dt * A of the state; derived, read-only', unavailable=needs_torch_kernels)
class-attribute
instance-attribute
¶
states = DerivedEProperty(_scan_states, description="The state after every token of this call; needs route_kernels(family, 'torch')", unavailable=needs_token_loop)
state_output(value: Tensor) -> ScanState
The state after this call's last token, [batch, channels, state_dim]: what the next decode step starts from.
The scan's returned state on a prompt; on a decode step the updated state as the kernel copies it into the cache. Either way it is what the cache receives: assigning it changes what the next step starts from, not this call's output.
Bases: Envoy
An envoy carrying standard values: what Layer, Attention and Mlp share.
A value that is an operation inside a forward is served only on a call
whose forward was instrumented before the call began. Instrumenting on
first read is enough when the read comes before the module runs, which
is the usual case; it is too late when the worker is already parked
inside the module, as it is when a block's mlp_output (an operation
in the block's forward) is read after the block's attention_output.
A family whose value sits in such a forward sets sourced = True on
that envoy, and the forward is instrumented when the envoy is built and
again when real weights replace meta ones, which reinstalls the plain
forward.
TokenEProperty(key: str | Callable[[Envoy], str] | None = None, description: str | None = None, unavailable: str | Callable[[Envoy], str | None] | None = None, select: int | str | Callable[[Envoy], int | str | None] | None = None)
Bases: EProperty
An EProperty whose tensor the model may hold flat over tokens, [batch * seq, ...], one rank below its layout.
The same descriptor in every other respect (path, select,
unavailable, layout, support, the repr line). The preprocess,
postprocess and transform see the model's own tensor; the user sees and
writes this invoke's [batch, seq, ...]:
__get__reads asEPropertydoes, then servesrowsof the result: a view, so an in-place edit lands on the model's tensor.__set__splices the assigned[batch, seq, ...]into the whole tensor the location holds now (splice), then writes asEPropertydoes.
Under two or more invokes the spliced tensor's leading dim is not the
batch; nnsight keeps such an edit from PR #738 on (Batcher._widen_tensor).
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].
Bases: Standard
The wrapper's inner model (model.model): the module whose forward scatters the image features into the token embeddings.
A family keys it on the wrapper model's type (LlavaModel: ImageScatter);
Vision.image_features reads the scatter_argument of the scatter
operation, the features in inputs_embeds.masked_scatter(image_mask,
image_features), which is what the text model receives at the image
tokens whatever the wrapper did after its projector (LLaVA-NeXT's
unpadding and newline tokens). A wrapper that writes the features in
elsewhere subclasses it with its own operation and argument. The forward
is instrumented at build (sourced), so the scatter is served after the
tower's values, which run inside the same forward. Carries no values of
its own.
A wrapper that scatters in the root's own forward (Llama 4's
Llama4ForConditionalGeneration) has no child to key: its family sets
ROOT_SCATTER to the operation's name, the root is the host
(scatter_host), and StandardizedTransformer instruments its own
forward at build.
Bases: Vision
Pixtral's tower (Mistral 3, Llava-Pixtral): packed, every image's patches in one row.
The convolution runs on the batch padded to its largest image, each image
is cropped to its own grid, and the grids are flattened and concatenated
into [1, all patches, vision_hidden_size] before ln_pre; the blocks
attend under a block-diagonal mask, so no patch sees another image's. So
Patches has a leading 1 here, and patch_embeddings is that packed
row as it enters ln_pre (the convolution's own output, patch_embed.output,
is still the padded grid). The processor's image_sizes split the row
per image: (height // patch_size) * (width // patch_size) patches each,
in order. The attention interior is whole: one interface call over the
packed row, so attention_probabilities is [1, heads, all patches,
all patches], zero between images. image_size is Unavailable: the
config's is the largest side the processor resizes to, not one size every
image has.
patch_embeddings(value: Tensor) -> Patches
Every image's patch embeddings, [1, all patches, vision_hidden_size], image after image, each in raster order.
What enters ln_pre: the convolution's output cropped per image and
concatenated. Assign a tensor of the same shape to replace it.
Bases: Vision
The Qwen ViT (Qwen2-VL, Qwen2.5-VL, Qwen3-VL, Qwen3.5): packed, its merger inside, at model.visual.
The tower runs on [patches, vision_hidden_size], every image's patches
concatenated, so its stream values are Patches with 1 in the images
axis, [1, patches, vision_hidden_size], and its blocks' attention is a
QwenVisionAttention.
The processor cuts each image into patch_size squares (two frames
deep), and lays them out in spatial_merge_size x spatial_merge_size
blocks, each block's patches consecutive, images one after another; the
merger folds each block into one image token. So patches runs over
every image of the invoke, image_grid_thw ([t, h, w] per image, in
patches) splits it, and an image's patches are in merge-block order, not
raster order. Qwen2.5-VL's tower further permutes them into attention
windows at entry: its blocks' values, tower_output and the merger's
output are in window order, and the tower restores the merge-block order
after the merger.
The tower has no final norm, so tower_output is the last block's
stream, the merger's input. It is served at the tower's output, after the
merger has run, so projector.input and projector.output are read
before tower_output.
Sizes, off the vision config: vision_hidden_size is the tower's width
(embed_dim on Qwen2-VL, whose config's hidden_size is the merger's
output width), num_heads, intermediate_size (Qwen2-VL's is
embed_dim * mlp_ratio), patch_size, spatial_merge_size and
window_size (Qwen2.5-VL's, in pixels; None on the others).
image_size is Unavailable: the tower takes any resolution.
Bases: VisionAttention
The Qwen ViT's attention: one interface call per image (or window), so the interior is read whole around them.
attention_queries, attention_keys and attention_values are the
whole [1, heads, patches, head_dim] tensors (the queries and keys after
the 2D rotary embedding) before the module splits them per image, and
attention_head_outputs is the per-image outputs concatenated back,
[1, patches, heads, head_dim], under any implementation but flash.
attention_scores and attention_probabilities are Unavailable
(PER_IMAGE): loading eager does not help, since the module splits
under every implementation.
attention_queries(value: Tensor) -> Queries
Every image's queries, [1, heads, patches, head_dim], before the per-image split; assign or edit in place.
attention_keys(value: Tensor) -> Keys
Every image's keys, [1, heads, patches, head_dim], before the per-image split; assign or edit in place.
attention_values(value: Tensor) -> Values
Every image's values, [1, heads, patches, head_dim], before the per-image split; assign or edit in place.
attention_head_outputs(value: Tensor) -> HeadOutputs
Every image's per-head outputs concatenated back, [1, patches, heads, head_dim], before the output projection.
Bases: Standard
model.vision: a vision tower's root. Its blocks are layers; its sizes are its own config's.
Values: patch_embeddings, the patch embedding's output one row per patch
(before any position embedding, CLS token or pre-norm the tower adds), and
tower_output, the last block's stream after the tower's final norm
where it has one, before any pooling, CLS dropping or adapter (the tower's
last_hidden_state, unless the tower returns something after those).
The host reads the tower its own way: Gemma 3 pools tower_output,
Llava 1.5 takes vision.layers[-2].layer_output and drops the CLS
token, so what reaches the projector is projector.input, not
necessarily tower_output. Where the image meets the text model:
image_token_mask, read off the model's inputs, and image_features,
read at the scatter (ImageScatter, keyed from the root), so
layers[0].input[image_token_mask] == image_features. Both need an image
to reach the model (no_images).
The sizes (num_layers, vision_hidden_size, num_heads, head_dim,
intermediate_size, patch_size, image_size) are read off the tower's
own config, the plain spellings SigLIP's and CLIP's use; a tower that
spells one its own way overrides the property on a subclass. The tower's
width is a different quantity from the text model's hidden_size and
is named apart, so axis_size on image_features (text width, served
here) walks past the tower to the root; num_heads, head_dim,
num_layers and intermediate_size keep the plain names and shadow
the root's on purpose, because the tower's values want the tower's numbers.
| ATTRIBUTE | DESCRIPTION |
|---|---|
layers |
The tower's blocks, each a
TYPE:
|
patch_embed |
The patch embedding (a convolution on SigLIP and CLIP).
TYPE:
|
norm |
The final norm over the patches, where the tower has one (SigLIP's
TYPE:
|
property
¶
Width of one head: the config's head_dim when it says, else vision_hidden_size // num_heads.
property
¶
Side of the square image the tower is configured for, in pixels; Unavailable on a tower with no fixed resolution (variable_resolution).
Why no image reaches the model this tower belongs to, or None: the one place that decides.
A wrapper whose family names no projector (where the image enters
the text model is unknown), or a load without a processor (a
task="text-generation" load of the wrapper).
Which tower values this checkpoint has, the way StandardizedTransformer.support reports the text model's.
The tower's own values (image_token_mask and image_features
among them) plus every block value over layers ("layer_output",
"self_attn.attention_probabilities", ...): None when available
on every block, else {layer: reason}. With layer, that block's
values alone. Empty on a load no image reaches (no_images): the tower
never runs, so the model's support has no vision rows, and a
read of an image value raises Unavailable saying why.
Why a value read inside the eager attention forward is unavailable, or None.
[batch, num_heads, seq, head_dim] <-> [batch, seq, num_heads, head_dim], as a view.
The standard layout of attention_head_outputs is what the shared
interface returns, sequence before heads. A family whose own arithmetic
keeps heads first serves a transposed view on read, so in-place edits
still land, and transposes back on write; the transpose is its own inverse.
unavailable(reason: str) -> EProperty
A value a family does not have: assign it in the class body in place of the inherited one.
attention_probabilities = unavailable("no softmax: the attention is linear")
keeps the name in the tree and in the repr, with the reason, and makes any
access raise Unavailable with it.
Why this module runs no mixture of experts, or None: the one place a family decides it (Moe.no_mixture).
Why the per-slot expert outputs are unavailable, or None: they exist only in the grouped_mm / batched_mm forwards.
Why shared_expert_output is unavailable on this mixture, or None: it has no shared expert.
Why the per-token state is unavailable: only the token-by-token kernel materializes it.
Why a recurrent mixer's kernel values are unavailable: the kernel has no Python source to read inside.
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.
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").
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.
Why a value read inside the kernels' own bodies is unavailable: they are not bound to the pure-torch functions.
Why the state after each token is unavailable: the kernels are not routed, or the scan runs mambapy's parallel scan.
The size of layout axis axis on a value envoy serves.
An axis that is a fact about the model is named as the size that holds
it (HIDDEN is "hidden_size", HEADS "num_heads"), so its size
is that attribute: on envoy when its class declares it (an
attention's num_heads, a mixture's num_experts), else on the
first envoy above it that does, up to the model's root (hidden_size,
vocab_size)::
axis_size(model.layers[3], HIDDEN) # model.hidden_size
axis_size(model.layers[3].self_attn, HEADS) # this block's heads
Only an envoy's class counts (a property, the root's
StandardizedProperty), never an attribute of the wrapped module, which
can share a name and hold another number.
| RAISES | DESCRIPTION |
|---|---|
ValueError
|
no envoy from |
The block values of a block list, the way support reports them.
With layer, that block's values by dotted name ("layer_output",
"self_attn.attention_probabilities"). Without, each value -> None
when available on every block, else {layer: reason} for the blocks
where it is not. Shared by the text model's blocks and a vision tower's.
The tensor of a module output: the first element when it is a tuple.
value in the shape the module returned: back in its tuple, if any.
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.
The id of the token the processor puts where an image's features go: image_token_id (image_token_index on older configs).