Coverage for transformer_lens/model_bridge/generalized_components/position_embeddings_attention.py: 83%
320 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-01 16:23 +0000
1"""Position embeddings attention bridge with full hook support.
3Reimplements attention for models using RoPE (Llama, Gemma, Qwen, OLMo, etc.)
4so that all hook points fire at the correct computation stage:
5- hook_q/hook_k/hook_v: after projection
6- hook_rot_q/hook_rot_k: after RoPE rotation
7- hook_attn_scores: PRE-softmax (matching HookedTransformer convention)
8- hook_pattern: POST-softmax
9"""
10from __future__ import annotations
12import weakref
13from typing import Any, Callable, Dict, Optional
15import torch
16import transformers.models.gemma2.modeling_gemma2 as gemma2_module
18from transformer_lens.hook_points import HookPoint
19from transformer_lens.model_bridge.generalized_components.attention import (
20 AttentionBridge,
21)
22from transformer_lens.model_bridge.generalized_components.position_embedding_hooks_mixin import (
23 PositionEmbeddingHooksMixin,
24)
25from transformer_lens.utilities.attention import clamp_qkv
26from transformer_lens.utilities.heterogeneous_config import safe_config_get
27from transformer_lens.utilities.hf_utils import get_rotary_pct_from_config
29# Global registry mapping HF attention modules to their bridge instances
30# Uses WeakValueDictionary to avoid preventing garbage collection of bridges
31_ATTENTION_BRIDGE_REGISTRY: weakref.WeakValueDictionary = weakref.WeakValueDictionary()
33# Track whether we've already wrapped eager_attention_forward
34_EAGER_ATTENTION_WRAPPED = False
36# Store the original function for restoration
37_ORIGINAL_EAGER_ATTENTION_FORWARD: Optional[Callable] = None
40def _apply_rotary_pos_emb_adjacent_pairs(
41 q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
42) -> tuple[torch.Tensor, torch.Tensor]:
43 """GLM/ERNIE-style RoPE: rotate adjacent element pairs in full precision.
45 cos/sin arrive in the standard half-duplicated layout; this convention
46 takes the first half and expands it by repeat_interleave(2) so rotation
47 pairs are (0,1), (2,3), ... instead of (i, i + d/2).
48 """
49 original_dtype = q.dtype
50 cos = cos.unsqueeze(1)
51 sin = sin.unsqueeze(1)
52 cos = cos[..., : cos.shape[-1] // 2].repeat_interleave(2, dim=-1)
53 sin = sin[..., : sin.shape[-1] // 2].repeat_interleave(2, dim=-1)
55 def _rotate(x: torch.Tensor) -> torch.Tensor:
56 x1 = x[..., 0::2]
57 x2 = x[..., 1::2]
58 return torch.stack((-x2, x1), dim=-1).flatten(-2)
60 q_embed = (q.float() * cos) + (_rotate(q).float() * sin)
61 k_embed = (k.float() * cos) + (_rotate(k).float() * sin)
62 return q_embed.to(original_dtype), k_embed.to(original_dtype)
65def _setup_eager_attention_hook_wrapper() -> None:
66 """Wrap gemma2's eager_attention_forward to fire hook_rot_q and hook_rot_k.
68 This function monkey-patches the module-level eager_attention_forward function
69 to intercept query and key tensors (which have already had rotary embeddings applied)
70 and fire the corresponding hooks on the registered bridge instance.
72 This is safe to call multiple times - it will only wrap once.
73 """
74 global _EAGER_ATTENTION_WRAPPED, _ORIGINAL_EAGER_ATTENTION_FORWARD
76 if _EAGER_ATTENTION_WRAPPED:
77 return
79 # Store the original function
80 _ORIGINAL_EAGER_ATTENTION_FORWARD = gemma2_module.eager_attention_forward
82 def hooked_eager_attention_forward(
83 module: torch.nn.Module,
84 query: torch.Tensor,
85 key: torch.Tensor,
86 value: torch.Tensor,
87 attention_mask: Optional[torch.Tensor],
88 **kwargs: Any,
89 ) -> tuple:
90 """Wrapped eager_attention_forward that fires rotary hooks.
92 Args:
93 module: The HF attention module (used to look up the bridge)
94 query: Query tensor AFTER rotary embeddings applied
95 key: Key tensor AFTER rotary embeddings applied
96 value: Value tensor
97 attention_mask: Attention mask
98 **kwargs: Additional arguments (dropout, scaling, etc.)
100 Returns:
101 Tuple of (attn_output, attn_weights)
102 """
103 # Look up the bridge instance for this attention module
104 bridge = _ATTENTION_BRIDGE_REGISTRY.get(id(module))
106 if bridge is not None: 106 ↛ 108line 106 didn't jump to line 108 because the condition on line 106 was never true
107 # Fire hook_rot_q and hook_rot_k with the post-rotary Q/K
108 if hasattr(bridge, "hook_rot_q"):
109 query = bridge.hook_rot_q(query)
110 if hasattr(bridge, "hook_rot_k"):
111 key = bridge.hook_rot_k(key)
113 # Call the original function
114 assert _ORIGINAL_EAGER_ATTENTION_FORWARD is not None
115 return _ORIGINAL_EAGER_ATTENTION_FORWARD(
116 module, query, key, value, attention_mask, **kwargs
117 )
119 # Replace the module-level function for both Gemma 2 and Gemma 3
120 gemma2_module.eager_attention_forward = hooked_eager_attention_forward # type: ignore[assignment]
122 try:
123 import transformers.models.gemma3.modeling_gemma3 as gemma3_module
125 gemma3_module.eager_attention_forward = hooked_eager_attention_forward # type: ignore[assignment]
126 except ImportError:
127 pass # Gemma 3 not available in this transformers version
129 _EAGER_ATTENTION_WRAPPED = True
132class PositionEmbeddingsAttentionBridge(PositionEmbeddingHooksMixin, AttentionBridge):
133 """Attention bridge for models that require position embeddings (e.g., Gemma-3).
135 Some models use specialized position embedding systems (like Gemma-3's dual RoPE)
136 which require position_embeddings to be generated in a specific format that differs
137 from standard RoPE models.
139 The position_embeddings are generated by calling the model's rotary_emb
140 component with dummy Q/K tensors and position_ids.
141 """
143 supports_attn_result: bool = True
145 # NoPE architectures (EXAONE-4, SmolLM3, Cohere2) deliberately null
146 # position_embeddings on their non-rotary layers to match HF. Their bridge
147 # subclasses set this so the missing-RoPE warning stays meaningful.
148 rope_optional: bool = False
150 def __init__(
151 self,
152 name: str,
153 config: Any,
154 submodules: Optional[Dict[str, Any]] = None,
155 optional: bool = False,
156 # Accepted for caller compatibility (Granite passes these explicitly)
157 # but always forced to True — this bridge reimplements attention.
158 requires_attention_mask: bool = True,
159 requires_position_embeddings: bool = True,
160 is_causal: bool = True,
161 **kwargs, # absorb any other AttentionBridge kwargs callers may pass
162 ):
163 super().__init__(
164 name,
165 config,
166 submodules,
167 requires_position_embeddings=True,
168 requires_attention_mask=True,
169 maintain_native_attention=True,
170 is_causal=is_causal,
171 optional=optional,
172 )
173 self._init_position_embedding_hooks()
174 if getattr(config, "gated_q_proj", False):
175 self.hook_q_gate = HookPoint()
176 # Gate on adapter intent; HF-vs-adapter mismatches surface in set_original_component.
177 if submodules is not None and "gate" in submodules:
178 self.hook_gate = HookPoint()
179 if submodules is not None and "q_norm" in submodules:
180 self.hook_q_normed = HookPoint()
181 if submodules is not None and "k_norm" in submodules:
182 self.hook_k_normed = HookPoint()
183 self._qk_norm_phase: Optional[str] = None
185 def set_original_component(self, component: torch.nn.Module) -> None:
186 """Wire HF module, register for rotary hooks, validate adapter declarations."""
187 super().set_original_component(component)
188 _ATTENTION_BRIDGE_REGISTRY[id(component)] = self
189 _setup_eager_attention_hook_wrapper()
190 self._validate_submodule_declarations(component)
191 self._qk_norm_phase = self._decide_qk_norm_phase(component)
192 self._own_scaled_hook_k(component)
194 def _own_scaled_hook_k(self, hf_attn: torch.nn.Module) -> None:
195 """Replace the ``hook_k`` alias with a real HookPoint when K is scaled.
197 Aliased hook_k reported Falcon-H1's pre-key_multiplier tensor and
198 silently rescaled writes; the owned hook carries the scaled value while
199 ``k.hook_out`` stays raw (Granite's split). No-op elsewhere.
200 """
201 if getattr(hf_attn, "key_multiplier", None) is None:
202 return
203 if self.hook_aliases is type(self).hook_aliases: 203 ↛ 205line 203 didn't jump to line 205 because the condition on line 203 was always true
204 self.hook_aliases = dict(self.hook_aliases)
205 self.hook_aliases.pop("hook_k", None)
206 # The per-head hook_conversion is attached later, by
207 # _setup_qkv_hook_reshaping — component binding always precedes hook
208 # compatibility setup (bridge.py wires components, then calls it).
209 self.hook_k = HookPoint()
211 def _fire_scaled_hook_k(self, key_states: torch.Tensor) -> torch.Tensor:
212 """Fire an owned ``hook_k`` on the flat 3D tensor, preserving input rank.
214 A 4D input must flatten first: the conversion's revert only fires on 4D
215 returns, so an edited tensor would reach RoPE with the wrong rank.
216 """
217 if "hook_k" in self.hook_aliases or not hasattr(self, "hook_k"): 217 ↛ 218line 217 didn't jump to line 218 because the condition on line 217 was never true
218 return key_states
219 if key_states.dim() == 4: 219 ↛ 220line 219 didn't jump to line 220 because the condition on line 219 was never true
220 b, s, n_h, d_h = key_states.shape
221 flat = self.hook_k(key_states.reshape(b, s, n_h * d_h))
222 return flat.reshape(b, s, n_h, d_h)
223 return self.hook_k(key_states)
225 def _validate_submodule_declarations(self, hf_attn: torch.nn.Module) -> None:
226 """Raise if adapter omits q/k/v/o or a QK-norm the HF module has."""
227 # Silent fallback to raw HF linears is exactly what caused hook_q/k/v/z
228 # to never fire on 25 adapters; require explicit declaration.
229 missing = [req for req in ("q", "k", "v", "o") if req not in self.submodules]
230 if missing: 230 ↛ 231line 230 didn't jump to line 231 because the condition on line 230 was never true
231 raise RuntimeError(
232 f"{type(self).__name__} at '{self.name}' is missing required "
233 f"submodules: {missing}. Declare them in the adapter's "
234 f"component_mapping, e.g. submodules={{'q': LinearBridge(name='q_proj'), "
235 f"'k': LinearBridge(name='k_proj'), 'v': LinearBridge(name='v_proj'), "
236 f"'o': LinearBridge(name='o_proj')}}."
237 )
238 # Reverse mismatch (adapter declares, HF lacks) surfaces at norm forward.
239 # HF spells the per-head variants q_layernorm/k_layernorm (StableLM).
240 for declared, hf_names in (
241 ("q_norm", ("q_norm", "q_layernorm")),
242 ("k_norm", ("k_norm", "k_layernorm")),
243 ):
244 hf_present = next((n for n in hf_names if getattr(hf_attn, n, None) is not None), None)
245 if hf_present is not None and declared not in self.submodules: 245 ↛ 246line 245 didn't jump to line 246 because the condition on line 245 was never true
246 raise RuntimeError(
247 f"{type(self).__name__} at '{self.name}': HF module has "
248 f"'{hf_present}' but adapter did not declare '{declared}'. "
249 f"Forward would skip the norm, producing wrong logits vs HF. "
250 f"Add '{declared}' (name='{hf_present}') to the attention "
251 f"submodules."
252 )
254 def _decide_qk_norm_phase(self, hf_attn: torch.nn.Module) -> Optional[str]:
255 """Dispatch pre/post-reshape norm from weight shape; raise on ambiguity."""
256 if "q_norm" not in self.submodules:
257 return None
259 hf_norm_name = self.submodules["q_norm"].name
261 if hf_norm_name is None: 261 ↛ 262line 261 didn't jump to line 262 because the condition on line 261 was never true
262 raise RuntimeError(f"{self.name}: q_norm submodule declared without a name.")
264 q_norm = getattr(hf_attn, hf_norm_name, None)
266 if q_norm is None:
267 # Config-gated norms (use_qk_norm, qk_layernorm) declare optional;
268 # setup pops them after this runs.
269 if getattr(self.submodules["q_norm"], "optional", False):
270 return None
271 raise RuntimeError(f"{self.name}: q_norm declared but HF module has none.")
273 weight = getattr(q_norm, "weight", None)
274 head_dim = int(getattr(hf_attn, "head_dim"))
275 n_heads = int(getattr(self.config, "n_heads", 0))
277 # Non-learnable norm (Gemma-3 style) broadcasts over head_dim.
278 if weight is None or weight.ndim == 0:
279 return "post_reshape"
280 shape = tuple(weight.shape)
281 if shape == (head_dim,):
282 return "post_reshape"
283 if n_heads and shape == (n_heads * head_dim,): 283 ↛ 286line 283 didn't jump to line 286 because the condition on line 283 was always true
284 return "pre_reshape"
285 # Per-head norm (Cohere) broadcasts on the reshaped [B,H,S,D] tensor.
286 if n_heads and shape == (n_heads, head_dim):
287 return "post_reshape"
288 raise RuntimeError(
289 f"{self.name}: cannot determine QK-norm phase from q_norm weight "
290 f"shape {shape} (head_dim={head_dim}, n_heads={n_heads}). Expected "
291 f"(head_dim,), (n_heads*head_dim,), or (n_heads, head_dim)."
292 )
294 @staticmethod
295 def _apply_pre_reshape_qk_norm(
296 tensor: torch.Tensor,
297 norm_module: Any,
298 hook: Any,
299 head_dim: int,
300 ) -> torch.Tensor:
301 """Apply an OLMo-2-style pre-reshape QK norm, shape-preserving.
303 The norm computes RMS over the flattened (n_heads * d_head) dim. When
304 the split path hands us a 4D [B, S, H, d_head], flatten, norm, and
305 re-split so the result matches what the default 3D path produces at
306 this point.
307 """
308 if tensor.ndim == 4:
309 b, s, h, d = tensor.shape
310 flat = tensor.reshape(b, s, h * d)
311 normed = hook(norm_module(flat))
312 return normed.view(b, s, h, d)
313 return hook(norm_module(tensor))
315 def forward(self, *args: Any, **kwargs: Any) -> Any:
316 """Reimplemented forward pass with hooks at correct computation stages.
318 Instead of delegating to the HF attention module (which returns post-softmax
319 weights), this reimplements attention step-by-step so that:
320 - hook_attn_scores fires on PRE-softmax scores (matching HookedTransformer)
321 - hook_pattern fires on POST-softmax weights
322 - hook_rot_q/hook_rot_k fire after RoPE application
324 Handles RoPE, GQA, Q/K norms, sliding window, and softcapping.
325 """
326 if self.original_component is None: 326 ↛ 327line 326 didn't jump to line 327 because the condition on line 326 was never true
327 raise RuntimeError(
328 f"Original component not set for {self.name}. "
329 "Call set_original_component() first."
330 )
332 # Type as Any — the HF attention module's interface (q_proj, k_proj, etc.)
333 # varies by architecture and isn't captured by nn.Module's type signature.
334 hf_attn: Any = self.original_component
336 # Extract hidden_states and kwargs
337 if "hidden_states" in kwargs:
338 hidden_states = kwargs.pop("hidden_states")
339 elif len(args) > 0 and isinstance(args[0], torch.Tensor): 339 ↛ 343line 339 didn't jump to line 343 because the condition on line 339 was always true
340 hidden_states = args[0]
341 args = args[1:]
342 else:
343 raise ValueError("Could not find hidden_states in args or kwargs")
345 position_embeddings = kwargs.pop("position_embeddings", None)
346 attention_mask = kwargs.pop("attention_mask", None)
348 # Apply input hook
349 hidden_states = self.hook_in(hidden_states)
351 input_shape = hidden_states.shape[:-1]
352 head_dim = hf_attn.head_dim
353 hidden_shape = (*input_shape, -1, head_dim)
355 use_split_qkv = bool(getattr(self.config, "use_split_qkv_input", False))
356 use_attn_in = bool(getattr(self.config, "use_attn_in", False))
357 has_head_count = (
358 self.config is not None and hasattr(self.config, "n_heads") and self.config.n_heads
359 )
360 split_active = (use_split_qkv or use_attn_in) and has_head_count
362 # Qwen3.5/Qwen3-Next interleave [Q|gate] per head in q_proj output.
363 # The 2×-width output breaks per-head W slicing, so the split path is
364 # not supported for gated q_proj. Raise explicitly rather than
365 # producing silently wrong logits.
366 if split_active and getattr(self.config, "gated_q_proj", False):
367 raise NotImplementedError(
368 "use_split_qkv_input / use_attn_in are not supported on gated "
369 "q_proj architectures (Qwen3.5 / Qwen3-Next). The 2×-width "
370 "q_proj output breaks per-head weight routing. If you need "
371 "this combination, file a bug describing the workflow."
372 )
374 if split_active:
375 assert self.config is not None # narrowed by `has_head_count`
376 n_heads = int(self.config.n_heads)
377 n_kv_heads = int(getattr(self.config, "n_key_value_heads", None) or n_heads)
378 # #1317: fork pre-LN when available so hook patches match legacy.
379 captured = self._captured_pre_ln_residual
380 source = captured if captured is not None else hidden_states
381 if use_split_qkv:
382 q_in = self._fork_and_norm_per_head(source, self.hook_q_input, n_heads)
383 k_in = self._fork_and_norm_per_head(source, self.hook_k_input, n_kv_heads)
384 v_in = self._fork_and_norm_per_head(source, self.hook_v_input, n_kv_heads)
385 else:
386 attn_in = self._fork_and_norm_per_head(source, self.hook_attn_in, n_heads)
387 q_in = attn_in
388 if n_kv_heads != n_heads:
389 k_in = attn_in[..., :n_kv_heads, :].contiguous()
390 v_in = attn_in[..., :n_kv_heads, :].contiguous()
391 else:
392 k_in = v_in = attn_in
393 query_states = self._project_per_head_qkv(self.q, q_in, n_heads, head_dim)
394 key_states = self._project_per_head_qkv(self.k, k_in, n_kv_heads, head_dim)
395 value_states = self._project_per_head_qkv(self.v, v_in, n_kv_heads, head_dim)
396 q_gate = None
397 else:
398 # Route through LinearBridges so hook_q/k/v/z (aliased to
399 # q/k/v.hook_out, o.hook_in) fire on the live path.
400 query_states = self.q(hidden_states)
401 key_states = self.k(hidden_states)
402 value_states = self.v(hidden_states)
404 # Qwen3.5/Qwen3-Next interleave [Q|gate] per head in q_proj output.
405 # Processed-weights mode slices q_proj to standard width beforehand,
406 # so the 2×-width path only triggers on unprocessed state dicts.
407 q_gate = None
408 if getattr(self.config, "gated_q_proj", False):
409 q_dim = query_states.shape[-1]
410 n_heads_gated = getattr(self.config, "n_heads", q_dim // head_dim)
411 standard_q_dim = n_heads_gated * head_dim
412 if q_dim == standard_q_dim * 2:
413 query_states, q_gate = torch.chunk(
414 query_states.view(*input_shape, -1, head_dim * 2), 2, dim=-1
415 )
416 q_gate = q_gate.reshape(*input_shape, -1)
417 query_states = query_states.reshape(*input_shape, -1)
419 # Falcon-H1 scales K by a learned mup scalar between projection and RoPE.
420 # hook_k fires after the scale (see _own_scaled_hook_k) so it carries the
421 # tensor that actually reaches attention; k.hook_out kept the raw one.
422 key_multiplier = getattr(hf_attn, "key_multiplier", None)
423 if key_multiplier is not None:
424 key_states = key_states * key_multiplier
425 key_states = self._fire_scaled_hook_k(key_states)
427 has_q_norm = "q_norm" in self.submodules
428 has_k_norm = "k_norm" in self.submodules
430 # Pre-reshape phase (OLMo-2): norm is RMS over the flattened H*d_head
431 # dim. When the split path produced 4D [B, S, H, d_head], flatten for
432 # the norm then re-split so the post-norm tensors share shape with the
433 # non-split path going into the transpose below.
434 if has_q_norm and self._qk_norm_phase == "pre_reshape":
435 query_states = self._apply_pre_reshape_qk_norm(
436 query_states, self.q_norm, self.hook_q_normed, head_dim
437 )
438 if has_k_norm: 438 ↛ 446line 438 didn't jump to line 446 because the condition on line 438 was always true
439 key_states = self._apply_pre_reshape_qk_norm(
440 key_states, self.k_norm, self.hook_k_normed, head_dim
441 )
443 # OLMo v1 / OLMoE clamp Q/K/V when clip_qkv is set — after the
444 # pre-reshape qk-norm (OLMoE norms first) and before RoPE, matching HF
445 # order. HF gates on `is not None` for these archs.
446 clip_qkv = getattr(getattr(hf_attn, "config", None), "clip_qkv", None)
447 if clip_qkv is not None:
448 if self._qk_norm_phase == "post_reshape": 448 ↛ 451line 448 didn't jump to line 451 because the condition on line 448 was never true
449 # No arch pairs clip_qkv with a post-reshape norm; the HF
450 # ordering is unknowable here, so refuse rather than guess.
451 raise NotImplementedError(
452 "clip_qkv with a post-reshape qk-norm has no reference "
453 "ordering; add the architecture's HF order before enabling."
454 )
455 query_states, key_states, value_states = clamp_qkv(
456 query_states, key_states, value_states, clip_qkv
457 )
459 # For the split path, tensors are already [B, S, H, d_head]; for the
460 # default path they're flat [B, S, H*d_head] and need the view.
461 if split_active:
462 query_states = query_states.transpose(1, 2)
463 key_states = key_states.transpose(1, 2)
464 value_states = value_states.transpose(1, 2)
465 else:
466 query_states = query_states.view(hidden_shape).transpose(1, 2)
467 key_states = key_states.view(hidden_shape).transpose(1, 2)
468 value_states = value_states.view(hidden_shape).transpose(1, 2)
470 # Post-reshape phase (Gemma-3/Cohere): norm on [B, H, S, D].
471 if has_q_norm and self._qk_norm_phase == "post_reshape":
472 query_states = self.hook_q_normed(self.q_norm(query_states))
473 if has_k_norm: 473 ↛ 477line 473 didn't jump to line 477 because the condition on line 473 was always true
474 key_states = self.hook_k_normed(self.k_norm(key_states))
476 # --- RoPE ---
477 if ( 477 ↛ 486line 477 didn't jump to line 486 because the condition on line 477 was never true
478 position_embeddings is None
479 and not self.rope_optional
480 and getattr(self.config, "positional_embedding_type", None) == "rotary"
481 ):
482 # Silent skipping is how internlm2 ran without positional encoding
483 # while reporting perfect parity: its per-layer rotary means the HF
484 # decoder layer passes nothing down. Adapters in that shape must
485 # supply position_embeddings themselves (see internlm2.py).
486 import warnings
488 warnings.warn(
489 f"{type(self).__name__}({self.name}) reconstructed attention without "
490 "position_embeddings on a rotary architecture — RoPE was NOT applied. "
491 "The adapter must supply them (e.g. from the layer's own rotary_emb).",
492 RuntimeWarning,
493 stacklevel=2,
494 )
495 if position_embeddings is not None:
496 position_embeddings = self._apply_position_embedding_hooks(position_embeddings)
497 cos, sin = position_embeddings
498 if getattr(self.config, "rotary_adjacent_pairs", False):
499 # GLM/ERNIE convention: rotate adjacent element pairs in fp32,
500 # with cos/sin halves expanded by repeat_interleave.
501 apply_rotary_pos_emb = _apply_rotary_pos_emb_adjacent_pairs
502 else:
503 from transformers.models.llama.modeling_llama import (
504 apply_rotary_pos_emb,
505 )
507 rotary_dim = cos.shape[-1]
508 if rotary_dim * 2 == head_dim and get_rotary_pct_from_config(self.config) == 1.0:
509 # GPT-OSS convention: full rotation with un-duplicated half-width
510 # cos/sin. Duplicating the halves reduces it to llama's
511 # rotate-half formula over the full head_dim.
512 cos = torch.cat([cos, cos], dim=-1)
513 sin = torch.cat([sin, sin], dim=-1)
514 query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
515 elif rotary_dim < head_dim:
516 # Partial rotary (e.g., GPT-NeoX/Phi) where cos/sin cover only a
517 # portion of head_dim. Split Q/K, rotate the partial dims, recombine.
518 q_rot, q_pass = query_states[..., :rotary_dim], query_states[..., rotary_dim:]
519 k_rot, k_pass = key_states[..., :rotary_dim], key_states[..., rotary_dim:]
520 q_rot, k_rot = apply_rotary_pos_emb(q_rot, k_rot, cos, sin)
521 query_states = torch.cat([q_rot, q_pass], dim=-1)
522 key_states = torch.cat([k_rot, k_pass], dim=-1)
523 else:
524 query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
526 # Ministral-3's llama-4 query scale: HF multiplies Q by
527 # 1 + beta*log(1 + floor(pos/original_max)) right after RoPE.
528 rope_params = getattr(getattr(hf_attn, "config", None), "rope_parameters", None)
529 if isinstance(rope_params, dict) and rope_params.get("llama_4_scaling_beta") is not None:
530 from transformers.models.ministral3.modeling_ministral3 import (
531 get_llama_4_attn_scale,
532 )
534 position_ids = kwargs.get("position_ids")
535 if position_ids is None: 535 ↛ 536line 535 didn't jump to line 536 because the condition on line 535 was never true
536 position_ids = torch.arange(
537 query_states.shape[-2], device=query_states.device
538 ).unsqueeze(0)
539 query_states = query_states * get_llama_4_attn_scale(
540 position_ids,
541 rope_params["llama_4_scaling_beta"],
542 int(rope_params["original_max_position_embeddings"]),
543 ).to(query_states.dtype)
545 # Fire hook_rot_q/hook_rot_k (post-rotation)
546 if hasattr(self, "hook_rot_q"):
547 query_states = self.hook_rot_q(query_states)
548 if hasattr(self, "hook_rot_k"):
549 key_states = self.hook_rot_k(key_states)
551 # --- KV cache: extend K/V with cached positions ---
552 key_states, value_states = self._update_kv_cache(key_states, value_states, **kwargs)
554 # --- GQA: Expand K/V ---
555 num_key_value_groups = getattr(hf_attn, "num_key_value_groups", 1)
556 if num_key_value_groups > 1:
557 from transformers.models.llama.modeling_llama import repeat_kv
559 key_states_expanded = repeat_kv(key_states, num_key_value_groups)
560 value_states_expanded = repeat_kv(value_states, num_key_value_groups)
561 else:
562 key_states_expanded = key_states
563 value_states_expanded = value_states
565 # --- Attention Scores ---
566 scaling = getattr(hf_attn, "scaling", head_dim**-0.5)
567 attn_scores = torch.matmul(query_states, key_states_expanded.transpose(-2, -1)) * scaling
569 # --- Softcapping (Gemma 2) ---
570 softcap = getattr(hf_attn, "attn_logit_softcapping", None)
571 if softcap is not None:
572 attn_scores = attn_scores / softcap
573 attn_scores = torch.tanh(attn_scores)
574 attn_scores = attn_scores * softcap
576 # --- Causal / Sliding Window Mask ---
577 kv_seq_len = key_states_expanded.shape[-2]
578 q_seq_len = query_states.shape[-2]
579 attn_scores = self._apply_reconstruct_attention_mask(
580 attn_scores=attn_scores,
581 attention_mask=attention_mask,
582 seq_len=kv_seq_len,
583 q_seq_len=q_seq_len,
584 )
586 # --- hook_attn_scores: PRE-softmax (matching HookedTransformer) ---
587 attn_scores = self.hook_attn_scores(attn_scores)
589 # --- Softmax (in float32 for numerical stability) ---
590 # GPT-OSS attention sinks: a learned per-head logit joins the softmax as
591 # an extra key column and is dropped afterward, so every real position's
592 # weight is scaled down by the sink's share. Appended after
593 # hook_attn_scores so hooks keep the [batch, head, q_pos, kv_pos] shape
594 # and score patches still flow into the softmax.
595 sinks = getattr(hf_attn, "sinks", None)
596 if sinks is not None:
597 sink_col = (
598 sinks.reshape(1, -1, 1, 1)
599 .expand(attn_scores.shape[0], -1, attn_scores.shape[-2], -1)
600 .to(attn_scores.dtype)
601 )
602 combined = torch.cat([attn_scores, sink_col], dim=-1)
603 combined = combined - combined.max(dim=-1, keepdim=True).values
604 attn_weights = torch.nn.functional.softmax(combined, dim=-1, dtype=torch.float32).to(
605 query_states.dtype
606 )[..., :-1]
607 else:
608 attn_weights = torch.nn.functional.softmax(attn_scores, dim=-1, dtype=torch.float32).to(
609 query_states.dtype
610 )
611 attn_weights = self._scrub_compatibility_pattern_nans(attn_weights)
613 # --- Dropout ---
614 dropout_rate = getattr(hf_attn, "attention_dropout", 0.0)
615 if self.training and dropout_rate > 0.0: 615 ↛ 616line 615 didn't jump to line 616 because the condition on line 615 was never true
616 attn_weights = torch.nn.functional.dropout(attn_weights, p=dropout_rate, training=True)
618 # --- hook_pattern: POST-softmax ---
619 attn_weights = self.hook_pattern(attn_weights)
621 # --- Attention Output ---
622 attn_output = torch.matmul(attn_weights, value_states_expanded)
623 attn_output = attn_output.transpose(1, 2).contiguous()
624 attn_output = attn_output.reshape(*input_shape, -1)
626 # --- Gated attention (Qwen3.5/Qwen3Next) ---
627 if q_gate is not None:
628 if hasattr(self, "hook_q_gate"): 628 ↛ 630line 628 didn't jump to line 630 because the condition on line 628 was always true
629 q_gate = self.hook_q_gate(q_gate)
630 attn_output = attn_output * torch.sigmoid(q_gate)
632 # --- Gated attention (HRM-Text: separate gate_proj on hidden_states) ---
633 gate_comp = self._modules.get("gate")
634 if gate_comp is not None and gate_comp.original_component is not None and q_gate is None:
635 gate_states = gate_comp(hidden_states)
636 if hasattr(self, "hook_gate"): 636 ↛ 638line 636 didn't jump to line 638 because the condition on line 636 was always true
637 gate_states = self.hook_gate(gate_states)
638 attn_output = attn_output * torch.sigmoid(gate_states)
640 # Adapter seam: sub-layer transforms between attention output and the
641 # o projection (e.g. BitNet's attn_sub_norm).
642 attn_output = self._pre_output_projection(attn_output)
644 # Some rotary modules (FlexOlmo) return fp32 cos/sin without casting to
645 # the input dtype, so RoPE promotes q/k
646 # to fp32 while the projection weights stay in the model dtype. Match
647 # the projection rather than the activations; no-op when they agree.
648 o_module = getattr(self, "o", None)
649 o_weight = getattr(getattr(o_module, "original_component", None), "weight", None)
650 if ( 650 ↛ 659line 650 didn't jump to line 659 because the condition on line 650 was never true
651 isinstance(o_weight, torch.Tensor)
652 # Quantized weights (bnb int8/uint8, GPTQ int32) dequantize inside
653 # the matmul; casting activations to an integer storage dtype would
654 # destroy them (same guard as base.py's compute-dtype selection).
655 and o_weight.dtype.is_floating_point
656 and attn_output.is_floating_point()
657 and attn_output.dtype != o_weight.dtype
658 ):
659 attn_output = attn_output.to(dtype=o_weight.dtype)
661 if (
662 bool(getattr(self.config, "use_attn_result", False))
663 and hasattr(self, "o")
664 and self.o.original_component is not None
665 ):
666 # Per-head output pre-sum across heads. Fire hook_z on the pre-
667 # projection tensor first so any patch at hook_z flows into the
668 # per-head computation below — matches the default path where
669 # `self.o(attn_output)` calls o.hook_in before the linear.
670 n_heads = int(getattr(self.config, "n_heads"))
671 attn_output = self.o.hook_in(attn_output)
672 z_4d = attn_output.view(*input_shape, n_heads, head_dim)
673 attn_output = self._compute_per_head_result(z_4d, n_heads, head_dim)
674 attn_output = self.hook_out(attn_output)
675 else:
676 # Route through LinearBridge so hook_z (aliased to o.hook_in) fires.
677 # LinearBridge wraps whichever HF attr the adapter mapped (o_proj,
678 # dense, out_proj).
679 attn_output = self.o(attn_output)
680 attn_output = self.hook_out(attn_output)
682 return attn_output, attn_weights
684 def _pre_output_projection(self, attn_output: torch.Tensor) -> torch.Tensor:
685 """Overridable seam applied before the output projection."""
686 return attn_output
688 def get_random_inputs(
689 self,
690 batch_size: int = 2,
691 seq_len: int = 8,
692 device: Optional[torch.device] = None,
693 dtype: Optional[torch.dtype] = None,
694 ) -> Dict[str, Any]:
695 """Generate random inputs for Gemma-3 attention testing.
697 Gemma-3's position_embeddings are generated by calling rotary_emb(seq_len, device)
698 which returns a tuple of (cos, sin) tensors with shape [seq_len, head_dim].
700 Args:
701 batch_size: Batch size for generated inputs
702 seq_len: Sequence length for generated inputs
703 device: Device to place tensors on
704 dtype: Dtype for generated tensors
706 Returns:
707 Dictionary with keys: hidden_states, position_embeddings, attention_mask
708 """
709 if device is None:
710 device = torch.device("cpu")
711 if dtype is None:
712 dtype = torch.float32
713 d_model = self.config.d_model if self.config and hasattr(self.config, "d_model") else 1152
714 inputs: Dict[str, Any] = {
715 "hidden_states": torch.randn(batch_size, seq_len, d_model, device=device, dtype=dtype)
716 }
717 num_heads = safe_config_get(self.config, "num_attention_heads", 4) if self.config else 4
718 head_dim = safe_config_get(self.config, "head_dim", 256) if self.config else 256
719 dummy_qk = torch.randn(1, seq_len, num_heads, head_dim, device=device, dtype=dtype)
720 position_ids = torch.arange(seq_len, device=device).unsqueeze(0)
721 if self._rotary_emb is not None:
722 try:
723 position_embeddings = self._rotary_emb(dummy_qk, position_ids)
724 inputs["position_embeddings"] = position_embeddings
725 except Exception as e:
726 cos = torch.ones(1, seq_len, head_dim, device=device, dtype=dtype)
727 sin = torch.zeros(1, seq_len, head_dim, device=device, dtype=dtype)
728 inputs["position_embeddings"] = (cos, sin)
729 else:
730 cos = torch.ones(1, seq_len, head_dim, device=device, dtype=dtype)
731 sin = torch.zeros(1, seq_len, head_dim, device=device, dtype=dtype)
732 inputs["position_embeddings"] = (cos, sin)
733 inputs["attention_mask"] = None
734 return inputs