Skip to content

[Proposal] Add Jamba attention+Mamba interleave adapter (JambaForCausalLM) #1468

Description

@jlarson4

Proposal

Add a TransformerBridge adapter for JambaForCausalLM (AI21 Jamba), a decoder that interleaves vanilla attention and Mamba (SSM) layers, with optional MoE.

Motivation

Jamba is the canonical openly-released decoder that interleaves attention and Mamba layer-by-layer, letting researchers study, layer by layer, what attention adds over a state-space layer in a single model. The adjacent Mamba/hybrid interp literature is active (hidden-attention views of Mamba, Mamba knockout for factual-recall flow, Mamba-Transformer hybrid analysis), and the new 3B reasoning checkpoints make Jamba laptop-runnable for the first time.

Honest caveats for scoping: The attention+Mamba interleave surface partially overlaps the already-supported NemotronHForCausalLM and GraniteMoeHybridForCausalLM hybrids. Jamba's distinction is the specific interleave pattern and its status as the original public exemplar. Secondarily, there is no Jamba-specific interp paper yet (the citations are general Mamba/hybrid work). Finally, the laptop-runnable 3B checkpoints have num_experts=1, so the MoE-routing surface only appears on the 52B+ models. Worth adding for the interleave surface, but lower-priority than the non-overlapping novel surfaces.

Gap scan (2026-06-25): 21 models, ~128K downloads.

Pitch

Implement the Jamba block schedule (interleaved attention and Mamba mixer layers; MoE MLP where present), exposing the attention and Mamba streams as separate hookable paths so the two can be ablated and compared per layer.

  • Claude Code users can scaffold with /add-model-support ai21labs/AI21-Jamba-Reasoning-3B.
  • Register at the four sites listed in contributing.md.
  • Verify smallest-first: ai21labs/Jamba-tiny-random, then ai21labs/AI21-Jamba-Reasoning-3B.

Additional context

Checklist

  • I have checked that there is no similar issue in the repo (required)

Metadata

Metadata

Assignees

Labels

TransformerBridgeBug specific to the new TransformerBridge systemcomplexity-moderateModerately complicated issues for people who have intermediate experience with the codehelp wantedExtra attention is needednew-architectureThis card involves adding a new architecture .

Type

No type

Fields

No fields configured for issues without a type.

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions