attention¶
Attention: a softmax-attention module, its contribution, its pattern and its interior.
module-attribute
¶
NOT_ON_INTERFACE = "The attention does its own arithmetic rather than transformers' shared attention interface; not mapped for this family yet"
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.
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.