Coverage for transformer_lens/model_bridge/generalized_components/attention.py: 90%

1"""Attention bridge component. 

2 

3This module contains the bridge component for attention layers. 

4""" 

5import logging 

6from typing import Any, Dict, Optional, Tuple 

7 

8import einops 

9import torch 

10from transformers.pytorch_utils import Conv1D 

11 

12logger = logging.getLogger(__name__) 

13 

14from transformer_lens.conversion_utils.conversion_steps.attention_auto_conversion import ( 

15 AttentionAutoConversion, 

16) 

17from transformer_lens.conversion_utils.conversion_steps.base_tensor_conversion import ( 

18 BaseTensorConversion, 

19) 

20from transformer_lens.hook_points import HookPoint 

21from transformer_lens.model_bridge.generalized_components.base import ( 

22 GeneralizedComponent, 

23) 

24from transformer_lens.utilities.hf_utils import get_rotary_pct_from_config 

25from transformer_lens.utilities.quantization import require_readable_weight 

26 

27 

28class PerLayerGeometryError(ValueError): 

29 """A weight accessor refused a head split on per-layer attention geometry. 

30 

31 Dedicated type so callers with a raw-weight fallback (the centering 

32 benchmark) catch exactly this and nothing else — a generic ValueError 

33 catch would silently mask unrelated accessor regressions. 

34 """ 

35 

36 

37class AttentionBridge(GeneralizedComponent): 

38 """Bridge component for attention layers. 

39 

40 This component handles the conversion between Hugging Face attention layers 

41 and TransformerLens attention components. 

42 """ 

43 

44 hook_aliases = { 

45 "hook_q": "q.hook_out", 

46 "hook_k": "k.hook_out", 

47 "hook_v": "v.hook_out", 

48 "hook_z": "o.hook_in", 

49 } 

50 

51 # Override to False on variants without a pre-LN fork (e.g. MLA); skips 

52 # the split-qkv HookPoints and the BlockBridge pre-ln1 capture. 

53 supports_split_qkv_fork: bool = True 

54 # Reconstructed variants can opt in independently of Q/K/V input forks. 

55 supports_attn_result: bool = False 

56 property_aliases = { 

57 "W_Q": "q.weight", 

58 "W_K": "k.weight", 

59 "W_V": "v.weight", 

60 "W_O": "o.weight", 

61 "b_Q": "q.bias", 

62 "b_K": "k.bias", 

63 "b_V": "v.bias", 

64 "b_O": "o.bias", 

65 } 

66 

67 def __init__( 

68 self, 

69 name: Optional[str], 

70 config: Any, 

71 submodules: Optional[Dict[str, GeneralizedComponent]] = None, 

72 conversion_rule: Optional[BaseTensorConversion] = None, 

73 pattern_conversion_rule: Optional[BaseTensorConversion] = None, 

74 maintain_native_attention: bool = False, 

75 requires_position_embeddings: bool = False, 

76 requires_attention_mask: bool = False, 

77 attention_mask_4d: bool = False, 

78 requires_relative_position_bias: bool = False, 

79 is_cross_attention: bool = False, 

80 is_causal: bool = True, 

81 optional: bool = False, 

82 fused_qkv: bool = False, 

83 ): 

84 """Initialize the attention bridge. 

85 

86 Args: 

87 name: The name of this component, or None when projections live on the parent 

88 config: Model configuration (required for auto-conversion detection) 

89 submodules: Dictionary of submodules to register (e.g., q_proj, k_proj, etc.) 

90 conversion_rule: Optional conversion rule. If None, AttentionAutoConversion will be used 

91 pattern_conversion_rule: Optional conversion rule for attention patterns. If None, 

92 uses AttentionPatternConversion to ensure [n_heads, pos, pos] shape 

93 maintain_native_attention: If True, preserve the original HF attention implementation 

94 without wrapping. Use for models with custom attention 

95 (e.g., attention sinks, specialized RoPE). Defaults to False. 

96 requires_position_embeddings: If True, this attention requires position_embeddings argument 

97 (e.g., Gemma-3 with dual RoPE). Defaults to False. 

98 requires_attention_mask: If True, this attention requires attention_mask argument 

99 (e.g., GPTNeoX/Pythia). Defaults to False. 

100 attention_mask_4d: If True, generate 4D attention_mask [batch, 1, tgt_len, src_len] 

101 instead of 2D [batch, seq_len]. Required for OPT. Defaults to False. 

102 requires_relative_position_bias: T5/mT5-style relative attention; supplies a 

103 zero ``position_bias`` so HF's forward skips its ``cache_position[-1]`` fallback. 

104 is_cross_attention: Encoder-decoder cross-attention; supplies ``key_value_states``. 

105 is_causal: If True, apply a causal (lower-triangular) mask when reconstructing 

106 attention. Set False for bidirectional encoders (e.g. T5Gemma's encoder). 

107 """ 

108 if conversion_rule is None: 108 ↛ 110line 108 didn't jump to line 110 because the condition on line 108 was always true

109 conversion_rule = AttentionAutoConversion(config) 

