transformer_lens.model_bridge.supported_architectures.bamba module

Bamba (BambaForCausalLM) adapter: Jamba-lineage hybrid alternating Mamba-2 and llama-style GQA attention mixers per config.layers_block_type.

class transformer_lens.model_bridge.supported_architectures.bamba.BambaArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for BambaForCausalLM models.

Both mixers are mapped optional — each present only on its layer type. The Mamba-2 mixer is wired under the canonical .mixer slot (HF path .mamba) so SSM analyses reach it as on GraniteMoeHybrid / NemotronH.

__init__(cfg: Any) → None

Initialize the Bamba architecture adapter.

create_stateful_cache(hf_model: Any, batch_size: int, device: Any, dtype: Any) → Any

Unified DynamicCache carrying KV entries and SSM conv/recurrent state.