Coverage for transformer_lens/model_bridge/generalized_components/joint_qkv_attention.py: 86%
243 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-21 19:27 +0000
1"""Joint QKV attention bridge component.
3This module contains the bridge component for attention layers that use a fused qkv matrix.
4"""
5import copy
6from typing import Any, Callable, Dict, Optional
8import einops
9import torch
11from transformer_lens.conversion_utils.conversion_steps.base_tensor_conversion import (
12 BaseTensorConversion,
13)
14from transformer_lens.model_bridge.generalized_components.attention import (
15 AttentionBridge,
16)
17from transformer_lens.model_bridge.generalized_components.base import (
18 GeneralizedComponent,
19 align_offloaded_subtree,
20)
21from transformer_lens.model_bridge.generalized_components.linear import LinearBridge
22from transformer_lens.utilities.quantization import require_readable_weight
25class JointQKVAttentionBridge(AttentionBridge):
26 """Joint QKV attention bridge that wraps a joint qkv linear layer.
28 This component wraps attention layers that use a fused qkv matrix such that
29 the individual activations from the separated q, k, and v matrices are hooked and accessible.
30 """
32 supports_attn_result: bool = True
34 # property_aliases inherited from AttentionBridge (W_Q, W_K, W_V, W_O, b_Q, b_K, b_V, b_O)
36 def __init__(
37 self,
38 name: str,
39 config: Any,
40 split_qkv_matrix: Optional[Callable] = None,
41 submodules: Optional[Dict[str, GeneralizedComponent]] = None,
42 qkv_conversion_rule: Optional[BaseTensorConversion] = None,
43 attn_conversion_rule: Optional[BaseTensorConversion] = None,
44 pattern_conversion_rule: Optional[BaseTensorConversion] = None,
45 requires_position_embeddings: bool = False,
46 requires_attention_mask: bool = False,
47 ):
48 """Initialize the Joint QKV attention bridge.
50 Args:
51 name: The name of this component
52 config: Model configuration (required for auto-conversion detection)
53 split_qkv_matrix: Optional function to split the qkv matrix into q, k, and v linear transformations.
54 If None, uses the default implementation that splits a combined c_attn weight/bias.
55 submodules: Dictionary of submodules to register (e.g., q_proj, k_proj, etc.)
56 qkv_conversion_rule: Optional conversion rule for the individual q, k, and v matrices to convert their output shapes to HookedTransformer format. If None, uses default RearrangeTensorConversion
57 attn_conversion_rule: Optional conversion rule. Passed to parent AttentionBridge. If None, AttentionAutoConversion will be used
58 pattern_conversion_rule: Optional conversion rule for attention patterns. If None,
59 uses AttentionPatternConversion to ensure [n_heads, pos, pos] shape
60 requires_position_embeddings: Whether this attention requires position_embeddings as input
61 requires_attention_mask: Whether this attention requires attention_mask as input
62 """
63 super().__init__(
64 name,
65 config,
66 submodules=submodules,
67 conversion_rule=attn_conversion_rule,
68 pattern_conversion_rule=pattern_conversion_rule,
69 requires_position_embeddings=requires_position_embeddings,
70 requires_attention_mask=requires_attention_mask,
71 )
72 self.split_qkv_matrix = (
73 split_qkv_matrix if split_qkv_matrix is not None else self._default_split_qkv_matrix
74 )
75 if qkv_conversion_rule is not None:
76 self.qkv_conversion_rule = qkv_conversion_rule
77 else:
78 self.qkv_conversion_rule = self._create_qkv_conversion_rule()
79 self.q = LinearBridge(name="q")
80 self.k = LinearBridge(name="k")
81 self.v = LinearBridge(name="v")
82 for submodule_name, submodule in (submodules or {}).items():
83 if not hasattr(self, submodule_name): 83 ↛ 82line 83 didn't jump to line 82 because the condition on line 83 was always true
84 setattr(self, submodule_name, submodule)
85 self.submodules["q"] = self.q
86 self.submodules["k"] = self.k
87 self.submodules["v"] = self.v
88 self.q.hook_out.hook_conversion = self.qkv_conversion_rule
89 self.k.hook_out.hook_conversion = self.qkv_conversion_rule
90 self.v.hook_out.hook_conversion = self.qkv_conversion_rule
92 # Register q, k, v LinearBridges in real_components for weight distribution
93 # This allows set_processed_weights to distribute weights to these submodules
94 self.real_components["q"] = ("q", self.q)
95 self.real_components["k"] = ("k", self.k)
96 self.real_components["v"] = ("v", self.v)
97 if hasattr(self, "o"):
98 self.real_components["o"] = ("o", self.o)
100 self._reference_model: Optional[Any] = None
102 # Exclude stale qkv combined weights from state_dict after splitting.
103 self._register_state_dict_hook(JointQKVAttentionBridge._filter_qkv_state_dict)
104 self.register_load_state_dict_pre_hook(
105 JointQKVAttentionBridge._restore_filtered_qkv_state_dict
106 )
108 def __deepcopy__(self, memo):
109 """Share split_qkv_matrix across clones instead of copying.
111 split_qkv_matrix may be a bound method of the architecture adapter,
112 which transitively references the full HF model. Without this override,
113 deepcopy duplicates the entire model per block (~1GB x N_layers).
114 """
115 saved_split_fn = self.split_qkv_matrix
116 saved_config = self.config
118 self.split_qkv_matrix = None # type: ignore[assignment]
119 self.config = None
120 try:
121 # Remove override from defining class (not subclass) to avoid recursion.
122 owner = JointQKVAttentionBridge
123 override = owner.__dict__["__deepcopy__"]
124 del owner.__deepcopy__
125 try:
126 clone = copy.deepcopy(self, memo)
127 finally:
128 owner.__deepcopy__ = override # type: ignore[method-assign]
129 finally:
130 self.split_qkv_matrix = saved_split_fn
131 self.config = saved_config
133 clone.split_qkv_matrix = saved_split_fn
134 clone.config = self._resolve_cloned_config(saved_config, memo)
135 return clone
137 @staticmethod
138 def _resolve_cloned_config(config: Any, memo: Dict[int, Any]) -> Any:
139 """Pick the config a clone reads: the shared live one, or the cloning bridge's own.
141 Block replication deepcopies a block template while no Bridge owns the config
142 yet, and those clones must keep sharing the live one so every layer stays on a
143 single object. But a whole-bridge deepcopy gives the clone its own cfg, and an
144 attention still pointing at the original's config makes the clone's
145 ``use_attn_result`` rewire the ORIGINAL's forward -- a pristine control model
146 corrupted with no error. The owning Bridge sitting in ``memo`` is what tells the
147 two cases apart; copying through ``memo`` then hands back the very object that
148 Bridge clone adopts as its cfg instead of forking a third one.
149 """
150 if config is None: 150 ↛ 151line 150 didn't jump to line 151 because the condition on line 150 was never true
151 return None
152 bridge_ref = getattr(config, "_bridge_ref", None)
153 live_bridge = bridge_ref() if bridge_ref is not None else None
154 if live_bridge is None or id(live_bridge) not in memo:
155 return config
156 return copy.deepcopy(config, memo)
158 @staticmethod
159 def _filter_qkv_state_dict(
160 module: torch.nn.Module,
161 state_dict: Dict[str, Any],
162 prefix: str,
163 local_metadata: Dict[str, Any],
164 ) -> None:
165 """State dict hook that removes stale combined QKV entries."""
166 qkv_prefix = prefix + "qkv."
167 keys_to_remove = [k for k in state_dict if k.startswith(qkv_prefix)]
168 for k in keys_to_remove:
169 del state_dict[k]
171 @staticmethod
172 def _restore_filtered_qkv_state_dict(
173 module: torch.nn.Module,
174 state_dict: Dict[str, Any],
175 prefix: str,
176 local_metadata: Dict[str, Any],
177 strict: bool,
178 missing_keys: list[str],
179 unexpected_keys: list[str],
180 error_msgs: list[str],
181 ) -> None:
182 """Insert current combined weights only to satisfy strict key matching.
184 Production checkpoints restore authoritative values through the unfiltered
185 Hugging Face ``_original_component`` path.
186 """
187 del local_metadata, strict, missing_keys, unexpected_keys, error_msgs
188 qkv = module._modules.get("qkv")
189 if qkv is None: 189 ↛ 190line 189 didn't jump to line 190 because the condition on line 189 was never true
190 return
191 for key, value in qkv.state_dict(prefix=f"{prefix}qkv.").items():
192 state_dict.setdefault(key, value)
194 def _create_qkv_conversion_rule(self) -> BaseTensorConversion:
195 """Create the appropriate conversion rule for the individual q, k, and v matrices.
197 Returns:
198 BaseTensorConversion for individual q, k, and v matrices
199 """
200 assert self.config is not None
202 class ConditionalRearrangeConversion(BaseTensorConversion):
203 def __init__(self, n_heads: int):
204 super().__init__()
205 self.n_heads = n_heads
206 self.pattern = (
207 "batch seq (num_attention_heads d_head) -> batch seq num_attention_heads d_head"
208 )
210 def handle_conversion(self, input_value: torch.Tensor, *full_context) -> torch.Tensor:
211 if input_value.ndim == 4: 211 ↛ 212line 211 didn't jump to line 212 because the condition on line 211 was never true
212 return input_value
213 elif input_value.ndim == 3: 213 ↛ 218line 213 didn't jump to line 218 because the condition on line 213 was always true
214 return einops.rearrange(
215 input_value, self.pattern, num_attention_heads=self.n_heads
216 )
217 else:
218 raise ValueError(
219 f"Expected 3D or 4D tensor, got {input_value.ndim}D with shape {input_value.shape}"
220 )
222 def revert(self, input_value: torch.Tensor, *full_context) -> torch.Tensor:
223 if input_value.ndim == 3: 223 ↛ 224line 223 didn't jump to line 224 because the condition on line 223 was never true
224 return input_value
225 elif input_value.ndim == 4: 225 ↛ 232line 225 didn't jump to line 232 because the condition on line 225 was always true
226 return einops.rearrange(
227 input_value,
228 "batch seq num_attention_heads d_head -> batch seq (num_attention_heads d_head)",
229 num_attention_heads=self.n_heads,
230 )
231 else:
232 raise ValueError(
233 f"Expected 3D or 4D tensor, got {input_value.ndim}D with shape {input_value.shape}"
234 )
236 return ConditionalRearrangeConversion(self.config.n_heads)
238 def _default_split_qkv_matrix(
239 self, original_attention_component: Any
240 ) -> tuple[torch.nn.Module, torch.nn.Module, torch.nn.Module]:
241 """Default implementation to split the QKV matrix into separate linear transformations.
243 This uses the 'qkv' submodule defined in component_mapping to find the combined QKV weights.
244 Assumes combined QKV weights in the format [d_model, 3 * d_model] for weights
245 and [3 * n_head * d_head] for bias.
247 Args:
248 original_attention_component: The original attention layer component
249 Returns:
250 Tuple of nn.Linear modules for Q, K, and V transformations
251 """
252 assert self.config is not None
253 assert original_attention_component is not None
255 # Get the combined QKV component using the 'qkv' submodule name
256 if "qkv" not in self.submodules: 256 ↛ 257line 256 didn't jump to line 257 because the condition on line 256 was never true
257 raise ValueError(
258 "No 'qkv' submodule found in JointQKVAttentionBridge. "
259 "Please define a 'qkv' submodule or provide a custom split_qkv_matrix function."
260 )
262 # Get the actual qkv component name from the bridge
263 qkv_bridge = self.submodules["qkv"]
264 qkv_name = qkv_bridge.name
266 # Ensure qkv_name is not None
267 if qkv_name is None: 267 ↛ 268line 267 didn't jump to line 268 because the condition on line 267 was never true
268 raise ValueError(
269 "qkv bridge name is None. " "Please provide a custom split_qkv_matrix function."
270 )
272 # Navigate to the component using the name
273 if not hasattr(original_attention_component, qkv_name): 273 ↛ 274line 273 didn't jump to line 274 because the condition on line 273 was never true
274 raise ValueError(
275 f"Cannot find '{qkv_name}' in attention component. "
276 f"Available attributes: {dir(original_attention_component)}. "
277 f"Please provide a custom split_qkv_matrix function."
278 )
280 qkv_component = getattr(original_attention_component, qkv_name)
282 # Before the tensor_split/nn.Parameter below: int8 and uint8 would die
283 # there with an opaque "Only Tensors of floating point ... can require
284 # gradients", and float8 would split SILENTLY into scale-less pieces.
285 qkv_weights = require_readable_weight(
286 qkv_component.weight,
287 operation="split a fused QKV projection at boot",
288 owner=qkv_component,
289 )
291 # Original qkv_weights shape: [d_model, 3 * d_model]
292 # Split into three equal parts along dimension 1 to get Q, K, V weights
293 q_weight, k_weight, v_weight = torch.tensor_split(qkv_weights, 3, dim=1)
295 # Handle bias if it exists
296 has_bias = hasattr(qkv_component, "bias") and qkv_component.bias is not None
297 q_bias: torch.Tensor | None
298 k_bias: torch.Tensor | None
299 v_bias: torch.Tensor | None
300 if has_bias: 300 ↛ 309line 300 didn't jump to line 309 because the condition on line 300 was always true
301 qkv_bias = qkv_component.bias
302 assert isinstance(qkv_bias, torch.Tensor)
304 # Original qkv_bias shape: [3 * n_head * d_head]
305 # Reshape to [3, n_head * d_head] to split by Q, K, V
306 qkv_bias = qkv_bias.reshape(3, self.config.n_heads * self.config.d_head)
307 q_bias, k_bias, v_bias = qkv_bias[0, :], qkv_bias[1, :], qkv_bias[2, :]
308 else:
309 q_bias = k_bias = v_bias = None
311 # Create plain nn.Linear modules that output 3D tensors [batch, seq, d_model]
312 q_linear = torch.nn.Linear(q_weight.shape[0], q_weight.shape[1], bias=has_bias)
313 q_linear.weight = torch.nn.Parameter(q_weight.T)
314 if has_bias and q_bias is not None: 314 ↛ 317line 314 didn't jump to line 317 because the condition on line 314 was always true
315 q_linear.bias = torch.nn.Parameter(q_bias)
317 k_linear = torch.nn.Linear(k_weight.shape[0], k_weight.shape[1], bias=has_bias)
318 k_linear.weight = torch.nn.Parameter(k_weight.T)
319 if has_bias and k_bias is not None: 319 ↛ 322line 319 didn't jump to line 322 because the condition on line 319 was always true
320 k_linear.bias = torch.nn.Parameter(k_bias)
322 v_linear = torch.nn.Linear(v_weight.shape[0], v_weight.shape[1], bias=has_bias)
323 v_linear.weight = torch.nn.Parameter(v_weight.T)
324 if has_bias and v_bias is not None: 324 ↛ 327line 324 didn't jump to line 327 because the condition on line 324 was always true
325 v_linear.bias = torch.nn.Parameter(v_bias)
327 return q_linear, k_linear, v_linear
329 def set_original_component(self, original_component: torch.nn.Module) -> None:
330 """Set the original component that this bridge wraps and initialize LinearBridges for q, k, v, and o transformations.
332 Args:
333 original_component: The original attention layer to wrap
334 """
335 super().set_original_component(original_component)
337 # Capture HF-specific attention flags for faithful reconstruction
338 self._reorder_and_upcast_attn = getattr(
339 original_component, "reorder_and_upcast_attn", False
340 )
342 # split_qkv_matrix reads original_component's raw weight/bias once, here,
343 # to build independent q/k/v slices - not a live view, so under Accelerate
344 # offload this needs the real (not meta) data materialized for this one
345 # read. The resulting slices are real, standalone tensors that stay valid
346 # afterward regardless of what Accelerate later does to original_component
347 # (unlike whatever reads original_component itself on every forward call,
348 # e.g. GeneralizedComponent.__call__ / LinearBridge for "o"/c_proj, split
349 # q/k/v specifically stay permanently resident rather than re-offloading
350 # after each forward - a deliberate, small, documented memory trade-off
351 # for combined-qkv architectures).
352 with align_offloaded_subtree(original_component):
353 q_transformation, k_transformation, v_transformation = self.split_qkv_matrix(
354 original_component
355 )
356 self.q.set_original_component(q_transformation)
357 self.k.set_original_component(k_transformation)
358 self.v.set_original_component(v_transformation)
359 if hasattr(self, "o") and hasattr(original_component, "c_proj"):
360 self.o.set_original_component(original_component.c_proj)
362 def forward(self, *args: Any, **kwargs: Any) -> Any:
363 """Forward pass through the qkv linear transformation with hooks.
365 Args:
366 *args: Input arguments, where the first argument should be the input tensor
367 **kwargs: Additional keyword arguments
369 Returns:
370 Output tensor after qkv linear transformation
371 """
372 hooked_input = self._apply_attention_input_hook(*args, **kwargs)
373 if self._is_split_qkv_fork_active():
374 q_output, k_output, v_output = self._split_forward_qkv(hooked_input)
375 else:
376 q_output = self.q(hooked_input)
377 k_output = self.k(hooked_input)
378 v_output = self.v(hooked_input)
379 output = self._reconstruct_attention(q_output, k_output, v_output, **kwargs)
380 output = self._process_output(output)
381 return output
383 def _is_split_qkv_fork_active(self) -> bool:
384 cfg = self.config
385 if cfg is None or not getattr(cfg, "n_heads", 0): 385 ↛ 386line 385 didn't jump to line 386 because the condition on line 385 was never true
386 return False
387 return bool(
388 getattr(cfg, "use_split_qkv_input", False) or getattr(cfg, "use_attn_in", False)
389 )
391 def _split_forward_qkv(
392 self, hidden_states: torch.Tensor
393 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
394 """Fork the residual into independent Q/K/V copies, apply per-head projection.
396 After `split_qkv_matrix` runs in `set_original_component`, q/k/v are
397 separate `nn.Linear` modules whose weights partition the output dim by
398 head (output row h*d_head + i ↔ head h, dim i). Plain nn.Linear applied
399 to a 4D [B, S, H, d_model] copy would broadcast the full weight over
400 every head's copy and then we'd keep only the diagonal — n_heads× extra
401 compute. The per-head einsum in `_project_per_head_qkv` slices W per
402 head directly, producing the same 4D [B, S, H, d_head] result that
403 `_reconstruct_attention` expects.
404 """
405 cfg = self.config
406 assert cfg is not None, "config required for split QKV fork"
407 n_heads = int(cfg.n_heads)
408 n_kv_heads = int(getattr(cfg, "n_key_value_heads", None) or n_heads)
409 d_head = int(getattr(cfg, "d_head", 0) or (int(cfg.d_model) // n_heads))
410 use_split = bool(getattr(cfg, "use_split_qkv_input", False))
411 # #1317: fork pre-LN when available so hook patches match legacy.
412 captured = self._captured_pre_ln_residual
413 source = captured if captured is not None else hidden_states
414 if use_split:
415 q_in = self._fork_and_norm_per_head(source, self.hook_q_input, n_heads)
416 k_in = self._fork_and_norm_per_head(source, self.hook_k_input, n_kv_heads)
417 v_in = self._fork_and_norm_per_head(source, self.hook_v_input, n_kv_heads)
418 else:
419 attn_in = self._fork_and_norm_per_head(source, self.hook_attn_in, n_heads)
420 q_in = attn_in
421 if n_kv_heads != n_heads: 421 ↛ 422line 421 didn't jump to line 422 because the condition on line 421 was never true
422 k_in = attn_in[..., :n_kv_heads, :].contiguous()
423 v_in = attn_in[..., :n_kv_heads, :].contiguous()
424 else:
425 k_in = v_in = attn_in
426 q_4d = self._project_per_head_qkv(self.q, q_in, n_heads, d_head)
427 k_4d = self._project_per_head_qkv(self.k, k_in, n_kv_heads, d_head)
428 v_4d = self._project_per_head_qkv(self.v, v_in, n_kv_heads, d_head)
429 return q_4d, k_4d, v_4d
431 def _process_output(self, output: Any) -> Any:
432 """Process the output from _reconstruct_attention.
434 This override skips the duplicate hook_pattern call since
435 _reconstruct_attention already applies both hook_attn_scores
436 and hook_pattern correctly.
438 Args:
439 output: Output tuple from _reconstruct_attention (attn_output, attn_weights)
441 Returns:
442 Processed output with hook_out applied
443 """
444 attn_pattern = None
445 if isinstance(output, tuple) and len(output) >= 2: 445 ↛ 447line 445 didn't jump to line 447 because the condition on line 445 was always true
446 attn_pattern = output[1]
447 if attn_pattern is not None:
448 self._pattern = attn_pattern
449 if isinstance(output, tuple) and len(output) > 0 and isinstance(output[0], torch.Tensor): 449 ↛ 453line 449 didn't jump to line 453 because the condition on line 449 was always true
450 processed_output = list(output)
451 processed_output[0] = self.hook_hidden_states(output[0])
452 output = tuple(processed_output)
453 if isinstance(output, torch.Tensor): 453 ↛ 454line 453 didn't jump to line 454 because the condition on line 453 was never true
454 output = self.hook_out(output)
455 elif isinstance(output, tuple) and len(output) > 0: 455 ↛ 462line 455 didn't jump to line 462 because the condition on line 455 was always true
456 processed_tuple = list(output)
457 if isinstance(output[0], torch.Tensor): 457 ↛ 459line 457 didn't jump to line 459 because the condition on line 457 was always true
458 processed_tuple[0] = self.hook_out(output[0])
459 if len(processed_tuple) == 1: 459 ↛ 460line 459 didn't jump to line 460 because the condition on line 459 was never true
460 return processed_tuple[0]
461 output = tuple(processed_tuple)
462 return output
464 def _apply_attention_input_hook(self, *args: Any, **kwargs: Any) -> torch.Tensor:
465 """Apply attention input hook to the input tensor.
467 This method extracts the input tensor from args/kwargs and applies the attention
468 input hook in the same way as the super class.
470 Args:
471 *args: Input arguments, where the first argument should be the input tensor
472 **kwargs: Additional keyword arguments that might contain input
474 Returns:
475 Input tensor with attention input hook applied
477 Raises:
478 ValueError: If no input tensor is found in args or kwargs
479 """
480 input_tensor = None
481 if "query_input" in kwargs: 481 ↛ 482line 481 didn't jump to line 482 because the condition on line 481 was never true
482 input_tensor = kwargs["query_input"]
483 elif "hidden_states" in kwargs:
484 input_tensor = kwargs["hidden_states"]
485 elif len(args) > 0 and isinstance(args[0], torch.Tensor): 485 ↛ 488line 485 didn't jump to line 488 because the condition on line 485 was always true
486 input_tensor = args[0]
487 else:
488 raise ValueError("No input tensor found in args or kwargs")
489 return self.hook_in(input_tensor)
491 def _reconstruct_attention(
492 self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, **kwargs
493 ) -> tuple:
494 """Manual attention reconstruction used by the bridge after splitting fused QKV projections."""
495 assert self.original_component is not None
496 assert self.config is not None
497 num_heads = self.config.n_heads
498 num_kv_heads = getattr(self.config, "n_key_value_heads", None) or num_heads
500 q, k, v, batch_size, seq_len, head_dim = self._reshape_qkv_to_heads(
501 q, k, v, num_heads, num_kv_heads
502 )
504 # KV cache: extend K/V with cached positions.
505 k, v = self._update_kv_cache(k, v, **kwargs)
507 # GQA/MQA: expand K/V heads to match Q heads
508 if num_kv_heads != num_heads:
509 n_rep = num_heads // num_kv_heads
510 k = k.repeat_interleave(n_rep, dim=1)
511 v = v.repeat_interleave(n_rep, dim=1)
513 # Attention scale: 1/sqrt(d_head) with optional inverse-layer scaling
514 scale = head_dim ** (-0.5)
515 if ( 515 ↛ 520line 515 didn't jump to line 520 because the condition on line 515 was never true
516 hasattr(self.config, "scale_attn_by_inverse_layer_idx")
517 and self.config.scale_attn_by_inverse_layer_idx
518 and self._layer_idx is not None
519 ):
520 scale /= float(self._layer_idx + 1)
522 # When reorder_and_upcast_attn is True, HF computes attention in float32.
523 reorder_and_upcast = getattr(self, "_reorder_and_upcast_attn", False)
524 if reorder_and_upcast: 524 ↛ 525line 524 didn't jump to line 525 because the condition on line 524 was never true
525 q_scores = q.to(torch.float32)
526 k_scores = k.to(torch.float32)
527 else:
528 q_scores = q
529 k_scores = k
531 kv_seq_len = k.shape[-2] # Includes cached positions
532 attn_scores = torch.matmul(q_scores, k_scores.transpose(-2, -1)) * scale
533 attention_mask = kwargs.get("attention_mask", None)
534 attn_scores = self._apply_reconstruct_attention_mask(
535 attn_scores=attn_scores,
536 attention_mask=attention_mask,
537 seq_len=kv_seq_len,
538 q_seq_len=seq_len,
539 )
541 attn_scores = self.hook_attn_scores(attn_scores)
543 attn_weights = self._softmax_dropout_pattern(
544 attn_scores,
545 target_dtype=v.dtype if reorder_and_upcast else None,
546 )
547 attn_output = torch.matmul(attn_weights, v)
548 attn_output = self._reshape_attn_output(
549 attn_output, batch_size, seq_len, num_heads, head_dim
550 )
551 if (
552 bool(getattr(self.config, "use_attn_result", False))
553 and hasattr(self, "o")
554 and self.o.original_component is not None
555 ):
556 # Per-head output pre-sum. Fire hook_z on the pre-projection flat
557 # tensor first so patches at hook_z propagate into the per-head
558 # computation, matching how the default path's `self.o(...)` call
559 # fires o.hook_in before the linear.
560 attn_output = self.o.hook_in(attn_output)
561 z_4d = attn_output.view(batch_size, seq_len, num_heads, head_dim)
562 attn_output = self._compute_per_head_result(z_4d, num_heads, head_dim)
563 else:
564 attn_output = self._apply_output_projection(attn_output)
565 return (attn_output, attn_weights)