110 super().__init__( 

111 name, 

112 config=config, 

113 submodules=submodules or {}, 

114 conversion_rule=conversion_rule, 

115 optional=optional, 

116 ) 

117 if fused_qkv: 

118 # A combined QKV projection has no q/k/v submodules for the class 

119 # aliases to resolve against; leaving them in place produces dead 

120 # aliases and resolution-audit noise (raven Wqkv, OpenELM qkv_proj). 

121 self.hook_aliases = { 

122 key: value 

123 for key, value in type(self).hook_aliases.items() 

124 if key not in {"hook_q", "hook_k", "hook_v"} 

125 } 

126 self.hook_attn_scores = HookPoint() 

127 self.hook_pattern = HookPoint() 

128 self.hook_hidden_states = HookPoint() 

129 # Per-head attention output, pre-sum across heads. 

130 # Shape [batch, pos, n_heads, d_model] when fired. Gated at fire time 

131 # by cfg.use_attn_result; the HookPoint exists unconditionally so 

132 # run_with_cache key lookups never miss. 

133 self.hook_result = HookPoint() 

134 # Pre-ln1 fork hooks ([B, S, H, D]) gated by use_split_qkv_input / 

135 # use_attn_in; fall back to post-ln1 if BlockBridge can't wire ln1. See #1317. 

136 if self.supports_split_qkv_fork: 

137 self.hook_attn_in = HookPoint() 

138 self.hook_q_input = HookPoint() 

139 self.hook_k_input = HookPoint() 

140 self.hook_v_input = HookPoint() 

141 self._captured_pre_ln_residual: Optional[torch.Tensor] = None 

142 self._ln1_module: Optional[torch.nn.Module] = None 

143 if ( 

144 hasattr(config, "positional_embedding_type") 

145 and config.positional_embedding_type == "rotary" 

146 ): 

147 self.hook_rot_k = HookPoint() 

148 self.hook_rot_q = HookPoint() 

149 self.hook_hidden_states.hook_conversion = conversion_rule 

150 if pattern_conversion_rule is not None: 150 ↛ 151line 150 didn't jump to line 151 because the condition on line 150 was never true

151 self.hook_pattern.hook_conversion = pattern_conversion_rule 

152 self._attn_scores = None 

153 self._pattern = None 

154 self._hf_forward_wrapped = False 

155 self.maintain_native_attention = maintain_native_attention 

156 self.requires_position_embeddings = requires_position_embeddings 

157 self.requires_attention_mask = requires_attention_mask 

158 self.attention_mask_4d = attention_mask_4d 

159 self.requires_relative_position_bias = requires_relative_position_bias 

160 self.is_cross_attention = is_cross_attention 

161 self.is_causal = is_causal 

162 self._layer_idx: Optional[int] = None 

163 

164 def set_original_component(self, original_component: torch.nn.Module) -> None: 

165 """Set original component and capture layer index for KV caching.""" 

166 super().set_original_component(original_component) 

167 layer_idx_raw = getattr(original_component, "layer_idx", None) 

168 if layer_idx_raw is not None: 

169 self._layer_idx = int(layer_idx_raw) 

170 

171 def _apply_ln1_per_head(self, x: torch.Tensor) -> torch.Tensor: 

172 """Apply ln1 to [B, S, H, D] with H folded into the batch. Identity if ln1 unwired. 

173 

174 Routes through the raw HF norm to avoid refiring ln1's internal hooks 

175 per-head — deliberate divergence from legacy's *Pre sub-hook firing. 

176 """ 

177 if self._ln1_module is None: 177 ↛ 178line 177 didn't jump to line 178 because the condition on line 177 was never true

178 return x 

179 b, s, h, d = x.shape 

180 return self._ln1_module(x.reshape(b * s * h, d)).reshape(b, s, h, d) 

181 

182 def _fork_and_norm_per_head( 

183 self, source: torch.Tensor, hook: HookPoint, n_heads: int 

184 ) -> torch.Tensor: 

185 """Repeat residual to [B, S, H, D], fire ``hook``, re-LN iff source is pre-LN.""" 

186 forked = einops.repeat(source, "b s d -> b s h d", h=n_heads).contiguous() 

187 forked = hook(forked) 

188 if self._captured_pre_ln_residual is not None: 

189 forked = self._apply_ln1_per_head(forked) 

190 return forked 

191 

192 def setup_hook_compatibility(self) -> None: 

193 """Setup hook compatibility transformations to match HookedTransformer behavior. 

194 

195 This sets up hook conversions that ensure Bridge hooks have the same shapes 

196 as HookedTransformer hooks. This includes reshaping Q/K/V/Z hooks from 

197 [batch, seq, d_model] to [batch, seq, n_heads, d_head] format. 

198 

199 This is called during Bridge.__init__ and should always be run. 

200 Note: This method is idempotent - can be called multiple times safely. 

201 """ 

202 if self._hf_forward_wrapped: 

203 return 

204 if hasattr(self.config, "n_heads"): 204 ↛ 206line 204 didn't jump to line 206 because the condition on line 204 was always true

205 self._setup_qkv_hook_reshaping() 

206 self._hf_forward_wrapped = True 

207 

