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