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