208 def get_random_inputs( 

209 self, 

210 batch_size: int = 2, 

211 seq_len: int = 8, 

212 device: Optional[torch.device] = None, 

213 dtype: Optional[torch.dtype] = None, 

214 ) -> Dict[str, Any]: 

215 """Get random inputs for testing this attention component. 

216 

217 Generates appropriate inputs based on the attention's requirements 

218 (position_embeddings, attention_mask, etc.). 

219 

220 Args: 

221 batch_size: Batch size for the test inputs 

222 seq_len: Sequence length for the test inputs 

223 device: Device to create tensors on (defaults to CPU) 

224 dtype: Dtype for generated tensors (defaults to float32) 

225 

226 Returns: 

227 Dictionary of keyword arguments to pass to forward() 

228 """ 

229 if device is None: 

230 device = torch.device("cpu") 

231 if dtype is None: 

232 dtype = torch.float32 

233 d_model = self.config.d_model if self.config and hasattr(self.config, "d_model") else 768 

234 inputs: Dict[str, Any] = { 

235 "hidden_states": torch.randn(batch_size, seq_len, d_model, device=device, dtype=dtype) 

236 } 

237 if self.requires_position_embeddings: 

238 if self.config: 238 ↛ 246line 238 didn't jump to line 246 because the condition on line 238 was always true

239 if hasattr(self.config, "d_head"): 239 ↛ 241line 239 didn't jump to line 241 because the condition on line 239 was always true

240 d_head = self.config.d_head 

241 elif hasattr(self.config, "head_dim"): 

242 d_head = self.config.head_dim 

243 else: 

244 d_head = 64 

245 else: 

246 d_head = 64 

247 rotary_pct = get_rotary_pct_from_config(self.config) 

248 rotary_ndims = int(rotary_pct * d_head) 

249 cos = torch.ones(1, seq_len, rotary_ndims, device=device, dtype=dtype) 

250 sin = torch.zeros(1, seq_len, rotary_ndims, device=device, dtype=dtype) 

251 inputs["position_embeddings"] = (cos, sin) 

252 # For models with internal rotary embeddings (e.g., GPT-J), the HF attention 

253 # forward expects position_ids to index into embed_positions. Models using 

254 # requires_position_embeddings get (cos, sin) tuples instead. 

255 if ( 255 ↛ 261line 255 didn't jump to line 261 because the condition on line 255 was never true

256 self.config 

257 and hasattr(self.config, "positional_embedding_type") 

258 and self.config.positional_embedding_type == "rotary" 

259 and not self.requires_position_embeddings 

260 ): 

261 inputs["position_ids"] = ( 

262 torch.arange(seq_len, device=device).unsqueeze(0).expand(batch_size, -1) 

263 ) 

264 if self.requires_attention_mask: 264 ↛ 265line 264 didn't jump to line 265 because the condition on line 264 was never true

265 if self.attention_mask_4d: 

266 # Generate 4D attention mask [batch, 1, tgt_len, src_len] for models like OPT 

267 inputs["attention_mask"] = torch.ones( 

268 batch_size, 1, seq_len, seq_len, device=device 

269 ) 

270 else: 

271 # Generate 2D attention mask [batch, seq_len] for most models 

272 inputs["attention_mask"] = torch.ones(batch_size, seq_len, device=device) 

273 if self.requires_relative_position_bias: 273 ↛ 275line 273 didn't jump to line 275 because the condition on line 273 was never true

274 # Zero bias short-circuits HF's None-cache_position fallback in T5Attention. 

275 n_heads = self.config.n_heads if self.config and hasattr(self.config, "n_heads") else 1 

276 inputs["position_bias"] = torch.zeros( 

277 1, n_heads, seq_len, seq_len, device=device, dtype=dtype 

278 ) 

279 if self.is_cross_attention: 

280 inputs["key_value_states"] = torch.randn( 

281 batch_size, seq_len, d_model, device=device, dtype=dtype 

282 ) 

283 return inputs 

284 

285 def _setup_qkv_hook_reshaping(self) -> None: 

286 """Setup hook reshaping for Q/K/V/Z to match HookedTransformer shapes. 

287 

288 Reshapes hooks from [batch, seq, d_model] to [batch, seq, n_heads, d_head] format. 

289 For models with Grouped Query Attention (GQA), K and V use n_kv_heads instead of n_heads. 

290 

291 Sets up conversions for: 

292 - q.hook_out (aliased as hook_q) 

293 - k.hook_out (aliased as hook_k) - uses n_kv_heads if GQA 

294 - v.hook_out (aliased as hook_v) - uses n_kv_heads if GQA 

295 - o.hook_in (aliased as hook_z) 

296 """ 

297 

298 class ReshapeForAttentionHeads(BaseTensorConversion): 

299 """Reshape tensors to split attention heads for Q/K/V/Z compatibility.""" 

300 

301 # Marks the [batch, pos, head, d_head] layout so consumers (e.g. run_with_cache's 

302 # pos_slice) know the position axis is two from the end, not one. 

303 splits_attention_heads = True 

304 

305 def __init__(self, n_heads: int, d_head: int): 

306 super().__init__() 

307 self.n_heads = n_heads 

