Coverage for transformer_lens/model_bridge/generalized_components/ssm2_mixer.py: 94%
127 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-01 16:23 +0000
1"""Wrap-don't-reimplement bridge for HF's Mamba2Mixer, plus SSD effective attention."""
2from typing import Any, NamedTuple, Optional
4import torch
6from transformer_lens.ActivationCache import ActivationCache
7from transformer_lens.hook_points import HookPoint
8from transformer_lens.model_bridge.generalized_components.base import (
9 GeneralizedComponent,
10)
11from transformer_lens.model_bridge.generalized_components.ssm_protocol import (
12 SSMStateHookMixin,
13)
16class _SSDTerms(NamedTuple):
17 """SSD intermediates reconstructed from cached Mamba-2 hooks (HF discretization)."""
19 dt: torch.Tensor # [batch, seq, num_heads]
20 L: torch.Tensor # [batch, num_heads, seq, seq] causal decay (upper triangle zero)
21 x: torch.Tensor # [batch, seq, num_heads, head_dim] SSM input per head
22 B: torch.Tensor # [batch, seq, num_heads, state] (group-expanded)
23 C: torch.Tensor # [batch, seq, num_heads, state] (group-expanded)
26class SSM2MixerBridge(SSMStateHookMixin, GeneralizedComponent):
27 """Opaque wrapper around Mamba-2's Mamba2Mixer.
29 Structural differences from Mamba-1:
30 - No x_proj/dt_proj; in_proj fuses gate, hidden_B_C, and dt into one output.
31 - Has an inner norm (``MambaRMSNormGated``) taking two inputs; exposed at
32 ``mixer.inner_norm`` (renamed from HF's ``norm``) to disambiguate from the
33 block-level norm.
34 - Multi-head with ``num_heads``, ``head_dim``, ``n_groups`` (GQA-like).
35 - ``A_log``, ``dt_bias``, ``D`` are ``[num_heads]`` parameters reached via
36 ``GeneralizedComponent.__getattr__`` delegation.
38 Decode-step caveat: ``conv1d.hook_out`` fires only on prefill during
39 stateful generation; see ``DepthwiseConv1DBridge`` for the reason.
40 """
42 hook_aliases = {
43 "hook_in_proj": "in_proj.hook_out",
44 "hook_conv": "conv1d.hook_out",
45 "hook_inner_norm": "inner_norm.hook_out",
46 # Canonical SSM vocabulary (additive). Mamba-2 fuses B/C/dt into the
47 # in_proj/conv1d outputs, so only the mixer output has a granular hook by
48 # default; B/C/dt/decay are reconstructed via compute_effective_attention /
49 # compute_ssm_state. hook_ssm_write / hook_ssm_state are real HookPoints
50 # that fire only on the opt-in eager-scan intervention path.
51 "hook_ssm_out": "hook_out",
52 }
54 def __init__(self, *args: Any, **kwargs: Any) -> None:
55 super().__init__(*args, **kwargs) # mixin adds hook_ssm_state + eager_scan
56 # Real per-step write term dt·(x⊗B); fires only on the eager-scan path.
57 self.hook_ssm_write = HookPoint()
59 def forward(self, *args: Any, **kwargs: Any) -> Any:
60 """Hook the input, run HF torch_forward (or the eager scan), hook the output."""
61 if self.original_component is None: 61 ↛ 62line 61 didn't jump to line 62 because the condition on line 61 was never true
62 raise RuntimeError(
63 f"Original component not set for {self.name}. "
64 "Call set_original_component() first."
65 )
67 hidden_states: Optional[torch.Tensor] = None
68 if len(args) > 0 and isinstance(args[0], torch.Tensor):
69 hidden_states = self.hook_in(args[0])
70 args = (hidden_states,) + args[1:]
71 elif "hidden_states" in kwargs and isinstance(kwargs["hidden_states"], torch.Tensor): 71 ↛ 77line 71 didn't jump to line 77 because the condition on line 71 was always true
72 hidden_states = self.hook_in(kwargs["hidden_states"])
73 kwargs["hidden_states"] = hidden_states
75 # Eager-scan slow path: opt-in flag, prefill only (no cache_params). Keyed
76 # on explicit state, never on hook-registry introspection.
77 if self.eager_scan and hidden_states is not None and kwargs.get("cache_params") is None:
78 output: Any = self._eager_scan_forward(hidden_states, kwargs.get("attention_mask"))
79 else:
80 output = self.original_component(*args, **kwargs)
82 if isinstance(output, tuple) and len(output) > 0:
83 first = output[0]
84 if isinstance(first, torch.Tensor): 84 ↛ 86line 84 didn't jump to line 86 because the condition on line 84 was always true
85 return (self.hook_out(first),) + output[1:]
86 return output
87 if isinstance(output, torch.Tensor): 87 ↛ 89line 87 didn't jump to line 89 because the condition on line 87 was always true
88 return self.hook_out(output)
89 return output
91 def _eager_scan_forward(
92 self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor]
93 ) -> torch.Tensor:
94 """Reimplement the Mamba-2 mixer prefill forward with an eager Python scan.
96 Reuses HF's in_proj / conv1d / inner-norm / out_proj submodules (so their
97 hooks still fire) and reimplements ONLY the SSD recurrence::
99 write_t = dt_t · (x_t ⊗ B_t) -> hook_ssm_write [b, seq, h, d, n]
100 S_t = exp(dt_t·A) · S_{t-1} + write_t
101 S = stack_t S_t -> hook_ssm_state [b, seq, h, d, n]
102 y_t = C_t · S_t (recomputed from the hooked state)
104 Intervening on hook_ssm_write re-runs the recurrence (the change propagates
105 to every later state); hook_ssm_state is the post-scan trajectory, so a
106 patch there changes only the same-position output y_t = C_t·S_t, not the
107 forward recurrence — patch hook_ssm_write for a propagating state edit.
109 Kernel-divergence caveat: this eager scan reproduces HF's fused chunk scan
110 only to floating-point tolerance (≈1e-6 fp32), never bit-for-bit, on any
111 family. It is O(seq) Python and materializes an O(b·seq·h·d·n) write/state
112 tensor — orders of magnitude slower/heavier than the kernel. Prefill only.
113 """
114 oc: Any = self.original_component
115 batch, seq_len, _ = hidden_states.shape
116 intermediate, num_heads, head_dim = oc.intermediate_size, oc.num_heads, oc.head_dim
117 n_groups, state = oc.n_groups, oc.ssm_state_size
119 def _mask_pad(states: torch.Tensor) -> torch.Tensor:
120 # Mirror HF apply_mask_to_padding_states (no-op for batch 1 / unpadded).
121 if attention_mask is not None and attention_mask.shape[1] > 1 and batch > 1:
122 return (states * attention_mask[:, :, None]).to(states.dtype)
123 return states
125 # 1-2. Projection + conv — reuse the HF submodule bridges (their hooks fire).
126 # Match HF: mask padding before in_proj and after conv (before B/C split).
127 projected = oc.in_proj(_mask_pad(hidden_states))
128 d_mlp = (projected.shape[-1] - 2 * intermediate - 2 * n_groups * state - num_heads) // 2
129 _, _, gate, hidden_B_C, dt = projected.split(
130 [d_mlp, d_mlp, intermediate, oc.conv_dim, num_heads], dim=-1
131 )
132 hidden_B_C = oc.act(oc.conv1d(hidden_B_C.transpose(1, 2))[..., :seq_len].transpose(1, 2))
133 hidden_B_C = _mask_pad(hidden_B_C)
134 x_flat, B_flat, C_flat = hidden_B_C.split(
135 [intermediate, n_groups * state, n_groups * state], dim=-1
136 )
138 x = x_flat.reshape(batch, seq_len, num_heads, head_dim).float()
139 heads_per_group = num_heads // n_groups
140 B = (
141 B_flat.reshape(batch, seq_len, n_groups, state)
142 .repeat_interleave(heads_per_group, dim=2)
143 .float()
144 )
145 C = (
146 C_flat.reshape(batch, seq_len, n_groups, state)
147 .repeat_interleave(heads_per_group, dim=2)
148 .float()
149 )
151 dt = torch.nn.functional.softplus(dt.float() + self.dt_bias.float())
152 dt = torch.clamp(dt, float(oc.time_step_limit[0]), float(oc.time_step_limit[1]))
153 A = -torch.exp(self.A_log.float()) # [num_heads]
155 # 3. Eager recurrence with intervention hooks.
156 writes = dt[:, :, :, None, None] * x[:, :, :, :, None] * B[:, :, :, None, :]
157 writes = self.hook_ssm_write(writes) # [batch, seq, heads, head_dim, state]
159 ssm_state = torch.zeros(
160 batch, num_heads, head_dim, state, dtype=writes.dtype, device=writes.device
161 )
162 states = []
163 for t in range(seq_len):
164 decay = torch.exp(dt[:, t, :] * A[None, :]) # [batch, heads]
165 ssm_state = decay[:, :, None, None] * ssm_state + writes[:, t]
166 states.append(ssm_state)
167 state_traj = torch.stack(states, dim=1) # [batch, seq, heads, head_dim, state]
168 state_traj = self.hook_ssm_state(state_traj)
170 y = torch.einsum("bthn,bthdn->bthd", C, state_traj)
171 y = y + self.D.float()[None, None, :, None] * x
172 y = y.reshape(batch, seq_len, intermediate)
174 # 4. Gated norm + output projection — reuse the HF submodule bridges.
175 scan_output = oc.norm(y.to(hidden_states.dtype), gate)
176 contextualized: torch.Tensor = oc.out_proj(scan_output)
177 return contextualized
179 def compute_effective_attention(
180 self,
181 cache: ActivationCache,
182 layer_idx: int,
183 include_dt_scaling: bool = False,
184 ) -> torch.Tensor:
185 """Materialize Mamba-2's effective attention matrix M = L ⊙ (C B^T).
187 Via State Space Duality (SSD), Mamba-2's SSM is equivalent to causal
188 attention with a per-step per-head learned decay — see "The Hidden
189 Attention of Mamba" (Ali et al., ACL 2025). Extracts B, C from
190 ``conv1d.hook_out`` (post conv + SiLU) and dt from ``in_proj.hook_out``,
191 then reads ``A_log`` and ``dt_bias`` via ``__getattr__`` delegation.
193 Args:
194 cache: ActivationCache from ``run_with_cache`` containing the
195 in_proj and conv1d hooks for this layer.
196 layer_idx: Block index for this mixer. Required because submodule
197 bridges don't know their own position in the block list.
198 include_dt_scaling: False (default) returns the attention-like
199 form M_att = L ⊙ (C B^T). True multiplies each column j by
200 dt[j], giving the strict reconstruction form that satisfies
201 ``y[i] = sum_j M[i,j] * x[j] + D * x[i]``.
203 Returns:
204 Tensor of shape ``[batch, num_heads, seq_len, seq_len]`` with the
205 upper triangle (j > i) zeroed.
207 Cost is O(batch · num_heads · seq_len²); use on short sequences (≤2k).
208 """
209 terms = self._ssd_terms(cache, layer_idx)
211 # CB[b, h, i, j] = <C[b, i, h], B[b, j, h]>
212 CB = torch.einsum("bihs,bjhs->bhij", terms.C, terms.B)
213 M = terms.L * CB # [batch, num_heads, seq, seq]
215 if include_dt_scaling:
216 # Multiply column j by dt[j, h] to absorb the B discretization
217 M = M * terms.dt.permute(0, 2, 1)[:, :, None, :]
219 return M
221 def compute_ssm_state(
222 self,
223 cache: ActivationCache,
224 layer_idx: int,
225 time_step: Optional[int] = None,
226 ) -> torch.Tensor:
227 """Reconstruct the recurrent SSM state ``S`` from cached hook values.
229 Mamba-2's recurrence is ``S_t = dA_t · S_{t-1} + dt_t · (x_t ⊗ B_t)`` with
230 ``dA_t[h] = exp(dt_t[h] · A[h])``, so post-hoc::
232 S_t[h, p, n] = sum_{j<=t} L[t, j] · dt_j[h] · x_j[h, p] · B_j[h, n]
234 where ``L`` is the same causal decay matrix used by
235 ``compute_effective_attention``. Read-only: no ``forward()`` re-run.
236 Verify with ``y_t = C_t · S_t + D · x_t == inner_norm.hook_in``.
238 On padded batches the cached hooks are unmasked, so ``S`` is exact only at
239 non-pad positions; pad-position state is out of contract.
241 Args:
242 cache: ActivationCache from ``run_with_cache`` with this layer's
243 in_proj and conv1d hooks.
244 layer_idx: Block index for this mixer.
245 time_step: If given, return only ``S`` at that position
246 (``[batch, num_heads, head_dim, state]``); avoids materializing
247 the full tensor. None returns every step.
249 Returns:
250 ``[batch, num_heads, seq, head_dim, state]`` for all steps, or
251 ``[batch, num_heads, head_dim, state]`` for a single ``time_step``.
253 Memory is O(batch · num_heads · seq · head_dim · state) for the full
254 tensor; pass ``time_step`` (or short sequences) when that is too large.
255 """
256 terms = self._ssd_terms(cache, layer_idx)
257 dtx = terms.dt[..., None] * terms.x # [batch, seq, num_heads, head_dim]
258 if time_step is not None:
259 # S_t = sum_{j<=t} L[t, j] · (dt_j x_j) ⊗ B_j
260 return torch.einsum("bhj,bjhp,bjhn->bhpn", terms.L[:, :, time_step, :], dtx, terms.B)
261 return torch.einsum("bhtj,bjhp,bjhn->bhtpn", terms.L, dtx, terms.B)
263 def _ssd_terms(self, cache: ActivationCache, layer_idx: int) -> _SSDTerms:
264 """Reconstruct the shared SSD intermediates (dt, decay L, per-head x/B/C).
266 Mirrors HF Mamba2Mixer's discretization from the cached in_proj/conv1d
267 outputs. Dims come from the wrapped HF mixer, not cfg: on a hybrid the
268 shared cfg holds the *attention* dims (cfg.n_heads etc.), while the HF
269 mixer always carries the true Mamba dims (as do A_log/dt_bias, already
270 read off the module via __getattr__). cfg is the fallback.
271 """
272 if self.config is None: 272 ↛ 273line 272 didn't jump to line 273 because the condition on line 272 was never true
273 raise RuntimeError("SSM2MixerBridge.config must be set")
275 in_proj_key = f"blocks.{layer_idx}.mixer.in_proj.hook_out"
276 conv1d_key = f"blocks.{layer_idx}.mixer.conv1d.hook_out"
277 if in_proj_key not in cache or conv1d_key not in cache:
278 raise RuntimeError(
279 f"SSD reconstruction needs {in_proj_key!r} and {conv1d_key!r} in "
280 "cache. Run `run_with_cache()` on the bridge before calling this."
281 )
283 cfg = self.config
284 oc = self.original_component
286 def _mamba_dim(module_attr: str, cfg_attr: str, default: Any) -> Any:
287 if oc is not None: 287 ↛ 291line 287 didn't jump to line 291 because the condition on line 287 was always true
288 val = getattr(oc, module_attr, None)
289 if val is not None: 289 ↛ 291line 289 didn't jump to line 291 because the condition on line 289 was always true
290 return val
291 return getattr(cfg, cfg_attr, default)
293 num_heads = int(_mamba_dim("num_heads", "n_heads", 0))
294 head_dim = int(_mamba_dim("head_dim", "d_head", 0))
295 intermediate_size = int(
296 _mamba_dim("intermediate_size", "intermediate_size", num_heads * head_dim)
297 )
298 state_size = int(_mamba_dim("ssm_state_size", "state_size", 128))
299 n_groups = int(_mamba_dim("n_groups", "n_groups", 1))
300 time_step_limit = _mamba_dim("time_step_limit", "time_step_limit", (0.0, float("inf")))
302 in_proj_out = cache[in_proj_key].float() # [batch, seq, proj_size]
303 conv1d_out = cache[conv1d_key].float() # [batch, conv_dim, seq + kernel - 1]
304 batch_size, seq_len = in_proj_out.shape[0], in_proj_out.shape[1]
306 dt = torch.nn.functional.softplus(in_proj_out[..., -num_heads:] + self.dt_bias.float())
307 dt = torch.clamp(dt, float(time_step_limit[0]), float(time_step_limit[1]))
309 # x, B, C from conv1d output (trim to seq, SiLU, channel-last split)
310 conv_activated = torch.nn.functional.silu(conv1d_out[..., :seq_len]).transpose(1, 2)
311 x_flat, B_flat, C_flat = conv_activated.split(
312 [intermediate_size, n_groups * state_size, n_groups * state_size], dim=-1
313 )
314 x = x_flat.view(batch_size, seq_len, num_heads, head_dim)
315 # GQA-style: each of n_groups B/C pairs covers n_heads // n_groups heads
316 heads_per_group = num_heads // n_groups
317 B = B_flat.view(batch_size, seq_len, n_groups, state_size).repeat_interleave(
318 heads_per_group, dim=2
319 )
320 C = C_flat.view(batch_size, seq_len, n_groups, state_size).repeat_interleave(
321 heads_per_group, dim=2
322 )
324 # L[i, j] = exp(sum_{k=j+1}^{i} A[h] · dt[k, h]) for i >= j, else 0.
325 # exp(cumsum[i] - cumsum[j]); cumsum[j] includes dt[j], so the sum is k>j.
326 A = -torch.exp(self.A_log.float()) # [num_heads]
327 cs = torch.cumsum(dt * A[None, None, :], dim=1).permute(0, 2, 1) # [batch, num_heads, seq]
328 L_log = cs[:, :, :, None] - cs[:, :, None, :]
329 causal_mask = torch.tril(
330 torch.ones(seq_len, seq_len, dtype=torch.bool, device=L_log.device)
331 )
332 L = torch.where(causal_mask[None, None], torch.exp(L_log), torch.zeros_like(L_log))
334 return _SSDTerms(dt=dt, L=L, x=x, B=B, C=C)