Coverage for transformer_lens/model_bridge/generalized_components/ssm_protocol.py: 100%

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

2 

3Defines the structural contract (``SSMMixerProtocol``) that ``SSM2MixerBridge`` 

4(Mamba-2), ``SSMMixerBridge`` (Mamba-1), and ``GatedDeltaNetBridge`` all satisfy, 

5plus a lookup that finds a block's SSM mixer regardless of which variant slot it 

6occupies (``.mixer`` for Mamba, ``.linear_attn`` for gated-delta-net). A Protocol 

7(not a base class) keeps each family's own hook set — only the discovery surface 

8is shared. ``SSMStateHookMixin`` is the one place the canonical eager-scan 

9intervention hooks (``hook_ssm_state`` + the ``eager_scan`` opt-in) are defined, so 

10the three families share a single definition rather than repeating it ad hoc. 

11""" 

12from __future__ import annotations 

13 

14from typing import TYPE_CHECKING, Any, Optional, Protocol, runtime_checkable 

15 

16import torch 

17import torch.nn as nn 

18 

19from transformer_lens.hook_points import HookPoint 

20from transformer_lens.model_bridge.generalized_components.block import ( 

21 VARIANT_SUBMODULE_NAMES, 

22) 

23 

24if TYPE_CHECKING: 

25 from transformer_lens.ActivationCache import ActivationCache 

26 

27 

28@runtime_checkable 

29class SSMMixerProtocol(Protocol): 

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

31 

32 Only ``compute_effective_attention(cache, layer_idx)`` is required; families 

33 add their own optional keyword options (e.g. ``include_dt_scaling``, 

34 ``per_state_coord``) which callers discover by signature introspection. 

35 """ 

36 

37 def compute_effective_attention(self, cache: "ActivationCache", layer_idx: int) -> torch.Tensor: 

38 ... 

39 

40 

41class SSMStateHookMixin: 

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

43 

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

45 

46 - ``hook_ssm_state`` — post-scan state trajectory ``S_t`` (created here for all); 

47 patching it changes only the same-position readout, not the recurrence. 

48 - ``hook_ssm_write`` — per-step write influence: a real HookPoint for the input- 

49 linear families (Mamba-1/2 ``dt·(x⊗B)``, added by the subclass), an alias onto 

50 ``hook_beta`` for the state-dependent delta rule; patching it re-runs the scan. 

51 

52 ``eager_scan`` (opt-in, prefill only) swaps HF's fused kernel for a Python scan so 

53 these hooks fire; default False leaves the cached path bit-for-bit untouched. Must 

54 precede ``GeneralizedComponent`` in the bases so ``super().__init__`` reaches it. 

55 """ 

56 

57 eager_scan: bool = False 

58 

59 def __init__(self, *args: Any, **kwargs: Any) -> None: 

60 super().__init__(*args, **kwargs) 

61 self.hook_ssm_state = HookPoint() 

62 

63 

64# Children present on *every* SSM2MixerBridge regardless of whether it wraps a 

65# real Mamba layer — universal I/O hooks, the wrapped module, and the eager-scan 

66# analysis hooks. None of these signal that the mixer is realized. 

67_PASSTHROUGH_CHILDREN = frozenset( 

68 {"hook_in", "hook_out", "_original_component", "hook_ssm_write", "hook_ssm_state"} 

69) 

70 

71 

72def _is_realized_ssm_mixer(mixer: object) -> bool: 

73 """True unless ``mixer`` is a no-op passthrough wrapper. 

74 

75 A hybrid like NemotronH wires a single ``SSM2MixerBridge`` ``.mixer`` slot on 

76 *every* layer; on attention / MLP / MoE layers its optional projection 

77 submodules are skipped, leaving only the universal hooks. A realized mixer 

78 always has something more — projection submodules (Mamba-1/2) or interior 

79 hooks like ``hook_q`` / ``hook_log_decay`` (gated-delta-net). This structural 

80 check needs no ``cfg.layers_block_type``. 

81 """ 

82 return any(name not in _PASSTHROUGH_CHILDREN for name in getattr(mixer, "_modules", {})) 

83 

84 

85def find_ssm_mixer(block: nn.Module) -> Optional[SSMMixerProtocol]: 

86 """Return the block's SSM mixer submodule, or None if it has none. 

87 

88 Scans the layer-type variant slots (``.mixer`` / ``.linear_attn`` / ``.mamba`` 

89 / ``.ssm``) and returns the first that conforms to ``SSMMixerProtocol`` and is 

90 a realized SSM mixer (not a passthrough wrapper). An attention layer (only 

91 ``.attn``), or a hybrid's passthrough ``.mixer`` on a non-SSM layer, returns 

92 None. 

93 """ 

94 modules = getattr(block, "_modules", {}) 

95 for name in VARIANT_SUBMODULE_NAMES: 

96 sub = modules.get(name) 

97 if isinstance(sub, SSMMixerProtocol) and _is_realized_ssm_mixer(sub): 

98 return sub 

99 return None