308 self.d_head = d_head 

309 

310 def handle_conversion(self, input_value, *full_context): 

311 """Convert from [batch, seq, d_model] to [batch, seq, n_heads, d_head].""" 

312 if len(input_value.shape) == 3: 312 ↛ 316line 312 didn't jump to line 316 because the condition on line 312 was always true

313 b, s, d = input_value.shape 

314 if d == self.n_heads * self.d_head: 

315 return input_value.view(b, s, self.n_heads, self.d_head) 

316 return input_value 

317 

318 def revert(self, input_value, *full_context): 

319 """Revert from [batch, seq, n_heads, d_head] to [batch, seq, d_model].""" 

320 if len(input_value.shape) == 4: 320 ↛ 325line 320 didn't jump to line 325 because the condition on line 320 was always true

321 b, s, n_h, d_h = input_value.shape 

322 if n_h == self.n_heads and d_h == self.d_head: 322 ↛ 325line 322 didn't jump to line 325 because the condition on line 322 was always true

323 # reshape (not view) — callers may pass non-contiguous tensors 

324 return input_value.reshape(b, s, n_h * d_h) 

325 return input_value 

326 

327 if self.config is None: 327 ↛ 328line 327 didn't jump to line 328 because the condition on line 327 was never true

328 raise RuntimeError(f"Config not set for {self.name}") 

329 

330 # Get n_heads (try n_heads first, then n_head) 

331 if hasattr(self.config, "n_heads"): 331 ↛ 333line 331 didn't jump to line 333 because the condition on line 331 was always true

332 n_heads = self.config.n_heads 

333 elif hasattr(self.config, "n_head"): 

334 n_heads = self.config.n_head 

335 else: 

336 # Can't setup reshaping without knowing number of heads 

337 return 

338 

339 # Get d_head (try d_head first, then compute from d_model or n_embd) 

340 if hasattr(self.config, "d_head"): 

341 d_head = self.config.d_head 

342 elif hasattr(self.config, "d_model"): 342 ↛ 343line 342 didn't jump to line 343 because the condition on line 342 was never true

343 d_head = self.config.d_model // n_heads 

344 elif hasattr(self.config, "n_embd"): 344 ↛ 345line 344 didn't jump to line 345 because the condition on line 344 was never true

345 d_head = self.config.n_embd // n_heads 

346 else: 

347 # Can't setup reshaping without knowing head dimension 

348 return 

349 n_kv_heads = n_heads 

350 if hasattr(self.config, "n_key_value_heads") and self.config.n_key_value_heads is not None: 

351 n_kv_heads = self.config.n_key_value_heads 

352 

353 # Per-layer d_head makes width division ambiguous (gemma4: 2048 = 8x256 

354 # = 4x512), so dividing by the majority scalar mis-factorizes silently. 

355 layer_d_head: Optional[int] = d_head 

356 per_layer_hd = getattr(self.config, "per_layer_head_dim", None) 

357 if per_layer_hd: 

358 entry: Any = None 

359 if self._layer_idx is not None and self._layer_idx < len(per_layer_hd): 

360 entry = per_layer_hd[self._layer_idx] 

361 if not (isinstance(entry, int) and entry > 0): 

362 entry = getattr(self.original_component, "head_dim", None) 

363 layer_d_head = entry if isinstance(entry, int) and entry > 0 else None 

364 

365 def conversion_dims( 

366 proj: Any, cfg_heads: int, width_attr: str = "out_features" 

367 ) -> Tuple[int, int]: 

368 """(n_heads, d_head) from the BOUND module width, not the cfg scalars. 

369 

370 OpenELM/laguna vary head counts per layer; cfg-only dims no-op there, 

371 flipping hook shapes mid-model. MLA excluded — its o-width is v_head_dim. 

372 """ 

373 if getattr(self, "_v_head_dim", None): 373 ↛ 374line 373 didn't jump to line 374 because the condition on line 373 was never true

374 return cfg_heads, d_head 

375 component = getattr(proj, "original_component", None) 

376 width = getattr(component, width_attr, None) 

377 if width and layer_d_head and width % layer_d_head == 0: 

378 return width // layer_d_head, layer_d_head 

379 return cfg_heads, layer_d_head or d_head 

380 

381 if hasattr(self, "q") and self.q is not None and hasattr(self.q, "hook_out"): 

382 q_reshape = ReshapeForAttentionHeads(*conversion_dims(self.q, n_heads)) 

383 self.q.hook_out.hook_conversion = q_reshape 

384 if hasattr(self, "k") and self.k is not None and hasattr(self.k, "hook_out"): 

385 k_reshape = ReshapeForAttentionHeads(*conversion_dims(self.k, n_kv_heads)) 

386 self.k.hook_out.hook_conversion = k_reshape 

387 # Subclasses that de-alias hook_k onto their own HookPoint (K-scaling 

388 # architectures) need the same conversion, whichever order runs first. 

389 if "hook_k" not in self.hook_aliases and isinstance( 

390 getattr(self, "hook_k", None), HookPoint 

391 ): 

392 self.hook_k.hook_conversion = k_reshape 

