Coverage for transformer_lens/model_bridge/generalized_components/attention.py: 88%
506 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-10-23 20:45 +0000
1"""Attention bridge component.
3This module contains the bridge component for attention layers.
4"""
5import inspect
6import logging
7import warnings
8from typing import Any, Dict, Optional, Tuple
10import einops
11import torch
12from transformers.pytorch_utils import Conv1D
14logger = logging.getLogger(__name__)
16from transformer_lens.conversion_utils.conversion_steps.attention_auto_conversion import (
17 AttentionAutoConversion,
18)
19from transformer_lens.conversion_utils.conversion_steps.base_tensor_conversion import (
20 BaseTensorConversion,
21)
22from transformer_lens.hook_points import HookPoint
23from transformer_lens.model_bridge.generalized_components.base import (
24 GeneralizedComponent,
25)
26from transformer_lens.utilities.hf_utils import get_rotary_pct_from_config
27from transformer_lens.utilities.quantization import require_readable_weight
30class PerLayerGeometryError(ValueError):
31 """A weight accessor refused a head split on per-layer attention geometry.
33 Dedicated type so callers with a raw-weight fallback (the centering
34 benchmark) catch exactly this and nothing else — a generic ValueError
35 catch would silently mask unrelated accessor regressions.
36 """
39class AttentionBridge(GeneralizedComponent):
40 """Bridge component for attention layers.
42 This component handles the conversion between Hugging Face attention layers
43 and TransformerLens attention components.
44 """
46 hook_aliases = {
47 "hook_q": "q.hook_out",
48 "hook_k": "k.hook_out",
49 "hook_v": "v.hook_out",
50 "hook_z": "o.hook_in",
51 }
53 # Override to False on variants without a pre-LN fork (e.g. MLA); skips
54 # the split-qkv HookPoints and the BlockBridge pre-ln1 capture.
55 supports_split_qkv_fork: bool = True
56 # Reconstructed variants can opt in independently of Q/K/V input forks.
57 supports_attn_result: bool = False
58 property_aliases = {
59 "W_Q": "q.weight",
60 "W_K": "k.weight",
61 "W_V": "v.weight",
62 "W_O": "o.weight",
63 "b_Q": "q.bias",
64 "b_K": "k.bias",
65 "b_V": "v.bias",
66 "b_O": "o.bias",
67 }
69 def __init__(
70 self,
71 name: Optional[str],
72 config: Any,
73 submodules: Optional[Dict[str, GeneralizedComponent]] = None,
74 conversion_rule: Optional[BaseTensorConversion] = None,
75 pattern_conversion_rule: Optional[BaseTensorConversion] = None,
76 maintain_native_attention: bool = False,
77 requires_position_embeddings: bool = False,
78 requires_attention_mask: bool = False,
79 attention_mask_4d: bool = False,
80 requires_relative_position_bias: bool = False,
81 is_cross_attention: bool = False,
82 is_causal: bool = True,
83 optional: bool = False,
84 fused_qkv: bool = False,
85 ):
86 """Initialize the attention bridge.
88 Args:
89 name: The name of this component, or None when projections live on the parent
90 config: Model configuration (required for auto-conversion detection)
91 submodules: Dictionary of submodules to register (e.g., q_proj, k_proj, etc.)
92 conversion_rule: Optional conversion rule. If None, AttentionAutoConversion will be used
93 pattern_conversion_rule: Optional conversion rule for attention patterns. If None,
94 uses AttentionPatternConversion to ensure [n_heads, pos, pos] shape
95 maintain_native_attention: If True, preserve the original HF attention implementation
96 without wrapping. Use for models with custom attention
97 (e.g., attention sinks, specialized RoPE). Defaults to False.
98 requires_position_embeddings: If True, this attention requires position_embeddings argument
99 (e.g., Gemma-3 with dual RoPE). Defaults to False.
100 requires_attention_mask: If True, this attention requires attention_mask argument
101 (e.g., GPTNeoX/Pythia). Defaults to False.
102 attention_mask_4d: If True, generate 4D attention_mask [batch, 1, tgt_len, src_len]
103 instead of 2D [batch, seq_len]. Required for OPT. Defaults to False.
104 requires_relative_position_bias: T5/mT5-style relative attention; supplies a
105 zero ``position_bias`` so HF's forward skips its ``cache_position[-1]`` fallback.
106 is_cross_attention: Encoder-decoder cross-attention; supplies ``key_value_states``.
107 is_causal: If True, apply a causal (lower-triangular) mask when reconstructing
108 attention. Set False for bidirectional encoders (e.g. T5Gemma's encoder).
109 """
110 if conversion_rule is None: 110 ↛ 112line 110 didn't jump to line 112 because the condition on line 110 was always true
111 conversion_rule = AttentionAutoConversion(config)
112 super().__init__(
113 name,
114 config=config,
115 submodules=submodules or {},
116 conversion_rule=conversion_rule,
117 optional=optional,
118 )
119 if fused_qkv:
120 # A combined QKV projection has no q/k/v submodules for the class
121 # aliases to resolve against; leaving them in place produces dead
122 # aliases and resolution-audit noise (raven Wqkv, OpenELM qkv_proj).
123 self.hook_aliases = {
124 key: value
125 for key, value in type(self).hook_aliases.items()
126 if key not in {"hook_q", "hook_k", "hook_v"}
127 }
128 self.hook_attn_scores = HookPoint()
129 self.hook_pattern = HookPoint()
130 self.hook_hidden_states = HookPoint()
131 # Per-head attention output, pre-sum across heads.
132 # Shape [batch, pos, n_heads, d_model] when fired. Gated at fire time
133 # by cfg.use_attn_result; the HookPoint exists unconditionally so
134 # run_with_cache key lookups never miss.
135 self.hook_result = HookPoint()
136 # Pre-ln1 fork hooks ([B, S, H, D]) gated by use_split_qkv_input /
137 # use_attn_in; fall back to post-ln1 if BlockBridge can't wire ln1. See #1317.
138 if self.supports_split_qkv_fork:
139 self.hook_attn_in = HookPoint()
140 self.hook_q_input = HookPoint()
141 self.hook_k_input = HookPoint()
142 self.hook_v_input = HookPoint()
143 self._captured_pre_ln_residual: Optional[torch.Tensor] = None
144 self._ln1_module: Optional[torch.nn.Module] = None
145 if (
146 hasattr(config, "positional_embedding_type")
147 and config.positional_embedding_type == "rotary"
148 ):
149 self.hook_rot_k = HookPoint()
150 self.hook_rot_q = HookPoint()
151 self.hook_hidden_states.hook_conversion = conversion_rule
152 if pattern_conversion_rule is not None: 152 ↛ 153line 152 didn't jump to line 153 because the condition on line 152 was never true
153 self.hook_pattern.hook_conversion = pattern_conversion_rule
154 self._attn_scores = None
155 self._pattern = None
156 self._hf_forward_wrapped = False
157 self.maintain_native_attention = maintain_native_attention
158 self.requires_position_embeddings = requires_position_embeddings
159 self.requires_attention_mask = requires_attention_mask
160 self.attention_mask_4d = attention_mask_4d
161 self.requires_relative_position_bias = requires_relative_position_bias
162 self.is_cross_attention = is_cross_attention
163 self.is_causal = is_causal
164 self._layer_idx: Optional[int] = None
166 def set_original_component(self, original_component: torch.nn.Module) -> None:
167 """Set original component and capture layer index for KV caching."""
168 super().set_original_component(original_component)
169 layer_idx_raw = getattr(original_component, "layer_idx", None)
170 if layer_idx_raw is not None:
171 self._layer_idx = int(layer_idx_raw)
173 def _apply_ln1_per_head(self, x: torch.Tensor) -> torch.Tensor:
174 """Apply ln1 to [B, S, H, D] with H folded into the batch. Identity if ln1 unwired.
176 Routes through the raw HF norm to avoid refiring ln1's internal hooks
177 per-head — deliberate divergence from legacy's *Pre sub-hook firing.
178 """
179 if self._ln1_module is None: 179 ↛ 180line 179 didn't jump to line 180 because the condition on line 179 was never true
180 return x
181 b, s, h, d = x.shape
182 return self._ln1_module(x.reshape(b * s * h, d)).reshape(b, s, h, d)
184 def _fork_and_norm_per_head(
185 self, source: torch.Tensor, hook: HookPoint, n_heads: int
186 ) -> torch.Tensor:
187 """Repeat residual to [B, S, H, D], fire ``hook``, re-LN iff source is pre-LN."""
188 forked = einops.repeat(source, "b s d -> b s h d", h=n_heads).contiguous()
189 forked = hook(forked)
190 if self._captured_pre_ln_residual is not None:
191 forked = self._apply_ln1_per_head(forked)
192 return forked
194 def setup_hook_compatibility(self) -> None:
195 """Setup hook compatibility transformations to match HookedTransformer behavior.
197 This sets up hook conversions that ensure Bridge hooks have the same shapes
198 as HookedTransformer hooks. This includes reshaping Q/K/V/Z hooks from
199 [batch, seq, d_model] to [batch, seq, n_heads, d_head] format.
201 This is called during Bridge.__init__ and should always be run.
202 Note: This method is idempotent - can be called multiple times safely.
203 """
204 if self._hf_forward_wrapped:
205 return
206 if hasattr(self.config, "n_heads"): 206 ↛ 208line 206 didn't jump to line 208 because the condition on line 206 was always true
207 self._setup_qkv_hook_reshaping()
208 self._hf_forward_wrapped = True
210 def get_random_inputs(
211 self,
212 batch_size: int = 2,
213 seq_len: int = 8,
214 device: Optional[torch.device] = None,
215 dtype: Optional[torch.dtype] = None,
216 ) -> Dict[str, Any]:
217 """Get random inputs for testing this attention component.
219 Generates appropriate inputs based on the attention's requirements
220 (position_embeddings, attention_mask, etc.).
222 Args:
223 batch_size: Batch size for the test inputs
224 seq_len: Sequence length for the test inputs
225 device: Device to create tensors on (defaults to CPU)
226 dtype: Dtype for generated tensors (defaults to float32)
228 Returns:
229 Dictionary of keyword arguments to pass to forward()
230 """
231 if device is None:
232 device = torch.device("cpu")
233 if dtype is None:
234 dtype = torch.float32
235 d_model = self.config.d_model if self.config and hasattr(self.config, "d_model") else 768
236 inputs: Dict[str, Any] = {
237 "hidden_states": torch.randn(batch_size, seq_len, d_model, device=device, dtype=dtype)
238 }
239 if self.requires_position_embeddings:
240 if self.config: 240 ↛ 248line 240 didn't jump to line 248 because the condition on line 240 was always true
241 if hasattr(self.config, "d_head"): 241 ↛ 243line 241 didn't jump to line 243 because the condition on line 241 was always true
242 d_head = self.config.d_head
243 elif hasattr(self.config, "head_dim"):
244 d_head = self.config.head_dim
245 else:
246 d_head = 64
247 else:
248 d_head = 64
249 rotary_pct = get_rotary_pct_from_config(self.config)
250 rotary_ndims = int(rotary_pct * d_head)
251 cos = torch.ones(1, seq_len, rotary_ndims, device=device, dtype=dtype)
252 sin = torch.zeros(1, seq_len, rotary_ndims, device=device, dtype=dtype)
253 inputs["position_embeddings"] = (cos, sin)
254 # For models with internal rotary embeddings (e.g., GPT-J), the HF attention
255 # forward expects position_ids to index into embed_positions. Models using
256 # requires_position_embeddings get (cos, sin) tuples instead.
257 if ( 257 ↛ 263line 257 didn't jump to line 263 because the condition on line 257 was never true
258 self.config
259 and hasattr(self.config, "positional_embedding_type")
260 and self.config.positional_embedding_type == "rotary"
261 and not self.requires_position_embeddings
262 ):
263 inputs["position_ids"] = (
264 torch.arange(seq_len, device=device).unsqueeze(0).expand(batch_size, -1)
265 )
266 if self.requires_attention_mask: 266 ↛ 267line 266 didn't jump to line 267 because the condition on line 266 was never true
267 if self.attention_mask_4d:
268 # Generate 4D attention mask [batch, 1, tgt_len, src_len] for models like OPT
269 inputs["attention_mask"] = torch.ones(
270 batch_size, 1, seq_len, seq_len, device=device
271 )
272 else:
273 # Generate 2D attention mask [batch, seq_len] for most models
274 inputs["attention_mask"] = torch.ones(batch_size, seq_len, device=device)
275 if self.requires_relative_position_bias: 275 ↛ 277line 275 didn't jump to line 277 because the condition on line 275 was never true
276 # Zero bias short-circuits HF's None-cache_position fallback in T5Attention.
277 n_heads = self.config.n_heads if self.config and hasattr(self.config, "n_heads") else 1
278 inputs["position_bias"] = torch.zeros(
279 1, n_heads, seq_len, seq_len, device=device, dtype=dtype
280 )
281 if self.is_cross_attention:
282 inputs["key_value_states"] = torch.randn(
283 batch_size, seq_len, d_model, device=device, dtype=dtype
284 )
285 return inputs
287 def _setup_qkv_hook_reshaping(self) -> None:
288 """Setup hook reshaping for Q/K/V/Z to match HookedTransformer shapes.
290 Reshapes hooks from [batch, seq, d_model] to [batch, seq, n_heads, d_head] format.
291 For models with Grouped Query Attention (GQA), K and V use n_kv_heads instead of n_heads.
293 Sets up conversions for:
294 - q.hook_out (aliased as hook_q)
295 - k.hook_out (aliased as hook_k) - uses n_kv_heads if GQA
296 - v.hook_out (aliased as hook_v) - uses n_kv_heads if GQA
297 - o.hook_in (aliased as hook_z)
298 """
300 class ReshapeForAttentionHeads(BaseTensorConversion):
301 """Reshape tensors to split attention heads for Q/K/V/Z compatibility."""
303 # Marks the [batch, pos, head, d_head] layout so consumers (e.g. run_with_cache's
304 # pos_slice) know the position axis is two from the end, not one.
305 splits_attention_heads = True
307 def __init__(self, n_heads: int, d_head: int):
308 super().__init__()
309 self.n_heads = n_heads
310 self.d_head = d_head
312 def handle_conversion(self, input_value, *full_context):
313 """Convert from [batch, seq, d_model] to [batch, seq, n_heads, d_head]."""
314 if len(input_value.shape) == 3: 314 ↛ 318line 314 didn't jump to line 318 because the condition on line 314 was always true
315 b, s, d = input_value.shape
316 if d == self.n_heads * self.d_head:
317 return input_value.view(b, s, self.n_heads, self.d_head)
318 return input_value
320 def revert(self, input_value, *full_context):
321 """Revert from [batch, seq, n_heads, d_head] to [batch, seq, d_model]."""
322 if len(input_value.shape) == 4: 322 ↛ 327line 322 didn't jump to line 327 because the condition on line 322 was always true
323 b, s, n_h, d_h = input_value.shape
324 if n_h == self.n_heads and d_h == self.d_head: 324 ↛ 327line 324 didn't jump to line 327 because the condition on line 324 was always true
325 # reshape (not view) — callers may pass non-contiguous tensors
326 return input_value.reshape(b, s, n_h * d_h)
327 return input_value
329 if self.config is None: 329 ↛ 330line 329 didn't jump to line 330 because the condition on line 329 was never true
330 raise RuntimeError(f"Config not set for {self.name}")
332 # Get n_heads (try n_heads first, then n_head)
333 if hasattr(self.config, "n_heads"): 333 ↛ 335line 333 didn't jump to line 335 because the condition on line 333 was always true
334 n_heads = self.config.n_heads
335 elif hasattr(self.config, "n_head"):
336 n_heads = self.config.n_head
337 else:
338 # Can't setup reshaping without knowing number of heads
339 return
341 # Get d_head (try d_head first, then compute from d_model or n_embd)
342 if hasattr(self.config, "d_head"):
343 d_head = self.config.d_head
344 elif hasattr(self.config, "d_model"): 344 ↛ 345line 344 didn't jump to line 345 because the condition on line 344 was never true
345 d_head = self.config.d_model // n_heads
346 elif hasattr(self.config, "n_embd"): 346 ↛ 347line 346 didn't jump to line 347 because the condition on line 346 was never true
347 d_head = self.config.n_embd // n_heads
348 else:
349 # Can't setup reshaping without knowing head dimension
350 return
351 n_kv_heads = n_heads
352 if hasattr(self.config, "n_key_value_heads") and self.config.n_key_value_heads is not None:
353 n_kv_heads = self.config.n_key_value_heads
355 # Per-layer d_head makes width division ambiguous (gemma4: 2048 = 8x256
356 # = 4x512), so dividing by the majority scalar mis-factorizes silently.
357 layer_d_head: Optional[int] = d_head
358 per_layer_hd = getattr(self.config, "per_layer_head_dim", None)
359 if per_layer_hd:
360 entry: Any = None
361 if self._layer_idx is not None and self._layer_idx < len(per_layer_hd):
362 entry = per_layer_hd[self._layer_idx]
363 if not (isinstance(entry, int) and entry > 0):
364 entry = getattr(self.original_component, "head_dim", None)
365 layer_d_head = entry if isinstance(entry, int) and entry > 0 else None
367 def conversion_dims(
368 proj: Any, cfg_heads: int, width_attr: str = "out_features"
369 ) -> Tuple[int, int]:
370 """(n_heads, d_head) from the BOUND module width, not the cfg scalars.
372 OpenELM/laguna vary head counts per layer; cfg-only dims no-op there,
373 flipping hook shapes mid-model. MLA excluded — its o-width is v_head_dim.
374 """
375 if getattr(self, "_v_head_dim", None): 375 ↛ 376line 375 didn't jump to line 376 because the condition on line 375 was never true
376 return cfg_heads, d_head
377 component = getattr(proj, "original_component", None)
378 width = getattr(component, width_attr, None)
379 if width and layer_d_head and width % layer_d_head == 0:
380 return width // layer_d_head, layer_d_head
381 return cfg_heads, layer_d_head or d_head
383 if hasattr(self, "q") and self.q is not None and hasattr(self.q, "hook_out"):
384 q_reshape = ReshapeForAttentionHeads(*conversion_dims(self.q, n_heads))
385 self.q.hook_out.hook_conversion = q_reshape
386 if hasattr(self, "k") and self.k is not None and hasattr(self.k, "hook_out"):
387 k_reshape = ReshapeForAttentionHeads(*conversion_dims(self.k, n_kv_heads))
388 self.k.hook_out.hook_conversion = k_reshape
389 # Subclasses that de-alias hook_k onto their own HookPoint (K-scaling
390 # architectures) need the same conversion, whichever order runs first.
391 if "hook_k" not in self.hook_aliases and isinstance(
392 getattr(self, "hook_k", None), HookPoint
393 ):
394 self.hook_k.hook_conversion = k_reshape
395 if hasattr(self, "v") and self.v is not None and hasattr(self.v, "hook_out"):
396 v_reshape = ReshapeForAttentionHeads(*conversion_dims(self.v, n_kv_heads))
397 self.v.hook_out.hook_conversion = v_reshape
398 if hasattr(self, "o") and self.o is not None and hasattr(self.o, "hook_in"): 398 ↛ 404line 398 didn't jump to line 404 because the condition on line 398 was always true
399 z_reshape = ReshapeForAttentionHeads(
400 *conversion_dims(self.o, n_heads, width_attr="in_features")
401 )
402 self.o.hook_in.hook_conversion = z_reshape
404 class TransposeRotaryHeads(BaseTensorConversion):
405 """Transpose rotary hook tensors from HF format to HookedTransformer format."""
407 def handle_conversion(self, input_value, *full_context):
408 """Convert from [batch, n_heads, seq, d_head] to [batch, seq, n_heads, d_head]."""
409 if len(input_value.shape) == 4: 409 ↛ 411line 409 didn't jump to line 411 because the condition on line 409 was always true
410 return input_value.transpose(1, 2)
411 return input_value
413 def revert(self, input_value, *full_context):
414 """Revert from [batch, seq, n_heads, d_head] to [batch, n_heads, seq, d_head]."""
415 if len(input_value.shape) == 4: 415 ↛ 417line 415 didn't jump to line 417 because the condition on line 415 was always true
416 return input_value.transpose(1, 2)
417 return input_value
419 if hasattr(self, "hook_rot_q"):
420 self.hook_rot_q.hook_conversion = TransposeRotaryHeads()
421 if hasattr(self, "hook_rot_k"):
422 self.hook_rot_k.hook_conversion = TransposeRotaryHeads()
424 def _update_kv_cache(
425 self, k: torch.Tensor, v: torch.Tensor, **kwargs: Any
426 ) -> tuple[torch.Tensor, torch.Tensor]:
427 """Update KV cache if provided, returning the (possibly extended) K and V.
429 Call this after K/V projections and any positional embeddings (e.g. RoPE)
430 have been applied, but before computing attention scores. If no cache is
431 present in kwargs, K and V are returned unchanged.
432 """
433 past_key_values = kwargs.get("past_key_values", None)
434 if past_key_values is None:
435 # GPT-NeoX/GPT-J/Bloom/Falcon/MPT/CodeGen/GPTBigCode still name the
436 # cache `layer_past`; missing it leaves the cache empty, so every
437 # decode step attends to itself alone and generation ignores the prompt.
438 past_key_values = kwargs.get("layer_past", None)
439 if past_key_values is None:
440 return k, v
441 layer_idx = getattr(self, "_layer_idx", None)
442 if layer_idx is None:
443 logger.warning(
444 "%s: past_key_values provided but _layer_idx is None "
445 "(HF component missing layer_idx attribute). KV cache update "
446 "skipped — cached generation will ignore earlier tokens.",
447 self.name,
448 )
449 return k, v
450 # Some architectures (e.g. HRM-Text's recurrent stacks) pass a cycle_offset
451 # so each stack invocation writes to a unique cache slot. Non-participating
452 # models simply leave it unset.
453 cycle_offset = kwargs.get("cycle_offset", 0)
454 k, v = past_key_values.update(k, v, layer_idx + cycle_offset)
455 return k, v
457 def _reshape_qkv_to_heads(
458 self,
459 q: torch.Tensor,
460 k: torch.Tensor,
461 v: torch.Tensor,
462 num_heads: int,
463 num_kv_heads: int | None = None,
464 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int, int, int]:
465 """Reshape Q/K/V from [batch, seq, hidden] or [batch, seq, heads, head_dim]
466 to [batch, heads, seq, head_dim]. Returns (q, k, v, batch_size, seq_len, head_dim).
468 Args:
469 num_kv_heads: If provided and differs from num_heads (GQA), K/V use
470 this head count for the 3D reshape path.
471 """
472 if num_kv_heads is None:
473 num_kv_heads = num_heads
474 if q.ndim == 3:
475 batch_size, seq_len, q_hidden = q.shape
476 head_dim: int = q_hidden // num_heads
477 q = q.view(batch_size, seq_len, num_heads, head_dim).transpose(1, 2)
478 k = k.view(batch_size, seq_len, num_kv_heads, head_dim).transpose(1, 2)
479 v = v.view(batch_size, seq_len, num_kv_heads, head_dim).transpose(1, 2)
480 elif q.ndim == 4: 480 ↛ 487line 480 didn't jump to line 487 because the condition on line 480 was always true
481 batch_size, seq_len = q.shape[0], q.shape[1]
482 head_dim = q.shape[-1]
483 q = q.transpose(1, 2)
484 k = k.transpose(1, 2)
485 v = v.transpose(1, 2)
486 else:
487 raise ValueError(f"Unexpected Q tensor shape: {q.shape}. Expected 3D or 4D.")
488 return q, k, v, batch_size, seq_len, head_dim
490 def _apply_attn_dropout(self, attn_weights: torch.Tensor) -> torch.Tensor:
491 """Apply attention dropout from the original HF component if present."""
492 if self.original_component is not None: 492 ↛ 498line 492 didn't jump to line 498 because the condition on line 492 was always true
493 dropout_fn = getattr(self.original_component, "attn_dropout", None)
494 if dropout_fn is None:
495 dropout_fn = getattr(self.original_component, "attention_dropout", None)
496 if dropout_fn is not None and callable(dropout_fn):
497 attn_weights = dropout_fn(attn_weights)
498 return attn_weights
500 def _apply_output_projection(self, attn_output: torch.Tensor) -> torch.Tensor:
501 """Apply the output projection (self.o) if present."""
502 if hasattr(self, "o") and self.o is not None:
503 attn_output = self.o(attn_output)
504 return attn_output
506 def _softmax_dropout_pattern(
507 self,
508 attn_scores: torch.Tensor,
509 target_dtype: torch.dtype | None = None,
510 upcast_to_fp32: bool = False,
511 ) -> torch.Tensor:
512 """Apply softmax, dropout, and hook_pattern to attention scores.
514 Args:
515 attn_scores: Raw attention scores [batch, heads, q_seq, kv_seq].
516 target_dtype: If set, cast weights to this dtype after softmax.
517 upcast_to_fp32: If True, compute softmax in float32 for numerical
518 stability, then cast to target_dtype.
519 """
520 if upcast_to_fp32:
521 attn_weights = self._masked_softmax(attn_scores, dtype=torch.float32)
522 if target_dtype is not None: 522 ↛ 528line 522 didn't jump to line 528 because the condition on line 522 was always true
523 attn_weights = attn_weights.to(target_dtype)
524 else:
525 attn_weights = self._masked_softmax(attn_scores)
526 if target_dtype is not None:
527 attn_weights = attn_weights.to(target_dtype)
528 attn_weights = self._scrub_compatibility_pattern_nans(attn_weights)
529 attn_weights = self._apply_attn_dropout(attn_weights)
530 attn_weights = self.hook_pattern(attn_weights)
531 return attn_weights
533 def _masked_softmax(
534 self, attn_scores: torch.Tensor, dtype: torch.dtype | None = None
535 ) -> torch.Tensor:
536 """Softmax over keys that zeroes fully masked rows in compatibility mode.
538 Compatibility mode masks with -inf, so a query with no visible key (e.g.
539 the leading pads of a left-padded row) softmaxes to NaN. Zeroing that NaN
540 afterwards only fixes the forward: softmax's backward still multiplies by
541 the NaN output and leaks it into Q/K and the residual stream.
542 """
543 if not self.compatibility_mode:
544 return torch.nn.functional.softmax(attn_scores, dim=-1, dtype=dtype)
545 fully_masked = torch.isneginf(attn_scores).all(dim=-1, keepdim=True)
546 pattern = torch.nn.functional.softmax(
547 attn_scores.masked_fill(fully_masked, 0.0), dim=-1, dtype=dtype
548 )
549 return pattern.masked_fill(fully_masked, 0.0)
551 def _scrub_compatibility_pattern_nans(self, pattern: torch.Tensor) -> torch.Tensor:
552 """Match HookedTransformer's NaN scrub of the attention pattern."""
553 if self.compatibility_mode:
554 pattern = torch.where(torch.isnan(pattern), torch.zeros_like(pattern), pattern)
555 return pattern
557 def _normalize_compatibility_mask_sentinel(self, attention_mask: torch.Tensor) -> torch.Tensor:
558 """Normalize additive mask sentinels before dtype conversion or addition."""
559 if self.compatibility_mode and attention_mask.is_floating_point():
560 attention_mask = attention_mask.masked_fill(
561 attention_mask <= torch.finfo(attention_mask.dtype).min, -torch.inf
562 )
563 return attention_mask
565 def _reshape_attn_output(
566 self,
567 attn_output: torch.Tensor,
568 batch_size: int,
569 seq_len: int,
570 num_heads: int,
571 head_dim: int,
572 ) -> torch.Tensor:
573 """Reshape attention output from [batch, heads, seq, dim] to [batch, seq, hidden]."""
574 attn_output = attn_output.transpose(1, 2).contiguous()
575 attn_output = attn_output.view(batch_size, seq_len, num_heads * head_dim)
576 return attn_output
578 def _apply_reconstruct_attention_mask(
579 self,
580 attn_scores: torch.Tensor,
581 attention_mask: torch.Tensor | None,
582 seq_len: int,
583 q_seq_len: int | None = None,
584 ) -> torch.Tensor:
585 """Apply causal and optional attention masking to reconstructed scores.
587 HuggingFace-style 4D masks already encode causal semantics, so they are
588 treated as authoritative. Lower-rank masks do not, so the local causal
589 mask is still applied before adding the caller-provided padding mask.
591 Args:
592 attn_scores: Attention scores [batch, heads, q_seq_len, kv_seq_len].
593 attention_mask: Optional mask from the caller.
594 seq_len: The KV sequence length (total positions including cache).
595 q_seq_len: The query sequence length. When using KV cache this is
596 shorter than seq_len. Defaults to seq_len when not provided.
597 """
598 if q_seq_len is None:
599 q_seq_len = seq_len
600 min_dtype = torch.finfo(attn_scores.dtype).min
601 mask_value = -torch.inf if self.compatibility_mode else min_dtype
602 use_direct_hf_mask = attention_mask is not None and attention_mask.ndim >= 4
603 # Bidirectional attention (encoders) and cross-attention have no causal
604 # structure, so only synthesize the triangular mask for causal self-attention.
605 apply_causal = self.is_causal and not self.is_cross_attention
606 if not use_direct_hf_mask and apply_causal:
607 # Rectangular causal mask: query i attends to KV 0..(offset+i)
608 # where offset = kv_seq_len - q_seq_len (cached positions).
609 causal_mask = torch.ones(
610 q_seq_len, seq_len, device=attn_scores.device, dtype=torch.bool
611 )
612 causal_mask = torch.tril(causal_mask, diagonal=seq_len - q_seq_len)
613 attn_scores = attn_scores.masked_fill(~causal_mask, mask_value)
615 if attention_mask is not None:
616 if attention_mask.shape[-1] != seq_len: 616 ↛ 617line 616 didn't jump to line 617 because the condition on line 616 was never true
617 attention_mask = attention_mask[..., :seq_len]
618 if attention_mask.ndim >= 3 and attention_mask.shape[-2] != q_seq_len:
619 # Extra query rows mean a full-sequence mask on a cached decode step,
620 # where the live queries are the LAST rows; taking the first hands
621 # every step position 0 (Baichuan-13B fuses ALiBi slopes in here).
622 attention_mask = attention_mask[..., -q_seq_len:, :]
624 if attention_mask.dtype == torch.bool:
625 attention_mask = torch.where(
626 attention_mask,
627 torch.zeros((), dtype=attn_scores.dtype, device=attn_scores.device),
628 torch.full((), mask_value, dtype=attn_scores.dtype, device=attn_scores.device),
629 )
630 else:
631 attention_mask = self._normalize_compatibility_mask_sentinel(attention_mask)
632 attention_mask = attention_mask.to(dtype=attn_scores.dtype)
633 attn_scores = attn_scores + attention_mask
635 return attn_scores
637 def _get_n_heads(self, use_kv: bool = False) -> int:
638 """Resolve the number of attention heads from config.
640 Args:
641 use_kv: If True, return n_key_value_heads (for GQA) when available.
642 """
643 assert self.config is not None, "config required to resolve n_heads"
644 if use_kv:
645 if hasattr(self.config, "n_key_value_heads") and self.config.n_key_value_heads:
646 return self.config.n_key_value_heads
647 if hasattr(self.config, "n_heads"):
648 return self.config.n_heads
649 return self.config.n_head
651 def _weight_layout_in_out(self, proj: Any) -> Optional[bool]:
652 """Whether ``proj``'s wrapped module stores its weight as [in, out].
654 Conv1D (GPT-2 style) stores [in_features, out_features]; nn.Linear stores
655 [out_features, in_features]. Returns None when the wrapped module is
656 neither, so callers can fall back to a shape heuristic.
657 """
658 component = getattr(proj, "original_component", None)
659 if isinstance(component, Conv1D):
660 return True
661 if isinstance(component, torch.nn.Linear):
662 return False
663 return None
665 def _reshape_weight_to_3d(
666 self,
667 weight: torch.Tensor,
668 n_heads: int,
669 pattern: str = "qkv",
670 in_out_layout: Optional[bool] = None,
671 ) -> torch.Tensor:
672 """Reshape a 2D weight to 3D by splitting heads.
674 Args:
675 weight: 2D weight tensor
676 n_heads: Number of heads to split into
677 pattern: "qkv" for [n_heads, d_model, d_head], "o" for [n_heads, d_head, d_model]
678 in_out_layout: True if the weight is stored [in, out] (Conv1D), False
679 for [out, in] (nn.Linear). None falls back to a shape heuristic,
680 which cannot distinguish the two when the weight is square.
681 """
682 if pattern == "o":
683 if in_out_layout is None:
684 # Heuristic: assumes [in, out] whenever the head split fits dim 0.
685 in_out_layout = weight.shape[0] == n_heads * (
686 weight.shape[1] // n_heads
687 if weight.shape[1] % n_heads == 0
688 else weight.shape[0] // n_heads
689 )
690 mat = weight if in_out_layout else weight.T
691 self._check_head_split_width(mat, n_heads)
692 return einops.rearrange(
693 mat, "(n_heads d_head) d_model -> n_heads d_head d_model", n_heads=n_heads
694 )
695 # QKV pattern
696 if in_out_layout is None:
697 in_out_layout = weight.shape[0] % n_heads != 0
698 mat = weight.T if in_out_layout else weight
699 self._check_head_split_width(mat, n_heads)
700 return einops.rearrange(
701 mat, "(n_heads d_head) d_model -> n_heads d_model d_head", n_heads=n_heads
702 )
704 def _check_head_split_width(self, mat: torch.Tensor, n_heads: int) -> None:
705 """Refuse a head split whose width disagrees with the model's geometry.
707 einops only needs divisibility, so per-layer geometry (gemma4 KV-shared
708 layers) would factorize into wrong-shaped heads silently. MLA's
709 bind-time ``_v_head_dim`` is a legitimate second width (o_proj).
710 """
711 d_head = getattr(self.config, "d_head", None)
712 if not d_head:
713 return
714 allowed_dims = {d_head}
715 v_head_dim = getattr(self, "_v_head_dim", None)
716 if v_head_dim:
717 allowed_dims.add(v_head_dim)
718 if mat.shape[0] % n_heads == 0 and mat.shape[0] // n_heads in allowed_dims:
719 return
720 raise PerLayerGeometryError(
721 f"Cannot split {tuple(mat.shape)} weight on '{self.name}' into "
722 f"{n_heads} heads of d_head in {sorted(allowed_dims)}: this "
723 "layer's attention geometry differs from the config-level scalars "
724 "(per-layer num_key_value_heads/head_dim, e.g. gemma4 KV-shared "
725 "or K==V layers). Read the wrapped module's weight directly for "
726 "this layer instead."
727 )
729 def _project_per_head_qkv(
730 self,
731 linear_bridge: "GeneralizedComponent",
732 input_4d: torch.Tensor,
733 n_heads: int,
734 d_head: int,
735 ) -> torch.Tensor:
736 """Per-head Q/K/V projection over a 4D residual fork.
738 Plain nn.Linear applied to [batch, pos, H, d_model] broadcasts the
739 same weight across heads' copies — which for the split-qkv fork means
740 head h's copy sees every head's W rows, not just head h's. This routes
741 head h's copy through head h's W slice only via a per-head einsum.
743 Fires `linear_bridge.hook_out` on the flat 3D tensor so the hook sees
744 the same shape as the default path and downstream code receives a
745 consistent 4D `[B, S, H, d_head]` regardless of whether the user's
746 hook modified the tensor (which would otherwise trigger the
747 `hook_conversion.revert` 4D→3D flatten).
748 """
749 component = linear_bridge.original_component
750 assert component is not None, "LinearBridge.original_component not set"
751 weight = require_readable_weight(
752 component.weight,
753 operation="project per head (use_split_qkv_input / use_attn_in)",
754 owner=component,
755 )
756 bias = component.bias
757 w3d = einops.rearrange(
758 weight,
759 "(n_heads d_head) d_model -> n_heads d_model d_head",
760 n_heads=n_heads,
761 d_head=d_head,
762 )
763 out = torch.einsum("bshd,hde->bshe", input_4d, w3d)
764 if bias is not None:
765 b2d = einops.rearrange(bias, "(n_heads d_head) -> n_heads d_head", n_heads=n_heads)
766 assert isinstance(b2d, torch.Tensor)
767 out = out + b2d
768 # Flatten to 3D for hook_out (matches default-path shape); the
769 # hook_conversion reshapes to 4D for the user's fwd_hook, then reverts
770 # to 3D if the hook returned a modified tensor. Return 4D always.
771 b, s = out.shape[0], out.shape[1]
772 out_flat = out.reshape(b, s, n_heads * d_head)
773 out_flat = linear_bridge.hook_out(out_flat)
774 return out_flat.reshape(b, s, n_heads, d_head)
776 def _compute_per_head_result(
777 self,
778 z_4d: torch.Tensor,
779 n_heads: int,
780 d_head: int,
781 ) -> torch.Tensor:
782 """Per-head attention output pre-sum across heads.
784 Computes (z[..., h, :] @ W_O_per_head[h]) for each head h, fires
785 hook_result on the resulting [batch, pos, n_heads, d_model], then sums
786 across heads and adds b_O. Distributive over weight folding
787 (`sum_h z_h @ W_O_h + b_O == z_flat @ W_O.T + b_O`), so compat-mode and
788 raw-weight paths produce identical logits.
789 """
790 o = self.o.original_component
791 weight = o.weight
792 bias = getattr(o, "bias", None)
793 # HF Conv1D (GPT-2, GPT-J, CodeGen) stores weight as [in, out]; nn.Linear
794 # stores [out, in]. When W_O is square (d_model == n_heads*d_head, which
795 # is the common case), shape alone is ambiguous — dispatch on module
796 # type instead.
797 weight_is_in_out = type(o).__name__ == "Conv1D"
798 if weight_is_in_out:
799 w_per_head = einops.rearrange(
800 weight,
801 "(n_heads d_head) d_model -> n_heads d_head d_model",
802 n_heads=n_heads,
803 d_head=d_head,
804 )
805 else:
806 w_per_head = einops.rearrange(
807 weight,
808 "d_model (n_heads d_head) -> n_heads d_head d_model",
809 n_heads=n_heads,
810 d_head=d_head,
811 )
812 per_head = torch.einsum("bshd,hdm->bshm", z_4d, w_per_head)
813 scale = getattr(self, "_result_contribution_scale", 1.0)
814 if scale == 1.0:
815 # Keep callable hook replacements on the original unscaled path.
816 per_head = self.hook_result(per_head)
817 else:
818 per_head = self._hook_scaled_output(per_head, self.hook_result, scale)
819 summed = per_head.sum(dim=-2)
820 if bias is not None:
821 summed = summed + bias
822 return summed
824 def forward(self, *args: Any, **kwargs: Any) -> Any:
825 """Simplified forward pass - minimal wrapping around original component.
827 This does minimal wrapping: hook_in → delegate to HF → hook_out.
828 This ensures we match HuggingFace's exact output without complex intermediate processing.
830 Args:
831 *args: Input arguments to pass to the original component
832 **kwargs: Input keyword arguments to pass to the original component
834 Returns:
835 The output from the original component, with only input/output hooks applied
836 """
837 if self.original_component is None: 837 ↛ 838line 837 didn't jump to line 838 because the condition on line 837 was never true
838 raise RuntimeError(
839 f"Original component not set for {self.name}. Call set_original_component() first."
840 )
841 if "query_input" in kwargs:
842 hooked = self.hook_in(kwargs["query_input"])
843 kwargs["query_input"] = hooked
844 elif "hidden_states" in kwargs:
845 hooked = self.hook_in(kwargs["hidden_states"])
846 kwargs["hidden_states"] = hooked
847 elif len(args) > 0 and isinstance(args[0], torch.Tensor): 847 ↛ 853line 847 didn't jump to line 853 because the condition on line 847 was always true
848 hooked = self.hook_in(args[0])
849 args = (hooked,) + args[1:]
850 # Modules exposing the pattern seam (NativeAttention) run the pattern
851 # through hook_pattern INSIDE the computation, so hook edits re-weight
852 # the output; post-hoc firing below cannot (see _fire_post_hoc_hook).
853 pattern_hooked_inside = bool(getattr(self.original_component, "accepts_pattern_fn", False))
854 if pattern_hooked_inside:
855 kwargs["pattern_fn"] = self.hook_pattern
856 # Relative-position-bias modules (T5 family) return position_bias, not
857 # attention weights, in output[1]: force the real weights into the
858 # output and capture Q/K to rebuild the pre-softmax scores.
859 reroute_position_bias = self.requires_relative_position_bias
860 caller_output_attentions = bool(kwargs.get("output_attentions", False))
861 qk_capture: Dict[str, torch.Tensor] = {}
862 capture_handles: list = []
863 if reroute_position_bias:
864 kwargs["output_attentions"] = True
865 capture_handles = self._capture_qk_outputs(qk_capture)
866 # try/finally so the captured tensor (and its autograd graph) is
867 # released even if original_component raises.
868 try:
869 output = self.original_component(*args, **kwargs)
870 finally:
871 self._captured_pre_ln_residual = None
872 for handle in capture_handles:
873 handle.remove()
874 if isinstance(output, tuple) and len(output) >= 2:
875 # output[0] is attention output
876 # output[1] may be attention weights (pattern) or position_bias (T5)
877 # Additional elements may include position_bias, attention weights, etc.
878 attn_output = self.hook_out(output[0])
879 if reroute_position_bias:
880 output = (attn_output,) + self._fire_position_bias_module_hooks(
881 output[1:], qk_capture, caller_output_attentions
882 )
883 else:
884 second_element = output[1]
885 # Fire hook_pattern if the second element is attention weights (4D tensor)
886 if isinstance(second_element, torch.Tensor) and second_element.dim() == 4:
887 # This looks like attention weights [batch, heads, seq, seq]
888 if not pattern_hooked_inside:
889 second_element = self._fire_post_hoc_hook(
890 self.hook_pattern, second_element, "hook_pattern"
891 )
892 # Also store for potential hook_attn_scores (before softmax)
893 # Note: Most HF implementations return post-softmax weights
894 self._fire_post_hoc_hook(
895 self.hook_attn_scores, second_element, "hook_attn_scores"
896 )
897 output = (attn_output, second_element) + output[2:]
898 elif isinstance(output, tuple) and len(output) == 1: 898 ↛ 899line 898 didn't jump to line 899 because the condition on line 898 was never true
899 output = (self.hook_out(output[0]),)
900 else:
901 output = self.hook_out(output)
902 return output
904 def _capture_qk_outputs(self, store: Dict[str, torch.Tensor]) -> list:
905 """Capture the Q/K projection outputs during the wrapped module's forward.
907 The relative-position-bias path rebuilds pre-softmax scores from these
908 (the wrapped module exposes neither the scores nor a seam to hook them).
909 Returns the handles; the caller removes them after the forward.
910 """
911 handles = []
912 for key in ("q", "k"):
913 proj = self._modules.get(key)
914 if proj is None: 914 ↛ 917line 914 didn't jump to line 917 because the condition on line 914 was always true
915 continue
917 def make_capture(name: str):
918 def capture(module: Any, inputs: Any, out: Any) -> None:
919 if isinstance(out, torch.Tensor): 919 ↛ exitline 919 didn't return from function 'capture' because the condition on line 919 was always true
920 store[name] = out
922 return capture
924 handles.append(proj.register_forward_hook(make_capture(key)))
925 return handles
927 def _fire_position_bias_module_hooks(
928 self,
929 rest: Tuple[Any, ...],
930 qk_capture: Dict[str, torch.Tensor],
931 caller_output_attentions: bool,
932 ) -> Tuple[Any, ...]:
933 """Fire pattern/scores hooks for a module whose output[1] is position_bias.
935 ``rest`` is the wrapped module's output after the attention output:
936 (position_bias, attn_weights?) with the weights present because forward
937 forced ``output_attentions``. The forced element is stripped again
938 unless the caller asked for it, so upstream tuple indexing (the HF
939 stack reads fixed positions) stays intact. The bias itself gets no
940 hooks: it is a broadcast [1, heads, q, k] additive term, not a pattern.
941 """
942 position_bias = rest[0]
943 weights = rest[1] if len(rest) >= 2 else None
944 if not (isinstance(weights, torch.Tensor) and weights.dim() == 4): 944 ↛ 945line 944 didn't jump to line 945 because the condition on line 944 was never true
945 return rest
946 scores = self._reconstruct_bias_scores(qk_capture, position_bias, weights)
947 if scores is not None:
948 self._fire_post_hoc_hook(self.hook_attn_scores, scores, "hook_attn_scores")
949 weights = self._fire_post_hoc_hook(self.hook_pattern, weights, "hook_pattern")
950 if caller_output_attentions or not self._output_attentions_is_optional():
951 return (position_bias, weights) + rest[2:]
952 return (position_bias,) + rest[2:]
954 def _output_attentions_is_optional(self) -> bool:
955 """Whether the wrapped module only returns weights when asked to.
957 Transformers >= 5.15 T5 attention always returns three values, which callers unpack.
958 """
959 if self.original_component is None: 959 ↛ 960line 959 didn't jump to line 960 because the condition on line 959 was never true
960 return True
961 try:
962 params = inspect.signature(self.original_component.forward).parameters
963 except (TypeError, ValueError):
964 return True
965 return "output_attentions" in params
967 def _reconstruct_bias_scores(
968 self,
969 qk_capture: Dict[str, torch.Tensor],
970 position_bias: Any,
971 attn_weights: torch.Tensor,
972 ) -> Optional[torch.Tensor]:
973 """Rebuild pre-softmax scores as q @ k^T + position_bias.
975 T5-family attention applies no 1/sqrt(d) scaling, and the returned bias
976 already folds in the attention mask, so this reproduces the exact tensor
977 HF softmaxed. Returns None when Q or K was not captured this forward
978 (e.g. cross-attention K served from an encoder-decoder cache) or the
979 captured shapes disagree with the weights (e.g. cached decode steps).
980 """
981 q = qk_capture.get("q")
982 k = qk_capture.get("k")
983 if q is None or k is None or q.dim() != 3 or k.dim() != 3:
984 return None
985 batch, n_heads, q_len, k_len = attn_weights.shape
986 if (
987 q.shape[0] != batch
988 or q.shape[1] != q_len
989 or k.shape[0] != batch
990 or k.shape[1] != k_len
991 or q.shape[2] != k.shape[2]
992 or q.shape[2] % n_heads != 0
993 ):
994 return None
995 q4 = q.view(batch, q_len, n_heads, -1).transpose(1, 2)
996 k4 = k.view(batch, k_len, n_heads, -1).transpose(1, 2)
997 scores = torch.matmul(q4, k4.transpose(3, 2))
998 if isinstance(position_bias, torch.Tensor): 998 ↛ 1000line 998 didn't jump to line 1000 because the condition on line 998 was always true
999 scores = scores + position_bias
1000 return scores
1002 def _fire_post_hoc_hook(
1003 self, hook: HookPoint, tensor: torch.Tensor, hook_label: str
1004 ) -> torch.Tensor:
1005 """Fire a hook on a tensor the wrapped HF module already consumed.
1007 An edit returned from such a hook cannot flow back into the attention
1008 output, so it earns a warning instead of silently doing nothing.
1009 Read-only hooks (e.g. caching) return the tensor unchanged and stay
1010 silent.
1011 """
1012 result = hook(tensor)
1013 if isinstance(result, torch.Tensor) and result is not tensor:
1014 try:
1015 unchanged = result.shape == tensor.shape and torch.equal(result, tensor)
1016 except RuntimeError:
1017 unchanged = False
1018 if not unchanged:
1019 warnings.warn(
1020 f"A hook on {self.name or type(self).__name__}.{hook_label} returned a "
1021 "modified tensor, but this attention delegates its computation to the "
1022 "wrapped HF module, which has already consumed the original values. The "
1023 "edit is discarded and does not affect the model's output.",
1024 UserWarning,
1025 stacklevel=2,
1026 )
1027 return result
1029 @property
1030 def W_Q(self) -> torch.Tensor:
1031 """Get W_Q in 3D format [n_heads, d_model, d_head].
1033 Gated query projections retain their live query-and-gate parameter;
1034 this analysis view selects the query rows interleaved within each head.
1035 """
1036 weight = require_readable_weight(
1037 self.q.weight, operation=f"read W_Q from {self.name}", owner=self.q
1038 )
1039 if weight.ndim == 2 and self.config is not None: 1039 ↛ 1060line 1039 didn't jump to line 1060 because the condition on line 1039 was always true
1040 n_heads = self._get_n_heads()
1041 in_out_layout = self._weight_layout_in_out(self.q)
1042 if getattr(self.config, "gated_q_proj", False):
1043 d_head = int(self.config.d_head)
1044 gated_width = n_heads * d_head * 2
1045 if in_out_layout is True: 1045 ↛ 1046line 1045 didn't jump to line 1046 because the condition on line 1045 was never true
1046 output_first_weight = weight.T
1047 elif in_out_layout is False or weight.shape[0] == gated_width: 1047 ↛ 1049line 1047 didn't jump to line 1049 because the condition on line 1047 was always true
1048 output_first_weight = weight
1049 elif weight.shape[1] == gated_width:
1050 output_first_weight = weight.T
1051 else:
1052 output_first_weight = None
1054 if output_first_weight is not None and output_first_weight.shape[0] == gated_width: 1054 ↛ 1059line 1054 didn't jump to line 1059 because the condition on line 1054 was always true
1055 # Preserve the live query-gate projection; W_Q is an analysis-only query view.
1056 per_head_weight = output_first_weight.reshape(n_heads, d_head * 2, -1)
1057 return per_head_weight[:, :d_head, :].transpose(-1, -2)
1059 return self._reshape_weight_to_3d(weight, n_heads, in_out_layout=in_out_layout)
1060 return weight
1062 @property
1063 def W_K(self) -> torch.Tensor:
1064 """Get W_K in 3D format [n_heads, d_model, d_head] (uses n_kv_heads for GQA)."""
1065 weight = require_readable_weight(
1066 self.k.weight, operation=f"read W_K from {self.name}", owner=self.k
1067 )
1068 if weight.ndim == 2 and self.config is not None: 1068 ↛ 1074line 1068 didn't jump to line 1074 because the condition on line 1068 was always true
1069 return self._reshape_weight_to_3d(
1070 weight,
1071 self._get_n_heads(use_kv=True),
1072 in_out_layout=self._weight_layout_in_out(self.k),
1073 )
1074 return weight
1076 @property
1077 def W_V(self) -> torch.Tensor:
1078 """Get W_V in 3D format [n_heads, d_model, d_head] (uses n_kv_heads for GQA)."""
1079 weight = require_readable_weight(
1080 self.v.weight, operation=f"read W_V from {self.name}", owner=self.v
1081 )
1082 if weight.ndim == 2 and self.config is not None: 1082 ↛ 1088line 1082 didn't jump to line 1088 because the condition on line 1082 was always true
1083 return self._reshape_weight_to_3d(
1084 weight,
1085 self._get_n_heads(use_kv=True),
1086 in_out_layout=self._weight_layout_in_out(self.v),
1087 )
1088 return weight
1090 @property
1091 def W_O(self) -> torch.Tensor:
1092 """Get W_O in 3D format [n_heads, d_head, d_model]."""
1093 weight = require_readable_weight(
1094 self.o.weight, operation=f"read W_O from {self.name}", owner=self.o
1095 )
1096 if weight.ndim == 2 and self.config is not None: 1096 ↛ 1103line 1096 didn't jump to line 1103 because the condition on line 1096 was always true
1097 return self._reshape_weight_to_3d(
1098 weight,
1099 self._get_n_heads(),
1100 pattern="o",
1101 in_out_layout=self._weight_layout_in_out(self.o),
1102 )
1103 return weight
1105 def _reshape_bias(
1106 self, bias: Optional[torch.Tensor], use_kv: bool = False
1107 ) -> Optional[torch.Tensor]:
1108 """Reshape 1D bias to [n_heads, d_head]."""
1109 if bias is not None and bias.ndim == 1 and self.config is not None:
1110 n_heads = self._get_n_heads(use_kv=use_kv)
1111 return einops.rearrange(bias, "(n_heads d_head) -> n_heads d_head", n_heads=n_heads)
1112 return bias
1114 @property
1115 def b_Q(self) -> Optional[torch.Tensor]:
1116 """Get b_Q in 2D format [n_heads, d_head]."""
1117 return self._reshape_bias(self.q.bias)
1119 @property
1120 def b_K(self) -> Optional[torch.Tensor]:
1121 """Get b_K in 2D format [n_heads, d_head] (uses n_kv_heads for GQA)."""
1122 return self._reshape_bias(self.k.bias, use_kv=True)
1124 @property
1125 def b_V(self) -> Optional[torch.Tensor]:
1126 """Get b_V in 2D format [n_heads, d_head] (uses n_kv_heads for GQA)."""
1127 return self._reshape_bias(self.v.bias, use_kv=True)
1129 @property
1130 def b_O(self) -> Optional[torch.Tensor]:
1131 """Get b_O bias from linear bridge."""
1132 return self.o.bias