transformer_lens.model_bridge.generalized_components.ssm_protocol module

Family-agnostic discovery + shared state-mutation surface for SSM/recurrent mixers.

Defines the structural contract (SSMMixerProtocol) that SSM2MixerBridge (Mamba-2), SSMMixerBridge (Mamba-1), and GatedDeltaNetBridge all satisfy, plus a lookup that finds a block’s SSM mixer regardless of which variant slot it occupies (.mixer for Mamba, .linear_attn for gated-delta-net). A Protocol (not a base class) keeps each family’s own hook set — only the discovery surface is shared. SSMStateHookMixin is the one place the canonical eager-scan intervention hooks (hook_ssm_state + the eager_scan opt-in) are defined, so the three families share a single definition rather than repeating it ad hoc.

class transformer_lens.model_bridge.generalized_components.ssm_protocol.SSMMixerProtocol(*args, **kwargs)

Bases: Protocol

An SSM/recurrent mixer that can materialize an effective-attention matrix.

Only compute_effective_attention(cache, layer_idx) is required; families add their own optional keyword options (e.g. include_dt_scaling, per_state_coord) which callers discover by signature introspection.

compute_effective_attention(cache: ActivationCache, layer_idx: int) → torch.Tensor
class transformer_lens.model_bridge.generalized_components.ssm_protocol.SSMStateHookMixin(*args: Any, **kwargs: Any)

Bases: object

Canonical state-mutation hooks shared by the SSM/recurrent mixers.

One definition point so Mamba-1/Mamba-2/gated-delta-net expose the same names:

  • hook_ssm_state — post-scan state trajectory S_t (created here for all); patching it changes only the same-position readout, not the recurrence.

  • hook_ssm_write — per-step write influence: a real HookPoint for the input- linear families (Mamba-1/2 dt·(x⊗B), added by the subclass), an alias onto hook_beta for the state-dependent delta rule; patching it re-runs the scan.

eager_scan (opt-in, prefill only) swaps HF’s fused kernel for a Python scan so these hooks fire; default False leaves the cached path bit-for-bit untouched. Must precede GeneralizedComponent in the bases so super().__init__ reaches it.

eager_scan: bool = False
transformer_lens.model_bridge.generalized_components.ssm_protocol.find_ssm_mixer(block: Module) → SSMMixerProtocol | None

Return the block’s SSM mixer submodule, or None if it has none.

Scans the layer-type variant slots (.mixer / .linear_attn / .mamba / .ssm) and returns the first that conforms to SSMMixerProtocol and is a realized SSM mixer (not a passthrough wrapper). An attention layer (only .attn), or a hybrid’s passthrough .mixer on a non-SSM layer, returns None.