393 if hasattr(self, "v") and self.v is not None and hasattr(self.v, "hook_out"): 

394 v_reshape = ReshapeForAttentionHeads(*conversion_dims(self.v, n_kv_heads)) 

395 self.v.hook_out.hook_conversion = v_reshape 

396 if hasattr(self, "o") and self.o is not None and hasattr(self.o, "hook_in"): 

397 z_reshape = ReshapeForAttentionHeads( 

398 *conversion_dims(self.o, n_heads, width_attr="in_features") 

399 ) 

400 self.o.hook_in.hook_conversion = z_reshape 

401 

402 class TransposeRotaryHeads(BaseTensorConversion): 

403 """Transpose rotary hook tensors from HF format to HookedTransformer format.""" 

404 

405 def handle_conversion(self, input_value, *full_context): 

406 """Convert from [batch, n_heads, seq, d_head] to [batch, seq, n_heads, d_head].""" 

407 if len(input_value.shape) == 4: 407 ↛ 409line 407 didn't jump to line 409 because the condition on line 407 was always true

408 return input_value.transpose(1, 2) 

409 return input_value 

410 

411 def revert(self, input_value, *full_context): 

412 """Revert from [batch, seq, n_heads, d_head] to [batch, n_heads, seq, d_head].""" 

413 if len(input_value.shape) == 4: 413 ↛ 415line 413 didn't jump to line 415 because the condition on line 413 was always true

414 return input_value.transpose(1, 2) 

415 return input_value 

416 

417 if hasattr(self, "hook_rot_q"): 

418 self.hook_rot_q.hook_conversion = TransposeRotaryHeads() 

419 if hasattr(self, "hook_rot_k"): 

420 self.hook_rot_k.hook_conversion = TransposeRotaryHeads() 

421 

422 def _update_kv_cache( 

423 self, k: torch.Tensor, v: torch.Tensor, **kwargs: Any 

424 ) -> tuple[torch.Tensor, torch.Tensor]: 

425 """Update KV cache if provided, returning the (possibly extended) K and V. 

426 

427 Call this after K/V projections and any positional embeddings (e.g. RoPE) 

428 have been applied, but before computing attention scores. If no cache is 

429 present in kwargs, K and V are returned unchanged. 

430 """ 

431 past_key_values = kwargs.get("past_key_values", None) 

432 if past_key_values is None: 

433 # GPT-NeoX/GPT-J/Bloom/Falcon/MPT/CodeGen/GPTBigCode still name the 

434 # cache `layer_past`; missing it leaves the cache empty, so every 

435 # decode step attends to itself alone and generation ignores the prompt. 

436 past_key_values = kwargs.get("layer_past", None) 

437 if past_key_values is None: 

438 return k, v 

439 layer_idx = getattr(self, "_layer_idx", None) 

440 if layer_idx is None: 

441 logger.warning( 

442 "%s: past_key_values provided but _layer_idx is None " 

443 "(HF component missing layer_idx attribute). KV cache update " 

444 "skipped — cached generation will ignore earlier tokens.", 

445 self.name, 

446 ) 

447 return k, v 

448 # Some architectures (e.g. HRM-Text's recurrent stacks) pass a cycle_offset 

449 # so each stack invocation writes to a unique cache slot. Non-participating 

450 # models simply leave it unset. 

451 cycle_offset = kwargs.get("cycle_offset", 0) 

452 k, v = past_key_values.update(k, v, layer_idx + cycle_offset) 

453 return k, v 

454 

455 def _reshape_qkv_to_heads( 

456 self, 

457 q: torch.Tensor, 

458 k: torch.Tensor, 

459 v: torch.Tensor, 

460 num_heads: int, 

461 num_kv_heads: int | None = None, 

462 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int, int, int]: 

463 """Reshape Q/K/V from [batch, seq, hidden] or [batch, seq, heads, head_dim] 

464 to [batch, heads, seq, head_dim]. Returns (q, k, v, batch_size, seq_len, head_dim). 

465 

466 Args: 

467 num_kv_heads: If provided and differs from num_heads (GQA), K/V use 

468 this head count for the 3D reshape path. 

469 """ 

470 if num_kv_heads is None: 

471 num_kv_heads = num_heads 

472 if q.ndim == 3: 

473 batch_size, seq_len, q_hidden = q.shape 

474 head_dim: int = q_hidden // num_heads 

475 q = q.view(batch_size, seq_len, num_heads, head_dim).transpose(1, 2) 

476 k = k.view(batch_size, seq_len, num_kv_heads, head_dim).transpose(1, 2) 

477 v = v.view(batch_size, seq_len, num_kv_heads, head_dim).transpose(1, 2) 

478 elif q.ndim == 4: 478 ↛ 485line 478 didn't jump to line 485 because the condition on line 478 was always true

479 batch_size, seq_len = q.shape[0], q.shape[1] 

480 head_dim = q.shape[-1] 

481 q = q.transpose(1, 2) 

482 k = k.transpose(1, 2) 

483 v = v.transpose(1, 2) 

484 else: 

485 raise ValueError(f"Unexpected Q tensor shape: {q.shape}. Expected 3D or 4D.") 

