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