Coverage for transformer_lens/model_bridge/generalized_components/gated_delta_net.py: 89%
197 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-10-10 05:10 +0000
1"""GatedDeltaNet bridge for Qwen3.5/Qwen3Next linear-attention layers.
3Reimplements forward (prefill only) to expose mech-interp-relevant intermediate
4states. Falls back to HF native forward during autoregressive generation where
5cache state management is required.
6"""
7import sys
8from typing import Any, Dict, Optional
10import torch
11import torch.nn.functional as F
13from transformer_lens.ActivationCache import ActivationCache
14from transformer_lens.hook_points import HookPoint
15from transformer_lens.model_bridge.generalized_components.base import (
16 GeneralizedComponent,
17)
18from transformer_lens.model_bridge.generalized_components.ssm_protocol import (
19 SSMStateHookMixin,
20)
22_MISSING = object()
25def _resolve_hf_fn(hf: torch.nn.Module, instance_attr: str, module_attr: str) -> Any:
26 """Resolve the function HF's GatedDeltaNet forward calls.
28 transformers < 5.15 stores it on the instance; >= 5.15 calls a module-level
29 function of the modeling file, which dispatches to the hub kernel when installed.
30 """
31 if hasattr(hf, instance_attr):
32 return getattr(hf, instance_attr)
33 fn = getattr(sys.modules[type(hf).__module__], module_attr, _MISSING)
34 if fn is _MISSING:
35 raise AttributeError(
36 f"{type(hf).__name__} has neither a `{instance_attr}` attribute "
37 f"(transformers < 5.15) nor a module-level `{module_attr}` in "
38 f"{type(hf).__module__} (transformers >= 5.15); GatedDeltaNetBridge "
39 "supports transformers >= 5.9 and was tested up to 5.17."
40 )
41 return fn
44class GatedDeltaNetBridge(SSMStateHookMixin, GeneralizedComponent):
45 """Bridge for GatedDeltaNet linear-attention with full hook decomposition.
47 Hooks (prefill, in execution order):
48 hook_in: input hidden_states [batch, seq, d_model]
49 hook_q_pre_conv: Q after projection, before conv [batch, seq, n_k_heads, head_k_dim]
50 hook_k_pre_conv: K after projection, before conv [batch, seq, n_k_heads, head_k_dim]
51 hook_v_pre_conv: V after projection, before conv [batch, seq, n_v_heads, head_v_dim]
52 hook_q: Q after conv, pre-GQA-expansion [batch, seq, n_k_heads, head_k_dim]
53 Note: on standard attn layers, hook_q is post-projection. Here it's
54 post-conv — use hook_q_pre_conv for the projection-only output.
55 hook_k: K after conv [batch, seq, n_k_heads, head_k_dim]
56 hook_v: V after conv [batch, seq, n_v_heads, head_v_dim]
57 hook_beta_logit: pre-sigmoid write gate logit, per v-head [batch, seq, n_v_heads]
58 hook_beta: write strength sigmoid(b), per v-head [batch, seq, n_v_heads]
59 hook_log_decay: log-space decay g (NEGATIVE; multiplicative decay = exp(g)),
60 per v-head [batch, seq, n_v_heads]
61 hook_recurrence_out: output of linear recurrence [batch, seq, n_v_heads, head_v_dim]
62 hook_gate_input: z tensor (pre-silu) for GatedRMSNorm [batch, seq, n_v_heads, head_v_dim]
63 hook_ssm_state: recurrent state trajectory S_t [batch, seq, n_v_heads, head_k_dim,
64 head_v_dim] — fires ONLY on the opt-in eager-scan path (eager_scan=True),
65 which swaps the fused kernel for a Python delta-rule scan so S_t can be
66 read/patched. hook_ssm_write (alias -> hook_beta) is the write strength and
67 propagates through the scan.
68 hook_out: final output to residual stream [batch, seq, d_model]
70 During generation (cache_params present), only hook_in/hook_out fire.
72 Property aliases:
73 W_in_proj_qkvz, W_in_proj_ba, W_out_proj, A_log, dt_bias
74 """
76 hook_aliases = {
77 "hook_linear_attn_in": "hook_in",
78 "hook_linear_attn_out": "hook_out",
79 # Canonical SSM vocabulary (additive) — maps onto GDN's existing hooks so
80 # interp tools can find these quantities by the same name across families.
81 # Semantic mapping: q reads the state (~C), k writes to it (~B), the gate
82 # g is the decay, beta is the per-step write strength.
83 "hook_ssm_out": "hook_out",
84 "hook_ssm_C": "hook_q",
85 "hook_ssm_B": "hook_k",
86 "hook_ssm_decay": "hook_log_decay",
87 "hook_ssm_write": "hook_beta",
88 }
90 property_aliases = {
91 "W_in_proj_qkvz": "in_proj_qkvz.weight",
92 "W_in_proj_ba": "in_proj_ba.weight",
93 "W_out_proj": "out_proj.weight",
94 "A_log": "A_log",
95 "dt_bias": "dt_bias",
96 }
98 def __init__(
99 self,
100 name: str,
101 config: Optional[Any] = None,
102 submodules: Optional[Dict[str, GeneralizedComponent]] = None,
103 **kwargs,
104 ):
105 super().__init__(name, config=config, submodules=submodules or {}, **kwargs)
106 # Pre-conv (after projection split, before causal conv mixes positions)
107 self.hook_q_pre_conv = HookPoint()
108 self.hook_k_pre_conv = HookPoint()
109 self.hook_v_pre_conv = HookPoint()
110 # Post-conv (pre-GQA-expansion, pre-recurrence)
111 self.hook_q = HookPoint()
112 self.hook_k = HookPoint()
113 self.hook_v = HookPoint()
114 # Gate parameters (per v-head)
115 self.hook_beta_logit = HookPoint()
116 self.hook_beta = HookPoint()
117 self.hook_log_decay = HookPoint()
118 # Recurrence output + gated norm input
119 self.hook_recurrence_out = HookPoint()
120 self.hook_gate_input = HookPoint()
121 # hook_ssm_state + eager_scan come from SSMStateHookMixin; hook_ssm_write is
122 # an alias onto hook_beta (the delta rule's write is state-dependent).
124 def forward(self, *args: Any, **kwargs: Any) -> Any:
125 if self.original_component is None: 125 ↛ 126line 125 didn't jump to line 126 because the condition on line 125 was never true
126 raise RuntimeError(f"Original component not set for {self.name}.")
128 if kwargs.get("cache_params") is not None:
129 return self._native_forward(*args, **kwargs)
130 return self._hooked_forward(*args, **kwargs)
132 def _native_forward(self, *args: Any, **kwargs: Any) -> Any:
133 """Delegate to HF with hook_in/hook_out only (generation path)."""
134 assert self.original_component is not None
135 if "hidden_states" in kwargs: 135 ↛ 137line 135 didn't jump to line 137 because the condition on line 135 was always true
136 kwargs["hidden_states"] = self.hook_in(kwargs["hidden_states"])
137 elif len(args) > 0 and isinstance(args[0], torch.Tensor):
138 args = (self.hook_in(args[0]),) + args[1:]
140 output = self.original_component(*args, **kwargs)
142 if isinstance(output, tuple) and len(output) > 0: 142 ↛ 143line 142 didn't jump to line 143 because the condition on line 142 was never true
143 first = output[0]
144 if isinstance(first, torch.Tensor):
145 return (self.hook_out(first),) + output[1:]
146 return output
147 if isinstance(output, torch.Tensor): 147 ↛ 149line 147 didn't jump to line 149 because the condition on line 147 was always true
148 return self.hook_out(output)
149 return output
151 def _hooked_forward(self, *args: Any, **kwargs: Any) -> Any:
152 """Reimplemented forward exposing all intermediate states (prefill)."""
153 hf: Any = self.original_component
155 if "hidden_states" in kwargs:
156 hidden_states = kwargs["hidden_states"]
157 elif len(args) > 0 and isinstance(args[0], torch.Tensor): 157 ↛ 160line 157 didn't jump to line 160 because the condition on line 157 was always true
158 hidden_states = args[0]
159 else:
160 raise ValueError("Could not find hidden_states")
162 attention_mask = kwargs.get("attention_mask")
163 if attention_mask is not None: 163 ↛ 165line 163 didn't jump to line 165 because the condition on line 163 was never true
164 # Inline masking — avoids hard dependency on qwen3_next module
165 hidden_states = hidden_states * attention_mask.unsqueeze(-1)
167 hidden_states = self.hook_in(hidden_states)
168 batch_size, seq_len, _ = hidden_states.shape
170 # --- Projections (two layouts: fused vs split) ---
171 if hasattr(hf, "in_proj_qkvz"):
172 # Qwen3Next: fused Q+K+V+Z projection, fused beta+alpha
173 projected_qkvz = hf.in_proj_qkvz(hidden_states)
174 projected_ba = hf.in_proj_ba(hidden_states)
175 query, key, value, z, b, a = hf.fix_query_key_value_ordering(
176 projected_qkvz, projected_ba
177 )
178 else:
179 # Qwen3.5: separate projections (in_proj_qkv, in_proj_z, in_proj_b, in_proj_a)
180 mixed_qkv_flat = hf.in_proj_qkv(hidden_states)
181 z = hf.in_proj_z(hidden_states).reshape(batch_size, seq_len, -1, hf.head_v_dim)
182 b = hf.in_proj_b(hidden_states)
183 a = hf.in_proj_a(hidden_states)
184 # Split QKV and reshape to per-head for pre-conv hooks
185 q_flat, k_flat, v_flat = torch.split(
186 mixed_qkv_flat, [hf.key_dim, hf.key_dim, hf.value_dim], dim=-1
187 )
188 query = q_flat.reshape(batch_size, seq_len, -1, hf.head_k_dim)
189 key = k_flat.reshape(batch_size, seq_len, -1, hf.head_k_dim)
190 value = v_flat.reshape(batch_size, seq_len, -1, hf.head_v_dim)
192 # --- Pre-conv hooks (per-head shape, before conv mixes positions) ---
193 query = self.hook_q_pre_conv(query)
194 key = self.hook_k_pre_conv(key)
195 value = self.hook_v_pre_conv(value)
197 # Flatten for conv
198 query, key, value = (x.reshape(x.shape[0], x.shape[1], -1) for x in (query, key, value))
200 # --- Causal Convolution ---
201 mixed_qkv = torch.cat((query, key, value), dim=-1).transpose(1, 2)
202 # None only on transformers < 5.15 without the causal-conv1d kernel.
203 # Positional: the >= 5.15 signature takes no x=/seq_idx=.
204 causal_conv1d_fn = _resolve_hf_fn(hf, "causal_conv1d_fn", "causal_conv1d_fn")
205 if causal_conv1d_fn is not None:
206 mixed_qkv = causal_conv1d_fn(
207 mixed_qkv,
208 hf.conv1d.weight.squeeze(1),
209 hf.conv1d.bias,
210 activation=hf.activation,
211 )
212 else:
213 mixed_qkv = F.silu(hf.conv1d(mixed_qkv)[:, :, :seq_len])
214 mixed_qkv = mixed_qkv.transpose(1, 2)
216 # Split post-conv into per-head Q, K, V
217 query, key, value = torch.split(
218 mixed_qkv,
219 [hf.key_dim, hf.key_dim, hf.value_dim],
220 dim=-1,
221 )
222 query = query.reshape(batch_size, seq_len, -1, hf.head_k_dim)
223 key = key.reshape(batch_size, seq_len, -1, hf.head_k_dim)
224 value = value.reshape(batch_size, seq_len, -1, hf.head_v_dim)
226 # --- Post-conv hooks (pre-GQA-expansion, pre-recurrence) ---
227 query = self.hook_q(query)
228 key = self.hook_k(key)
229 value = self.hook_v(value)
231 # --- Gate parameters (per v-head) ---
232 b = self.hook_beta_logit(b)
233 beta = self.hook_beta(b.sigmoid())
235 # g is log-space decay (NEGATIVE); multiplicative decay = exp(g)
236 g = -hf.A_log.float().exp() * F.softplus(a.float() + hf.dt_bias)
237 g = self.hook_log_decay(g)
239 # GQA expansion (Q/K from n_k_heads → n_v_heads)
240 if hf.num_v_heads // hf.num_k_heads > 1:
241 repeat = hf.num_v_heads // hf.num_k_heads
242 query = query.repeat_interleave(repeat, dim=2)
243 key = key.repeat_interleave(repeat, dim=2)
245 # --- Core linear recurrence ---
246 # Default: HF's opaque fused kernel. eager_scan: a Python delta-rule scan
247 # that fires hook_ssm_state so the state trajectory can be intervened on.
248 if self.eager_scan:
249 _, core_out = self._gated_delta_scan(query, key, value, g, beta, fire_hook=True)
250 core_out = core_out.to(value.dtype)
251 else:
252 chunk_gated_delta_rule = _resolve_hf_fn(
253 hf, "chunk_gated_delta_rule", "torch_chunk_gated_delta_rule"
254 )
255 core_out, _ = chunk_gated_delta_rule(
256 query,
257 key,
258 value,
259 g=g,
260 beta=beta,
261 initial_state=None,
262 output_final_state=False,
263 use_qk_l2norm_in_kernel=True,
264 )
265 core_out = self.hook_recurrence_out(core_out)
267 # --- Gated RMSNorm: norm(core_out) * silu(z) ---
268 z = self.hook_gate_input(z)
269 z_shape = z.shape
270 core_out = hf.norm(
271 core_out.reshape(-1, core_out.shape[-1]),
272 z.reshape(-1, z.shape[-1]),
273 )
274 core_out = core_out.reshape(z_shape).reshape(batch_size, seq_len, -1)
276 # --- Output projection ---
277 output = hf.out_proj(core_out)
278 return self.hook_out(output)
280 @staticmethod
281 def _l2norm(x: torch.Tensor, eps: float = 1e-6) -> torch.Tensor:
282 """Match the kernel's internal Q/K L2 normalization (use_qk_l2norm_in_kernel)."""
283 return x / torch.sqrt((x * x).sum(-1, keepdim=True) + eps)
285 def _gated_delta_scan(
286 self,
287 query: torch.Tensor,
288 key: torch.Tensor,
289 value: torch.Tensor,
290 g: torch.Tensor,
291 beta: torch.Tensor,
292 fire_hook: bool,
293 ) -> tuple[torch.Tensor, torch.Tensor]:
294 """Eager gated-delta-rule recurrence — the readable form of the fused kernel.
296 Reproduces HF ``torch_recurrent_gated_delta_rule``: L2-normalize Q/K, scale
297 Q by ``1/sqrt(head_k_dim)``, then per step::
299 S_t = exp(g_t) · S_{t-1} (decay)
300 delta_t = (v_t - S_t^T k_t) · beta_t (delta rule: remove then write)
301 S_t = S_t + k_t ⊗ delta_t -> hook_ssm_state trajectory
302 o_t = S_t^T q_t
304 Q/K are assumed already GQA-expanded to ``n_v_heads``. Matches the fused
305 kernel only to fp tolerance (≈1e-6 fp32), never bit-for-bit; O(seq) Python.
307 Args:
308 query, key: ``[batch, seq, n_v_heads, head_k_dim]`` (post-conv, pre-norm).
309 value: ``[batch, seq, n_v_heads, head_v_dim]``.
310 g: log-space decay ``[batch, seq, n_v_heads]`` (NEGATIVE; decay = exp(g)).
311 beta: write strength ``[batch, seq, n_v_heads]``.
312 fire_hook: if True, fire ``hook_ssm_state`` on the trajectory and
313 recompute ``o_t`` from the (possibly patched) state.
315 Returns:
316 ``(state_traj, core_out)`` — state ``[batch, seq, n_v_heads, head_k_dim,
317 head_v_dim]`` and output ``[batch, seq, n_v_heads, head_v_dim]``.
318 """
319 q = self._l2norm(query.float())
320 k = self._l2norm(key.float())
321 v = value.float()
322 q = q * (q.shape[-1] ** -0.5)
323 g_f = g.float()
324 beta_f = beta.float()
326 batch, seq_len, n_v, k_dim = q.shape
327 v_dim = v.shape[-1]
328 state = torch.zeros(batch, n_v, k_dim, v_dim, dtype=q.dtype, device=q.device)
329 states = []
330 outs = []
331 for t in range(seq_len):
332 q_t, k_t, v_t = q[:, t], k[:, t], v[:, t] # [batch, n_v, dim]
333 decay = g_f[:, t].exp()[:, :, None, None] # [batch, n_v, 1, 1]
334 beta_t = beta_f[:, t][:, :, None] # [batch, n_v, 1]
335 state = state * decay
336 kv_mem = (state * k_t[:, :, :, None]).sum(dim=-2) # [batch, n_v, v_dim]
337 delta = (v_t - kv_mem) * beta_t
338 state = state + k_t[:, :, :, None] * delta[:, :, None, :]
339 states.append(state)
340 outs.append((state * q_t[:, :, :, None]).sum(dim=-2))
342 state_traj = torch.stack(states, dim=1) # [batch, seq, n_v, k_dim, v_dim]
343 if fire_hook:
344 state_traj = self.hook_ssm_state(state_traj)
345 # Recompute o_t = S_t^T q_t from the (possibly patched) trajectory.
346 core_out = torch.einsum("bshkv,bshk->bshv", state_traj, q)
347 else:
348 core_out = torch.stack(outs, dim=1) # [batch, seq, n_v, v_dim]
349 return state_traj, core_out
351 def compute_ssm_state(
352 self,
353 cache: ActivationCache,
354 layer_idx: int,
355 time_step: Optional[int] = None,
356 ) -> torch.Tensor:
357 """Reconstruct the recurrent state ``S`` of the gated delta rule from cache.
359 Read-only: replays the eager delta-rule scan (``_gated_delta_scan``) on the
360 cached hook_q/k/v/beta/log_decay — no ``forward()`` re-run. Faithful (the
361 full delta rule, key-removal included), unlike ``compute_effective_attention``
362 which is a gated-linear-attention heuristic that drops key removal.
364 Requires the interior hooks, which fire only on the hooked prefill path:
365 call ``run_with_cache(tokens, use_cache=False)`` so ``cache_params`` is None.
367 On padded batches the cached hooks are unmasked, so ``S`` is exact only at
368 non-pad positions; pad-position state is out of contract.
370 Args:
371 cache: ActivationCache from ``run_with_cache(..., use_cache=False)``.
372 layer_idx: Block index for this linear_attn layer.
373 time_step: If given, return only ``S`` at that position
374 (``[batch, n_v_heads, head_k_dim, head_v_dim]``). None returns
375 every step.
377 Returns:
378 ``[batch, seq, n_v_heads, head_k_dim, head_v_dim]`` for all steps, or
379 ``[batch, n_v_heads, head_k_dim, head_v_dim]`` for a single ``time_step``.
381 Memory is O(batch · n_v_heads · seq · head_k_dim · head_v_dim); pass
382 ``time_step`` (or short sequences) when that is too large.
383 """
384 prefix = f"blocks.{layer_idx}.linear_attn"
385 q_key = f"{prefix}.hook_q"
386 k_key = f"{prefix}.hook_k"
387 v_key = f"{prefix}.hook_v"
388 beta_key = f"{prefix}.hook_beta"
389 decay_key = f"{prefix}.hook_log_decay"
391 for key in (q_key, k_key, v_key, beta_key, decay_key):
392 if key not in cache:
393 raise RuntimeError(
394 f"compute_ssm_state needs {key!r} in cache. Run "
395 "run_with_cache(tokens, use_cache=False) on the bridge first."
396 )
398 q = cache[q_key].float() # [batch, seq, n_k_heads, head_k_dim]
399 k = cache[k_key].float()
400 v = cache[v_key].float() # [batch, seq, n_v_heads, head_v_dim]
401 beta = cache[beta_key].float() # [batch, seq, n_v_heads]
402 g = cache[decay_key].float() # [batch, seq, n_v_heads]
404 # GQA expansion to n_v_heads (Q/K carry n_k_heads).
405 n_v = v.shape[2]
406 if q.shape[2] < n_v: 406 ↛ 407line 406 didn't jump to line 407 because the condition on line 406 was never true
407 repeat = n_v // q.shape[2]
408 q = q.repeat_interleave(repeat, dim=2)
409 k = k.repeat_interleave(repeat, dim=2)
411 state_traj, _ = self._gated_delta_scan(q, k, v, g, beta, fire_hook=False)
412 if time_step is not None:
413 return state_traj[:, time_step]
414 return state_traj
416 def compute_effective_attention(
417 self,
418 cache: ActivationCache,
419 layer_idx: int,
420 ) -> torch.Tensor:
421 """Materialize a heuristic effective-attention matrix from cached hooks.
423 Uses the gated-linear-attention form of the recurrence (the exact gated
424 *delta* rule additionally removes the key being written)::
426 S_t ≈ exp(g_t) * S_{t-1} + beta_t * v_t @ k_t^T
427 o_t = S_t^T @ q_t
428 M[i,j] = (q_i^T @ k_j) * beta_j * prod_{t=j+1}^{i} exp(g_t)
430 so ``M`` is an interpretability heuristic, not a faithful output decomposition.
432 Requires the interior hooks (hook_q/k/beta/log_decay), which fire only on
433 the hooked prefill path: call ``run_with_cache(tokens, use_cache=False)``
434 so ``cache_params`` is None. The default cached path exposes only
435 hook_in/hook_out and this method then raises.
437 **Measured divergence (tiny random-init test fixture, seed-stable):**
439 - *L2-norm gap.* The fused kernel L2-normalizes Q/K internally
440 (``use_qk_l2norm_in_kernel=True``) but the hooked Q/K are pre-norm, so
441 ``M`` differs from the normalized form by ≈1.0 relative when Q/K norms
442 are small/non-uniform (the random-init regime); the gap shrinks toward 0
443 as norms equalize after training.
444 - *Delta-rule omission.* Even with normalized Q/K, ``M @ V`` reconstructs
445 the fused-kernel ``hook_recurrence_out`` only to O(1) relative error
446 because the key-removal term is dropped.
448 Args:
449 cache: ActivationCache from ``run_with_cache(..., use_cache=False)``.
450 layer_idx: Block index for this linear_attn layer.
452 Returns:
453 ``[batch, n_v_heads, seq, seq]`` causal matrix (upper triangle zero).
455 Cost is O(batch * n_heads * seq^2); use on short sequences.
456 """
457 prefix = f"blocks.{layer_idx}.linear_attn"
458 q_key = f"{prefix}.hook_q"
459 k_key = f"{prefix}.hook_k"
460 beta_key = f"{prefix}.hook_beta"
461 decay_key = f"{prefix}.hook_log_decay"
463 for key in [q_key, k_key, beta_key, decay_key]:
464 if key not in cache:
465 raise RuntimeError(
466 f"compute_effective_attention needs {key!r} in cache. "
467 "Run run_with_cache() on the bridge first."
468 )
470 # [batch, seq, n_k_heads, head_k_dim] — pre-GQA-expansion
471 q = cache[q_key].float()
472 k = cache[k_key].float()
473 beta = cache[beta_key].float() # [batch, seq, n_v_heads]
474 g = cache[decay_key].float() # [batch, seq, n_v_heads]
476 # GQA expansion to match n_v_heads
477 if q.shape[2] < beta.shape[-1]: 477 ↛ 478line 477 didn't jump to line 478 because the condition on line 477 was never true
478 repeat = beta.shape[-1] // q.shape[2]
479 q = q.repeat_interleave(repeat, dim=2)
480 k = k.repeat_interleave(repeat, dim=2)
482 batch, seq, n_heads, d_head = q.shape
484 # QK similarity: [batch, n_heads, seq_i, seq_j]
485 q_perm = q.permute(0, 2, 1, 3)
486 k_perm = k.permute(0, 2, 1, 3)
487 qk = torch.matmul(q_perm, k_perm.transpose(-2, -1))
489 # Cumulative decay: L[i,j] = exp(sum g[j+1..i])
490 g_perm = g.permute(0, 2, 1) # [batch, n_heads, seq]
491 cumsum_g = torch.cumsum(g_perm, dim=-1)
492 L_log = cumsum_g[:, :, :, None] - cumsum_g[:, :, None, :]
494 causal_mask = torch.tril(torch.ones(seq, seq, dtype=torch.bool, device=q.device))
495 L = torch.where(causal_mask[None, None], torch.exp(L_log), torch.zeros_like(L_log))
497 # M[i,j] = qk[i,j] * beta[j] * L[i,j]
498 beta_col = beta.permute(0, 2, 1)[:, :, None, :]
499 return qk * beta_col * L