486 return q, k, v, batch_size, seq_len, head_dim 

487 

488 def _apply_attn_dropout(self, attn_weights: torch.Tensor) -> torch.Tensor: 

489 """Apply attention dropout from the original HF component if present.""" 

490 if self.original_component is not None: 490 ↛ 496line 490 didn't jump to line 496 because the condition on line 490 was always true

491 dropout_fn = getattr(self.original_component, "attn_dropout", None) 

492 if dropout_fn is None: 

493 dropout_fn = getattr(self.original_component, "attention_dropout", None) 

494 if dropout_fn is not None and callable(dropout_fn): 

495 attn_weights = dropout_fn(attn_weights) 

496 return attn_weights 

497 

498 def _apply_output_projection(self, attn_output: torch.Tensor) -> torch.Tensor: 

499 """Apply the output projection (self.o) if present.""" 

500 if hasattr(self, "o") and self.o is not None: 

501 attn_output = self.o(attn_output) 

502 return attn_output 

503 

504 def _softmax_dropout_pattern( 

505 self, 

506 attn_scores: torch.Tensor, 

507 target_dtype: torch.dtype | None = None, 

508 upcast_to_fp32: bool = False, 

509 ) -> torch.Tensor: 

510 """Apply softmax, dropout, and hook_pattern to attention scores. 

511 

512 Args: 

513 attn_scores: Raw attention scores [batch, heads, q_seq, kv_seq]. 

514 target_dtype: If set, cast weights to this dtype after softmax. 

515 upcast_to_fp32: If True, compute softmax in float32 for numerical 

516 stability, then cast to target_dtype. 

517 """ 

518 if upcast_to_fp32: 

519 attn_weights = torch.nn.functional.softmax(attn_scores, dim=-1, dtype=torch.float32) 

520 if target_dtype is not None: 520 ↛ 526line 520 didn't jump to line 526 because the condition on line 520 was always true

521 attn_weights = attn_weights.to(target_dtype) 

522 else: 

523 attn_weights = torch.nn.functional.softmax(attn_scores, dim=-1) 

524 if target_dtype is not None: 

525 attn_weights = attn_weights.to(target_dtype) 

526 attn_weights = self._scrub_compatibility_pattern_nans(attn_weights) 

527 attn_weights = self._apply_attn_dropout(attn_weights) 

528 attn_weights = self.hook_pattern(attn_weights) 

529 return attn_weights 

530 

531 def _scrub_compatibility_pattern_nans(self, pattern: torch.Tensor) -> torch.Tensor: 

532 """Match HookedTransformer for fully masked attention rows.""" 

533 if self.compatibility_mode: 

534 pattern = torch.where(torch.isnan(pattern), torch.zeros_like(pattern), pattern) 

535 return pattern 

536 

537 def _normalize_compatibility_mask_sentinel(self, attention_mask: torch.Tensor) -> torch.Tensor: 

538 """Normalize additive mask sentinels before dtype conversion or addition.""" 

539 if self.compatibility_mode and attention_mask.is_floating_point(): 

540 attention_mask = attention_mask.masked_fill( 

541 attention_mask <= torch.finfo(attention_mask.dtype).min, -torch.inf 

542 ) 

543 return attention_mask 

544 

545 def _reshape_attn_output( 

546 self, 

547 attn_output: torch.Tensor, 

548 batch_size: int, 

549 seq_len: int, 

550 num_heads: int, 

551 head_dim: int, 

552 ) -> torch.Tensor: 

553 """Reshape attention output from [batch, heads, seq, dim] to [batch, seq, hidden].""" 

554 attn_output = attn_output.transpose(1, 2).contiguous() 

555 attn_output = attn_output.view(batch_size, seq_len, num_heads * head_dim) 

556 return attn_output 

557 

558 def _apply_reconstruct_attention_mask( 

559 self, 

560 attn_scores: torch.Tensor, 

561 attention_mask: torch.Tensor | None, 

562 seq_len: int, 

563 q_seq_len: int | None = None, 

564 ) -> torch.Tensor: 

565 """Apply causal and optional attention masking to reconstructed scores. 

566 

567 HuggingFace-style 4D masks already encode causal semantics, so they are 

568 treated as authoritative. Lower-rank masks do not, so the local causal 

569 mask is still applied before adding the caller-provided padding mask. 

570 

571 Args: 

572 attn_scores: Attention scores [batch, heads, q_seq_len, kv_seq_len]. 

573 attention_mask: Optional mask from the caller. 

574 seq_len: The KV sequence length (total positions including cache). 

575 q_seq_len: The query sequence length. When using KV cache this is 

576 shorter than seq_len. Defaults to seq_len when not provided. 

577 """ 

578 if q_seq_len is None: 

579 q_seq_len = seq_len 

580 min_dtype = torch.finfo(attn_scores.dtype).min 

581 mask_value = -torch.inf if self.compatibility_mode else min_dtype 

582 use_direct_hf_mask = attention_mask is not None and attention_mask.ndim >= 4 

583 # Bidirectional attention (encoders) and cross-attention have no causal 

584 # structure, so only synthesize the triangular mask for causal self-attention. 

