API Reference¶
Auto-generated reference for the nnterp package, organized to mirror the source tree.
Most users need StandardizedTransformer and the values on the components; the rest documents
the internals.
nnterp: one module vocabulary across transformer architectures.
StandardizedTransformer is an nnsight TransformersModel whose envoy tree
answers to the same names whatever the checkpoint's family. The vocabulary is
Llama's::
model.embed_tokens
model.layers[i].input_layernorm
model.layers[i].self_attn
model.layers[i].post_attention_layernorm
model.layers[i].mlp
model.norm
model.lm_head
and the root answers for the whole model: logits, token_embeddings,
next_token_probs and the sizes (num_layers, num_heads, ...).
Each family under nnterp.families says how its own names map onto those, and
nnterp.components holds the envoys that give standard modules standard values
(layer_output, attention_output, mlp_output,
attention_probabilities), which each family subclasses. On a multimodal
checkpoint loaded with task="image-text-to-text" the vision tower is
model.vision (its blocks model.vision.layers[i], its values
model.vision.image_token_mask and model.vision.image_features) and the
projector model.projector.
StandardizedVLLM is the same over nnsight's VLLM engine, with vLLM's own
implementation of each family under nnterp.families.vllm.
module-attribute
¶
__all__ = ['Attention', 'DerivedEProperty', 'EProperty', 'Layer', 'LinearAttention', 'Mlp', 'Moe', 'RecurrentMixer', 'SelectiveScan', 'Standard', 'Standardized', 'StandardizedTransformer', 'StandardizedVLLM', 'StateSpace', 'Unavailable', 'UnsupportedFamily', 'Vision', 'axis_size', 'chunk_per_token', 'route_delta_rule', 'route_kernels', 'unavailable']
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: 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.
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: 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
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.
Bases: ValueError
No toolkit standardizes the checkpoint: it has no family, and the best-effort default cannot standardize it.
What a standardized model is on any engine: a family, its sizes, its availability, and the methods over its values.
Mixed in ahead of the nnsight model class that runs the checkpoint
(StandardizedTransformer over TransformersModel,
nnterp.StandardizedVLLM over nnsight's VLLM). The leaf resolves
family and hands its RENAME and ENVOYS to nnsight, and defines
the root's values (logits, token_embeddings, ...), which are read
where its engine keeps them. Everything here is written against the
standard names and values only, so it holds on both.
Skip blocks start through end inclusive.
The residual stream entering block start (or skip_with) is
handed straight to block end + 1; the skipped blocks do not run.
Negative indices count from the end. Inside a trace::
with model.trace(prompt):
model.skip_layers(4, 7)
logits = model.logits.save()
steer(layers: int | list[int], vector: Tensor, factor: float = 1.0, token_positions: int | list[int] | slice | None = None, batch_index: int | None = None) -> None
Add factor * vector to the residual stream leaving the given blocks.
vector is [hidden_size] (or broadcastable to the selected slice).
token_positions restricts the positions and batch_index the
row; both default to all. The add is in place on layer_output, so
it reaches the model. Inside a trace, with layers ascending.
The k most likely tokens of one [vocab_size] distribution, as {token: probability}, most likely first.
Keyed by the decoded text. Where two or more of the k decode to
the same text (partial UTF-8 byte tokens all decode to '�'),
those entries are keyed by the tokenizer's raw vocabulary token
instead (convert_ids_to_tokens, unique per id), with the id
appended should even that repeat, so the dict always holds k
entries, each with its own probability.
The k most likely next tokens for each position of hidden, [..., hidden_size].
project_on_vocab then softmax, one {token: probability} per
position, in row-major order over the leading dimensions.
Which standard values this checkpoint has, without running anything.
With layer, that block's values by dotted name ("layer_output",
"self_attn.attention_probabilities", "mlp.mlp_output"): None
when available, else the reason, including "no <module> module"
when the block has no such module at all (a hybrid's blocks have either
self_attn or linear_attn). Without, the root's values
plus every block value: None when available on every block, else
{layer: reason} for the blocks where it is not, so a hybrid reads
as a short dict.
The tree decides what is listed: every child of a block that carries
standard values (a Standard envoy) is walked under its standard
name, so a value installed through envoys= or a registered family
appears here as it does in the envoy's own Standard.support, and a
module no block has (OPT's mlp) has no entry. A standard child of
the root (the vision tower, Vision) is listed the same way, its own
support rows under its standard name ("vision.image_features",
"vision.self_attn.attention_probabilities"); a tower no image
reaches lists none.
Key/value heads: num_key_value_heads under grouped-query attention, else num_heads.
Width of one attention head: the config's head_dim when it says (Qwen3, Gemma), else hidden_size // num_heads.
Width of one head's queries and keys: head_dim, unless the family separates them (DeepSeek's latent attention).
Width of the dense MLP's hidden layer, config.intermediate_size; a mixture of experts' experts are moe_intermediate_size wide.
StandardizedTransformer(repo_id: Any, *args: Any, family: ModuleType | None = None, rename: dict[str, str | list[str]] | None = None, envoys: dict | None = None, tokenizer_kwargs: dict | None = None, **kwargs: Any)
Bases: Standardized, TransformersModel
A causal language model whose modules answer to one set of names.
The checkpoint's config is read first (its model_type), the matching
family toolkit is looked up in nnterp.families.REGISTRY, and the family's
RENAME is handed to nnsight's rename so every envoy in the tree also
answers to the standard name. The original names keep working: an alias is
an extra attribute on the same envoy, not a replacement.
The family's ENVOYS wraps its decoder blocks in its Layer
(layer_output), its attention modules in its Attention
(attention_output, attention_probabilities) and its feed-forwards
in its Mlp (mlp_output). The pattern is read inside the eager
attention forward, so it is unavailable unless the model is loaded with
attn_implementation="eager"; support says which values this
checkpoint has. See nnterp.components.
| PARAMETER | DESCRIPTION |
|---|---|
repo_id
|
A HuggingFace repo id, or an already-loaded
TYPE:
|
family
|
The family to use instead of the one looked up: a module or
any object with
TYPE:
|
rename
|
Extra aliases, merged over the family's; a key given here wins.
TYPE:
|
envoys
|
Extra
TYPE:
|
tokenizer_kwargs
|
Attributes to set on the loaded tokenizer, such as
TYPE:
|
**kwargs
|
Passed through to
TYPE:
|
Over the values, skip_layers, steer, project_on_vocab and
get_topk_closest_tokens do the common things (see each). Inside a
trace input_ids, input_size and attention_mask are what the model
was called with.
The root envoy also answers for the whole model, inside a trace: logits
(the logits the model returns, after any softcap or scale past the head), token_embeddings
(the embedding's output) and next_token_probs; and outside one: the
sizes num_layers, num_heads, num_kv_heads, head_dim,
qk_head_dim, hidden_size, intermediate_size and vocab_size,
read off the config, each a StandardizedProperty the family can define
instead; project_on_vocab is a StandardizedCapability, a method the family
can define the same way.
On an image-text-to-text checkpoint loaded with task="image-text-to-text"
the text names keep their meaning (the language model's), the vision tower
is vision (a Vision, with its own blocks, sizes and values, among them
vision.image_token_mask and vision.image_features) and the last
module before the scatter into the text stream is projector.
| ATTRIBUTE | DESCRIPTION |
|---|---|
family |
The toolkit module the checkpoint resolved to, or the
|
layers |
The decoder blocks, each a
TYPE:
|
embed_tokens, |
The embedding, the final norm, the unembedding.
TYPE:
|
vision |
The vision tower, a
TYPE:
|
projector |
The last module before the wrapper scatters the image features into the text stream.
TYPE:
|
| RAISES | DESCRIPTION |
|---|---|
UnsupportedFamily
|
when no family is passed and none covers the checkpoint's |
property
¶
The checkpoint's tokenizer loaded with add_prefix_space=False, so "word" and " word" differ.
What nnterp.prompt_utils.get_first_tokens uses. Loaded once, on first use.
token_embeddings(value: Tensor) -> Residual
The embedding module's output, [batch, seq, hidden_size].
Positional embeddings and embedding norms (GPT-2's wpe, BLOOM's
word_embeddings_layernorm) are applied after this by the families
that have them. Assign to replace it.
final_norm_output(value: Tensor) -> Residual
The final norm's output, [batch, seq, hidden_size]: the stream the head projects. Assign to replace it; in-place edits reach the model.
lm_head_output(value: Tensor) -> Logits
lm_head's own output, [batch, seq, vocab_size].
logits is what the model returns; on most families the same tensor,
on Gemma-2 the softcapped one. Assign to replace it; in-place edits
reach the model.
Logits for a residual-stream tensor: the final norm, lm_head, and what the model does after the head.
The logit lens: applied to a block's layer_output it reads that
layer's prediction; applied to the last block's, it is logits.
Works inside a trace on a live value and outside on a saved one.
After the head the plain case is the text config's
final_logit_softcapping, when set (Gemma-2's own config; a
multimodal checkpoint's text_config). A family whose model does
something else there (Cohere multiplies by logit_scale, Granite
divides by logits_scaling) defines
def project_on_vocab(model, hidden) in its module, which is bound
in this one's place (StandardizedCapability).
The available values in forward order: name -> rank; equal ranks are one location, either read first.
With layer, that block's values by dotted name (support's), from
0. Without, the root's, numbered across the whole forward. Measured by
a probe run on first use (nnterp.order.probe) and kept on the model,
so route kernels (route_kernels, chunk_per_token) before the first.
Where a value falls in the forward: (side, r), which sorts reads into forward order.
side is -1 or num_layers for a root value served before or
after the blocks, else layer; r is order's rank.
| RAISES | DESCRIPTION |
|---|---|
Unavailable
|
the checkpoint has no such value ( |
KeyError
|
there is no value of that name (a block value needs |
StandardizedVLLM(repo_id: str, *args: Any, family: ModuleType | None = None, rename: dict[str, str | list[str]] | None = None, envoys: dict | None = None, **kwargs: Any)
Bases: Standardized, VLLM
A causal language model on the vLLM engine whose modules answer to the standard names.
StandardizedTransformer's counterpart for nnsight's VLLM: the same
vocabulary (model.layers[i].self_attn, .mlp, model.norm,
model.lm_head), the same values with the same layouts, and the same
methods, over the engine's continuous batching. The family is vLLM's own
implementation of the checkpoint's model_type, a module under
nnterp.families.vllm.
It is nnsight's VLLM otherwise, with vLLM's defaults: a trace is a
generation request (max_tokens=16 and temperature=1.0 unless you
say), sampling settings go on trace/invoke, one prompt per invoke,
and a block runs in the engine's worker, where nnterp has to be importable.
What that changes for the values:
- A request is one sequence, so every value's batch axis is 1:
layer_outputis[1, tokens, hidden_size], the whole prompt on the prefill and one token on each decode step (tracer.iter). logitsis this step's,[1, 1, vocab_size]: the engine computes the last position only.logits[:, -1]andnext_token_probsread the same on both engines.- Values are private copies, so a saved one stays what it was read as.
- The queries, keys, values and head outputs are what goes into and
comes out of the engine's attention layer, for this step's tokens. The
scores and the pattern are inside its kernel, so they are recomputed
from the queries and keys: read-only, and on the prefill only (a decode
step raises
Unavailable). There is noattention_mask. project_on_vocabruns the engine's modules, so it works inside a trace only.
| PARAMETER | DESCRIPTION |
|---|---|
repo_id
|
A HuggingFace repo id.
TYPE:
|
family
|
The family to use instead of the one looked up: a module or
any object with
TYPE:
|
rename
|
Extra aliases, merged over the family's; a key given here wins.
TYPE:
|
envoys
|
Extra
TYPE:
|
**kwargs
|
Passed through to
TYPE:
|
| ATTRIBUTE | DESCRIPTION |
|---|---|
family |
The vLLM family the checkpoint resolved to, or the
|
| RAISES | DESCRIPTION |
|---|---|
UnsupportedFamily
|
when no family is passed and vLLM's implementation
of the checkpoint's |
class-attribute
instance-attribute
¶
lm_head_output = unavailable("vLLM never calls lm_head: its logits processor projects with lm_head's weight and applies any softcap and scale in the same call, on the last position only; read logits")
class-attribute
instance-attribute
¶
attention_mask = unavailable('a vLLM request is one unpadded sequence; there is no mask')
token_embeddings(value: Tensor) -> Residual
The embedding module's output, [1, tokens, hidden_size]; what a family scales or adds positions to comes after.
final_norm_output(value: Tensor) -> Residual
The final norm's output, [1, tokens, hidden_size]: every position of the step, of which the engine projects the last.
A family whose blocks hand the stream on in two halves calls the final
norm with both, and the norm returns (normed, residual): this is
the normed half. Assign to replace it; in-place edits reach the model.
Logits for a residual-stream tensor: the final norm, lm_head, and vLLM's logits processor.
The logit lens, inside a trace: the modules are the engine's, so there
is nothing to call outside one. The processor applies the model's
softcapping and scaling where it has them. hidden is [...,
hidden]; the leading dimensions are kept. A family whose model
projects another way defines def project_on_vocab(model, hidden)
in its module.
Not supported on the vllm engine: the forward order is measured on StandardizedTransformer only.
Not supported on the vllm engine: the forward order is measured on StandardizedTransformer only.
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 |
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).
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.
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.