transformer_lens.model_bridge.supported_architectures.mamba2 module¶

Architecture adapter for HF’s Mamba2ForCausalLM, plus the effective attention helper.

class transformer_lens.model_bridge.supported_architectures.mamba2.Mamba2ArchitectureAdapter(cfg: Any)¶

Bases: ArchitectureAdapter

Wraps HF’s Mamba2ForCausalLM.

Differs from Mamba-1 at the mixer level: fused in_proj (no x_proj/dt_proj), two-input inner norm, multi-head structure with num_heads/head_dim/ n_groups, and an [num_heads]-shaped dt_bias. Shares SSMBlockBridge, DepthwiseConv1DBridge, and the stateful generation loop with Mamba-1.

applicable_phases: list[int] = [1, 2, 3, 4]¶

component_mapping: ComponentMapping | None¶

create_stateful_cache(hf_model: Any, batch_size: int, device: Any, dtype: dtype) → Any¶: Build a cache for the stateful generation loop.

uses_split_attention: bool¶

weight_processing_conversions: Dict[str, ParamProcessingConversion | str] | None¶

transformer_lens.model_bridge.supported_architectures.mamba2.compute_effective_attention(bridge: TransformerBridge, cache: ActivationCache, layer: int | None = None, include_dt_scaling: bool = False) → Tensor | Dict[int, Tensor]¶: Mamba-2 effective attention for one or all layers.

Deprecated since version Use: the family-agnostic cache.compute_ssm_effective_attention(layer=...) instead. This thin wrapper delegates to it and ignores bridge (the cache already knows its model).

transformer_lens.model_bridge.supported_architectures.mamba2.compute_ssm_state(bridge: TransformerBridge, cache: ActivationCache, layer: int | None = None, time_step: int | None = None) → Tensor | Dict[int, Tensor]¶: Reconstruct the recurrent SSM state S for one or all Mamba-2 layers.

Deprecated since version Use: cache.compute_ssm_state(layer=..., time_step=...) instead. This thin wrapper delegates to it and ignores bridge.