585 apply_causal = self.is_causal and not self.is_cross_attention 

586 if not use_direct_hf_mask and apply_causal: 

587 # Rectangular causal mask: query i attends to KV 0..(offset+i) 

588 # where offset = kv_seq_len - q_seq_len (cached positions). 

589 causal_mask = torch.ones( 

590 q_seq_len, seq_len, device=attn_scores.device, dtype=torch.bool 

591 ) 

592 causal_mask = torch.tril(causal_mask, diagonal=seq_len - q_seq_len) 

593 attn_scores = attn_scores.masked_fill(~causal_mask, mask_value) 

594 

595 if attention_mask is not None: 

596 if attention_mask.shape[-1] != seq_len: 596 ↛ 597line 596 didn't jump to line 597 because the condition on line 596 was never true

597 attention_mask = attention_mask[..., :seq_len] 

598 if attention_mask.ndim >= 3 and attention_mask.shape[-2] != q_seq_len: 

599 # Extra query rows mean a full-sequence mask on a cached decode step, 

600 # where the live queries are the LAST rows; taking the first hands 

601 # every step position 0 (Baichuan-13B fuses ALiBi slopes in here). 

602 attention_mask = attention_mask[..., -q_seq_len:, :] 

603 

604 if attention_mask.dtype == torch.bool: 

605 attention_mask = torch.where( 

606 attention_mask, 

607 torch.zeros((), dtype=attn_scores.dtype, device=attn_scores.device), 

608 torch.full((), mask_value, dtype=attn_scores.dtype, device=attn_scores.device), 

609 ) 

610 else: 

611 attention_mask = self._normalize_compatibility_mask_sentinel(attention_mask) 

612 attention_mask = attention_mask.to(dtype=attn_scores.dtype) 

613 attn_scores = attn_scores + attention_mask 

614 

615 return attn_scores 

616 

617 def _get_n_heads(self, use_kv: bool = False) -> int: 

618 """Resolve the number of attention heads from config. 

619 

620 Args: 

621 use_kv: If True, return n_key_value_heads (for GQA) when available. 

622 """ 

623 assert self.config is not None, "config required to resolve n_heads" 

624 if use_kv: 

625 if hasattr(self.config, "n_key_value_heads") and self.config.n_key_value_heads: 

626 return self.config.n_key_value_heads 

627 if hasattr(self.config, "n_heads"): 

628 return self.config.n_heads 

629 return self.config.n_head 

630 

631 def _weight_layout_in_out(self, proj: Any) -> Optional[bool]: 

632 """Whether ``proj``'s wrapped module stores its weight as [in, out]. 

633 

634 Conv1D (GPT-2 style) stores [in_features, out_features]; nn.Linear stores 

635 [out_features, in_features]. Returns None when the wrapped module is 

636 neither, so callers can fall back to a shape heuristic. 

637 """ 

638 component = getattr(proj, "original_component", None) 

639 if isinstance(component, Conv1D): 

640 return True 

641 if isinstance(component, torch.nn.Linear): 

642 return False 

643 return None 

644 

645 def _reshape_weight_to_3d( 

646 self, 

647 weight: torch.Tensor, 

648 n_heads: int, 

649 pattern: str = "qkv", 

650 in_out_layout: Optional[bool] = None, 

651 ) -> torch.Tensor: 

652 """Reshape a 2D weight to 3D by splitting heads. 

653 

654 Args: 

655 weight: 2D weight tensor 

656 n_heads: Number of heads to split into 

657 pattern: "qkv" for [n_heads, d_model, d_head], "o" for [n_heads, d_head, d_model] 

658 in_out_layout: True if the weight is stored [in, out] (Conv1D), False 

659 for [out, in] (nn.Linear). None falls back to a shape heuristic, 

660 which cannot distinguish the two when the weight is square. 

661 """ 

662 if pattern == "o": 

663 if in_out_layout is None: 

664 # Heuristic: assumes [in, out] whenever the head split fits dim 0. 

665 in_out_layout = weight.shape[0] == n_heads * ( 

666 weight.shape[1] // n_heads 

667 if weight.shape[1] % n_heads == 0 

668 else weight.shape[0] // n_heads 

669 ) 

670 mat = weight if in_out_layout else weight.T 

671 self._check_head_split_width(mat, n_heads) 

672 return einops.rearrange( 

673 mat, "(n_heads d_head) d_model -> n_heads d_head d_model", n_heads=n_heads 

674 ) 

675 # QKV pattern 

676 if in_out_layout is None: 

677 in_out_layout = weight.shape[0] % n_heads != 0 

678 mat = weight.T if in_out_layout else weight 

679 self._check_head_split_width(mat, n_heads) 

680 return einops.rearrange( 

681 mat, "(n_heads d_head) d_model -> n_heads d_model d_head", n_heads=n_heads 

682 ) 

683 

684 def _check_head_split_width(self, mat: torch.Tensor, n_heads: int) -> None: 

685 """Refuse a head split whose width disagrees with the model's geometry. 

686 

687 einops only needs divisibility, so per-layer geometry (gemma4 KV-shared 

688 layers) would factorize into wrong-shaped heads silently. MLA's 

689 bind-time ``_v_head_dim`` is a legitimate second width (o_proj). 

690 """ 

691 d_head = getattr(self.config, "d_head", None) 

692 if not d_head: 

693 return 

694 allowed_dims = {d_head} 

695 v_head_dim = getattr(self, "_v_head_dim", None) 

696 if v_head_dim: 

697 allowed_dims.add(v_head_dim) 

698 if mat.shape[0] % n_heads == 0 and mat.shape[0] // n_heads in allowed_dims: 

699 return 

700 raise PerLayerGeometryError( 

701 f"Cannot split {tuple(mat.shape)} weight on '{self.name}' into " 

702 f"{n_heads} heads of d_head in {sorted(allowed_dims)}: this " 

703 "layer's attention geometry differs from the config-level scalars " 

704 "(per-layer num_key_value_heads/head_dim, e.g. gemma4 KV-shared " 

705 "or K==V layers). Read the wrapped module's weight directly for " 

706 "this layer instead." 

707 ) 

708 

709 def _project_per_head_qkv( 

710 self, 

711 linear_bridge: "GeneralizedComponent", 

712 input_4d: torch.Tensor, 

713 n_heads: int, 

714 d_head: int, 

715 ) -> torch.Tensor: 

716 """Per-head Q/K/V projection over a 4D residual fork. 

717 

718 Plain nn.Linear applied to [batch, pos, H, d_model] broadcasts the 

719 same weight across heads' copies — which for the split-qkv fork means 

720 head h's copy sees every head's W rows, not just head h's. This routes 

721 head h's copy through head h's W slice only via a per-head einsum. 

722 

723 Fires `linear_bridge.hook_out` on the flat 3D tensor so the hook sees 

724 the same shape as the default path and downstream code receives a 

725 consistent 4D `[B, S, H, d_head]` regardless of whether the user's 

726 hook modified the tensor (which would otherwise trigger the 

727 `hook_conversion.revert` 4D→3D flatten). 

728 """ 

729 component = linear_bridge.original_component 

730 assert component is not None, "LinearBridge.original_component not set" 

731 weight = require_readable_weight( 

732 component.weight, 

733 operation="project per head (use_split_qkv_input / use_attn_in)", 

734 owner=component, 

735 ) 

736 bias = component.bias 

737 w3d = einops.rearrange( 

738 weight, 

739 "(n_heads d_head) d_model -> n_heads d_model d_head", 

740 n_heads=n_heads, 

741 d_head=d_head, 

742 ) 

743 out = torch.einsum("bshd,hde->bshe", input_4d, w3d) 

744 if bias is not None: 

745 b2d = einops.rearrange(bias, "(n_heads d_head) -> n_heads d_head", n_heads=n_heads) 

746 assert isinstance(b2d, torch.Tensor) 

747 out = out + b2d 

748 # Flatten to 3D for hook_out (matches default-path shape); the 

749 # hook_conversion reshapes to 4D for the user's fwd_hook, then reverts 

750 # to 3D if the hook returned a modified tensor. Return 4D always. 

751 b, s = out.shape[0], out.shape[1] 

752 out_flat = out.reshape(b, s, n_heads * d_head) 

753 out_flat = linear_bridge.hook_out(out_flat) 

754 return out_flat.reshape(b, s, n_heads, d_head) 

755 

756 def _compute_per_head_result( 

757 self, 

758 z_4d: torch.Tensor, 

759 n_heads: int, 

760 d_head: int, 

761 ) -> torch.Tensor: 

762 """Per-head attention output pre-sum across heads. 

763 

764 Computes (z[..., h, :] @ W_O_per_head[h]) for each head h, fires 

765 hook_result on the resulting [batch, pos, n_heads, d_model], then sums 

766 across heads and adds b_O. Distributive over weight folding 

767 (`sum_h z_h @ W_O_h + b_O == z_flat @ W_O.T + b_O`), so compat-mode and 

768 raw-weight paths produce identical logits. 

769 """ 

770 o = self.o.original_component 

771 weight = o.weight 

772 bias = getattr(o, "bias", None) 

773 # HF Conv1D (GPT-2, GPT-J, CodeGen) stores weight as [in, out]; nn.Linear 

774 # stores [out, in]. When W_O is square (d_model == n_heads*d_head, which 

775 # is the common case), shape alone is ambiguous — dispatch on module 

776 # type instead. 

777 weight_is_in_out = type(o).__name__ == "Conv1D" 

778 if weight_is_in_out: 

779 w_per_head = einops.rearrange( 

780 weight, 

781 "(n_heads d_head) d_model -> n_heads d_head d_model", 

782 n_heads=n_heads, 

783 d_head=d_head, 

784 ) 

785 else: 

786 w_per_head = einops.rearrange( 

787 weight, 

788 "d_model (n_heads d_head) -> n_heads d_head d_model", 

789 n_heads=n_heads, 

790 d_head=d_head, 

791 ) 

792 per_head = torch.einsum("bshd,hdm->bshm", z_4d, w_per_head) 

793 per_head = self.hook_result(per_head) 

794 summed = per_head.sum(dim=-2)