Add Jamba attention+Mamba interleave adapter (JambaForCausalLM)

[udapy]

Description

Adds TransformerBridge support for JambaForCausalLM (AI21 Jamba family).

• Jamba is a layer-interleaved hybrid: each layer is either GQA attention (no RoPE) or a Mamba-1 mixer, then a dense SwiGLU or sparse MoE FFN under a dual-norm residual (input_layernorm → mixer → pre_ff_layernorm → FFN).
• Attention and Mamba are mapped as independently hookable optional streams so researchers can ablate one path per layer and compare what attention adds over SSM.

Fixes #1468

// summary

Aspect	Choice
Block container	BlockBridge — input_layernorm + pre_ff_layernorm, real hook_resid_mid
Layer schedule	Optional attn or mixer per layers_block_type (attention / mamba)
SSM mixer	SSMMixerBridge (Mamba-1) under canonical .mixer (HF .mamba) so find_ssm_mixer / compute_ssm_* / eager_scan resolve
Jamba-only LNs	dt_layernorm / b_layernorm / c_layernorm mapped; SSMMixerBridge applies them on reconstruction + eager_scan (stock Mamba-1 unchanged)
Attention	GQA, no RoPE → AttentionBridge(maintain_native_attention=True)
FFN	GatedMLPBridge if num_experts≤1; MoEBridge (optional dense+router children) if num_experts>1
Generation	Standard past_key_values / DynamicCache — is_stateful=False
Weight folding	supports_fold_ln=False, weight_processing_conversions={} (heterogeneous layers + native attn)
verify_models	applicable_phases=[1,2,3,4]

Key design decision: no is_stateful

• With is_stateful=True, the bridge’s stateful loop uses the pure-Mamba cache_params convention.
• Jamba carries attention KV and Mamba conv/SSM state in one unified past_key_values / DynamicCache.
• Leaving is_stateful=False keeps greedy generation on the HF-compatible path.

Key design decision: Jamba-aware SSMMixerBridge

• Stock Mamba-1 reads B/C from x_proj; Jamba norms dt/B/C after x_proj before dt_proj.
• Without that, compute_ssm_state / compute_ssm_effective_attention / eager_scan are shaped correctly but numerically wrong for Jamba.
• Gate is opt-in via getattr / cache keys: only mixers that expose those LNs take the Jamba path; pure MambaForCausalLM stays on the old path (validated; mamba integration tests pass).

File Changes (14 Files)

File	Change
transformer_lens/model_bridge/supported_architectures/jamba.py	New adapter
tests/unit/.../test_jamba_adapter.py	Structural unit tests
tests/integration/.../test_jamba_adapter.py	Parity, ablation, SSM reconstruction, eager_scan, greedy generate
supported_architectures/init.py	Registration 1/4
factories/architecture_adapter_factory.py	Registration 2/4
tools/model_registry/init.py	Registration 3/4 (HF_SUPPORTED_ARCHITECTURES, CANONICAL_AUTHORS_BY_ARCH)
tools/model_registry/generate_report.py	Registration 4/4
sources/transformers.py + sources/_bridge_builder.py	Jamba schedule/Mamba attrs in both _HF_PASSTHROUGH_ATTRS lists
generalized_components/ssm_mixer.py	Selective-param LN support for reconstruction + eager_scan
docs/source/content/ssm_interpretability.md	Document Jamba + selective-param LNs
tools/model_registry/AGENTS.md	List Jamba among SSM hybrids
tools/model_registry/data/supported_models.json	Jamba-tiny-random + AI21-Jamba-Reasoning-3B entries
tools/model_registry/data/verification_history.json	Verification audit trail

Local Verification

Model	Notes
ai21labs/Jamba-tiny-random	status=1; P1/P2/P3=100 (P4 low expected for random weights)
ai21labs/AI21-Jamba-Reasoning-3B	status=1; P1=100, P2=100 (core + hooks/grads); text quality high

• P3 on Reasoning-3B is below 100 only from weight_magnitudes after limited weight processing (supports_fold_ln=False); forward/hook/grad checks in P2 are 100. P4 is a GPT-2 perplexity quality metric, not a logit-parity check.

• Integration coverage: HF logit parity (max_diff == 0 on programmatic configs), independent attn vs Mamba ablation, find_ssm_mixer / ssm_layers, post-LN state reconstruction, eager_scan vs fused, greedy generate parity.

• Local suite: 46 Jamba unit+integration tests passed; stock Mamba adapter integration (54) still passes after ssm_mixer change.

Expected warnings (not regressions)

• On Mamba-only layers, hook_attn_* aliases do not resolve (optional attn absent) — same class of warning as other hybrids with optional attention.

CI Notes

Integration tests are not marked slow — they use programmatic configs with no Hub download (same as GraniteMoeHybrid) and run in normal CI via make integration-test.

Not in scope

• Compatibility-mode LN folding for Jamba (supports_fold_ln=False by design)
• MoE-routing interp beyond passthrough hooks on sparse feed_forward layers
• 52B+ Jamba MoE checkpoints as verify targets (tiny-random covers MoE layout; Reasoning-3B is dense num_experts=1)

Type of change

• New feature (non-breaking change which adds functionality)
• This change requires a documentation update

Checklist:

• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes
• I have not rewritten tests relating to key interfaces which would affect backward compatibility

Related

• Closes [Proposal] Add Jamba attention+Mamba interleave adapter (JambaForCausalLM) #1468
• Uses SSM interp surface from SSM Interp Support #1481 (SSMMixerBridge)
• Adjacent hybrid patterns: Falcon-H1 (Add Falcon-H1 parallel hybrid architecture adapter (FalconH1ForCausalLM) #1480), Zamba2 (feat: add Zamba2ForCausalLM architecture adapter #1486), NemotronH / GraniteMoeHybrid
• Paper: https://arxiv.org/abs/2403.19887
• HF docs: https://huggingface.co/docs/transformers/model_doc/jamba

[jlarson4]

This looks great, thanks for putting it together @udapy!

[udapy]

This looks great, thanks for putting it together @udapy!

Thank you @jlarson4