fix: decompose Falcon-H1's parallel Mamba output and inline multipliers

tests/integration/model_bridge/test_ssm_hybrid_resid_decomposition.py

@@ -2,7 +2,8 @@
 
 Blocks without attention (gated-delta-net, Mamba) expose their sequence-mixer
 output instead of ``hook_attn_out``; blocks without an MLP (Mamba, NemotronH)
-write only through their mixer. Every decomposition must still sum to the
+write only through their mixer. Parallel hybrids (Falcon-H1) add a Mamba output
+beside attention in every block. Every decomposition must still sum to the
 residual stream it decomposes.
 """
 from unittest.mock import MagicMock
@@ -169,6 +170,51 @@ def _granite_dense():
     return _boot(AutoModelForCausalLM.from_config(c).eval(), "GraniteMoeHybridForCausalLM")
 
 
+def _falcon_h1():
+    from transformers import AutoModelForCausalLM
+    from transformers.models.falcon_h1 import FalconH1Config
+
+    torch.manual_seed(0)
+    c = FalconH1Config(
+        hidden_size=64,
+        intermediate_size=128,
+        num_hidden_layers=3,
+        num_attention_heads=4,
+        num_key_value_heads=2,
+        head_dim=16,
+        vocab_size=64,
+        max_position_embeddings=64,
+        mamba_d_ssm=128,
+        mamba_n_heads=4,
+        mamba_d_head=32,
+        mamba_n_groups=1,
+        mamba_d_state=16,
+        mamba_d_conv=4,
+        mamba_chunk_size=8,
+        attention_bias=True,
+        mlp_bias=True,
+        mamba_proj_bias=True,
+        projectors_bias=True,
+        # Released checkpoints scale every branch; 1.0 would hide a missing multiplier.
+        embedding_multiplier=2.0,
+        attention_in_multiplier=1.1,
+        attention_out_multiplier=0.7,
+        key_multiplier=0.8,
+        ssm_in_multiplier=1.4,
+        ssm_out_multiplier=1.9,
+        ssm_multipliers=[0.5, 1.5, 0.8, 1.2, 0.9],
+        mlp_multipliers=[1.3, 0.6],
+        pad_token_id=0,
+    )
+    c.architectures = ["FalconH1ForCausalLM"]
+    model = AutoModelForCausalLM.from_config(c).eval()
+    with torch.no_grad():
+        for param in model.parameters():
+            if param.ndim == 1:
+                param.normal_(std=0.3)  # nonzero biases and norm weights
+    return _boot(model, "FalconH1ForCausalLM")
+
+
 def _zamba2():
     """Zamba2 with hybrid layers: the shared transformer feeds the Mamba-2 input of layers 1
     and 3 but does not write to the residual stream."""
@@ -250,6 +296,7 @@ def _lfm2():
 
 
 _FACTORIES = {
+    "falcon_h1": _falcon_h1,
     "granite_dense": _granite_dense,
     "lfm2": _lfm2,
     "mamba2": _mamba2,
@@ -307,6 +354,7 @@ def test_labels_name_the_mixer(self, bridge_and_cache):
         name, bridge, cache = bridge_and_cache
         _, labels = cache.decompose_resid(return_labels=True)
         expected = {
+            "falcon_h1": ["0_attn_out", "0_mamba_out", "0_mlp_out", "1_attn_out", "1_mamba_out"],
             "qwen3_5": ["0_linear_attn_out", "0_mlp_out", "3_attn_out", "3_mlp_out"],
             "granite_dense": ["0_mixer_out", "0_mlp_out", "1_attn_out", "1_mlp_out"],
             "lfm2": ["0_conv_out", "0_mlp_out", "1_attn_out", "1_mlp_out", "2_conv_out"],
@@ -346,14 +394,29 @@ def test_mlp_type_mixers_fill_the_mlp_slot(self, bridge_and_cache):
 
     def test_mlp_input_sums_to_resid_mid(self, bridge_and_cache):
         name, bridge, cache = bridge_and_cache
-        if name not in ("qwen3_5", "recurrent_gemma"):
-            pytest.skip("only Qwen3.5 and RecurrentGemma blocks have a separate MLP input")
-        for layer in (0, 2, 3):
+        if name not in ("qwen3_5", "falcon_h1", "recurrent_gemma"):
+            pytest.skip("Qwen3.5, RecurrentGemma and Falcon-H1 blocks have a separate MLP input")
+        for layer in (0, 3) if name == "qwen3_5" else (0, 2):
             comps = cache.decompose_resid(layer=layer, mlp_input=True)
             torch.testing.assert_close(
                 comps.sum(0), cache[f"blocks.{layer}.hook_resid_mid"], atol=ATOL, rtol=0
             )
 
+    def test_parallel_branches_sum_to_their_residual_add(self, bridge_and_cache):
+        name, bridge, cache = bridge_and_cache
+        if name != "falcon_h1":
+            pytest.skip("only Falcon-H1 runs a mixer beside attention")
+        ssm_out_multiplier = bridge.cfg.ssm_out_multiplier
+        for layer in range(bridge.cfg.n_layers):
+            added = (
+                cache[f"blocks.{layer}.hook_resid_mid"] - cache[f"blocks.{layer}.hook_resid_pre"]
+            )
+            branches = (
+                cache[f"blocks.{layer}.hook_attn_out"]
+                + cache[f"blocks.{layer}.mamba.hook_out"] * ssm_out_multiplier
+            )
+            torch.testing.assert_close(branches, added, atol=ATOL, rtol=0)
+
     def test_logit_attrs_runs(self, bridge_and_cache):
         _, bridge, cache = bridge_and_cache
         attrs = cache.logit_attrs(cache.decompose_resid(), tokens=torch.tensor([3, 3]))
@@ -379,7 +442,9 @@ def test_head_results_sum_to_attn_out(self, bridge_and_cache):
             heads = cache[f"blocks.{layer}.attn.hook_result"].sum(-2)
             b_O = bridge.blocks[layer].attn.b_O
             if b_O is not None:
-                heads = heads + b_O
+                heads = heads + b_O * getattr(
+                    bridge.blocks[layer], "residual_contribution_scale", 1.0
+                )
             torch.testing.assert_close(
                 heads, cache[f"blocks.{layer}.hook_attn_out"], atol=ATOL, rtol=0
             )

transformer_lens/ActivationCache.py

@@ -402,20 +402,40 @@ def _has_attn(self, layer: int) -> bool:
         )
 
     def _mixer_out_names(self, layer: int) -> List[str]:
-        """Variant slots (``linear_attn``, ``mixer``, …) of a block without attention.
+        """Variant slots (``linear_attn``, ``mixer``, …) that write outside ``hook_attn_out``.
 
         Hybrid and SSM blocks write their sequence-mixer output in place of
-        ``hook_attn_out``; empty for attention blocks.
+        ``hook_attn_out``. Parallel hybrids (Falcon-H1) run a mixer beside attention
+        and add both outputs to the residual stream; that mixer is listed too.
         """
         from transformer_lens.model_bridge.generalized_components.block import (
             VARIANT_SUBMODULE_NAMES,
         )
 
-        if self._has_attn(layer):
-            return []
         modules = self._block_modules(layer) or {}
+        if self._has_attn(layer):
+            if "attn" not in modules:
+                return []
+            return [name for name in VARIANT_SUBMODULE_NAMES if name != "attn" and name in modules]
         return [name for name in VARIANT_SUBMODULE_NAMES if name in modules]
 
+    @staticmethod
+    def _mlp_contribution_scale(block: Any) -> float:
+        """Multiplier HF applies to the MLP output before the residual add."""
+        return getattr(
+            block, "mlp_contribution_scale", getattr(block, "residual_contribution_scale", 1.0)
+        )
+
+    @classmethod
+    def _neuron_contribution_scale(cls, block: Any) -> float:
+        """Multiplier from ``post @ W_out + b_out`` to the residual contribution.
+
+        Adds any multiplier inside the MLP after ``W_out`` (Falcon-H1's down multiplier),
+        which ``hook_mlp_out`` already includes but the neuron expansion does not.
+        """
+        mlp = block._modules.get("mlp") if isinstance(block._modules, dict) else None
+        return cls._mlp_contribution_scale(block) * getattr(mlp, "_neuron_contribution_scale", 1.0)
+
     def _has_mlp(self, layer: int) -> bool:
         """Whether block ``layer`` writes an MLP output separate from its mixer."""
         from transformer_lens.model_bridge.generalized_components.block import (
@@ -470,24 +490,28 @@ def _mixer_stack(
         """Residual contributions of a non-attention block's mixer, e.g. ``0_linear_attn_out``.
 
         Scaled by the block's ``residual_contribution_scale`` (Granite's ``residual_multiplier``
-        is applied after ``hook_out``), matching the head and neuron paths. ``subtract_bias``
-        drops the mixer's output bias first, for decompositions that carry it in their own
-        ``bias`` row.
+        is applied after ``hook_out``), matching the head and neuron paths, or by the mixer's
+        own entry in ``mixer_contribution_scales`` (Falcon-H1's ``ssm_out_multiplier``).
+        ``subtract_bias`` drops the mixer's output bias first, for decompositions that carry it
+        in their own ``bias`` row.
         """
         components: List[torch.Tensor] = []
         labels: List[str] = []
         names = self._mixer_out_names(layer)
         block = self.model.blocks[layer]
-        scale = getattr(block, "residual_contribution_scale", 1.0)
-        # Same bias accumulated_bias counts for this block, so the two cancel exactly.
+        default_scale = getattr(block, "residual_contribution_scale", 1.0)
+        mixer_scales = getattr(block, "mixer_contribution_scales", {})
+        # Same bias accumulated_bias counts for this block, so the two cancel exactly. On a
+        # parallel hybrid that bias is the attention's, so the mixer keeps its own.
         bias = None
-        if subtract_bias and names:
+        if subtract_bias and names and not self._has_attn(layer):
             get_bias = getattr(self.model, "_get_block_variant_bias", None)
             bias = get_bias(block) if callable(get_bias) else None
         for i, name in enumerate(names):
             out = pos_slice.apply(self[f"blocks.{layer}.{name}.hook_out"], dim=-2)
             if i == 0 and isinstance(bias, torch.Tensor):
                 out = out - bias.to(out.device, out.dtype)
+            scale = mixer_scales.get(name, default_scale)
             if scale != 1.0:
                 out = out * scale
             components.append(out)
@@ -845,7 +869,8 @@ def decompose_resid(
                 attention components. Defaults to "all". On SSM / hybrid models, blocks without
                 attention contribute their sequence-mixer output (e.g. ``0_linear_attn_out``,
                 ``0_mixer_out``) in the attention slot, and blocks without an MLP contribute no
-                ``mlp_out``.
+                ``mlp_out``. Parallel hybrids (Falcon-H1) add their mixer output
+                (``0_mamba_out``) after ``0_attn_out``.
             apply_ln:
                 Whether to apply LayerNorm to the stack.
             pos_slice:
@@ -891,6 +916,8 @@ def add_attn_slot(l: int) -> None:
             if self._has_attn(l):
                 components_list.append(self[("attn_out", l)])
                 labels.append(f"{l}_attn_out")
+                # Parallel hybrids add a mixer output beside attention.
+                add_mixer(l)
             elif self._mixer_role(l) == "attn":
                 add_mixer(l)
 
@@ -1230,7 +1257,11 @@ def _post_norm_bias_correction(
             mapped = self._apply_post_norm_map(raw, scale, weight, centers)
             if norm_bias is not None:
                 mapped = mapped + norm_bias.to(device=scale.device)
-            delta = (mapped - raw) * getattr(block, "residual_contribution_scale", 1.0)
+            delta = (mapped - raw) * (
+                getattr(block, "residual_contribution_scale", 1.0)
+                if kind == "attn"
+                else self._neuron_contribution_scale(block)
+            )
             correction = delta if correction is None else correction + delta
         return correction
 
@@ -1408,7 +1439,7 @@ def get_neuron_results(
         neuron_acts = self[("post", layer, "mlp")]
         block = self.model.blocks[layer]
         W_out = block.mlp.W_out
-        contribution_scale = getattr(block, "residual_contribution_scale", 1.0)
+        contribution_scale = self._neuron_contribution_scale(block)
         if pos_slice is not None:
             # Note - order is important, as Slice.apply *may* collapse a dimension, so this ensures
             # that position dimension is -2 when we apply position slice
@@ -1521,7 +1552,7 @@ def _stack_neuron_results_apply_ln_projected(
                 continue
             block = self.model.blocks[l]
             W_out_l = block.mlp.W_out  # [d_mlp, d_model]
-            contribution_scale = getattr(block, "residual_contribution_scale", 1.0)
+            contribution_scale = self._neuron_contribution_scale(block)
             post = self._post_norm_map(l, "ln2_post", pos_slice)
             if post is not None:
                 post_scale, post_weight, post_centers, _ = post
@@ -1851,7 +1882,8 @@ def get_full_resid_decomposition(
            layers only on SSM / hybrid models)
         2. Sequence-mixer results of blocks without attention heads (SSM / hybrid
            models only; one row per mixer, e.g. ``0_linear_attn_out``, output bias
-           excluded since it is in the ``bias`` row). MLP-type passthrough mixers
+           excluded since it is in the ``bias`` row), and of mixers that run beside
+           attention (Falcon-H1's ``0_mamba_out``, bias included). MLP-type passthrough mixers
            (NemotronH) appear here only when ``expand_neurons=True``; otherwise they
            are part of the MLP rows below.
         3. Neuron / MLP results (only if ``cfg.attn_only=False`` and

transformer_lens/model_bridge/generalized_components/block.py

@@ -520,10 +520,17 @@ def __init__(
         residual_contribution_scale: float = 1.0,
         scaled_attn_submodule: str = "attn",
         scaled_mlp_submodule: Optional[str] = "mlp",
+        mlp_contribution_scale: Optional[float] = None,
+        mixer_contribution_scales: Optional[Dict[str, float]] = None,
     ):
         """scaled_mlp_submodule=None when no single submodule feeds the MLP-side
         add (GraniteMoeHybrid sums moe + shared_mlp inline) — hook_mlp_out then
-        stays absent rather than firing with a partial tensor."""
+        stays absent rather than firing with a partial tensor.
+
+        ``residual_contribution_scale`` scales the attention side, and the MLP side
+        unless ``mlp_contribution_scale`` is given. ``mixer_contribution_scales`` maps
+        a sequence mixer that runs beside attention (Falcon-H1's ``mamba``) to the
+        multiplier HF applies after its ``hook_out``."""
         super().__init__(
             name,
             config=config,
@@ -538,6 +545,16 @@ def __init__(
                 f"must be nonzero (the write path divides by it)."
             )
         self.residual_contribution_scale = scale
+        mlp_scale = scale if mlp_contribution_scale is None else float(mlp_contribution_scale)
+        if mlp_scale == 0.0:
+            raise ValueError(
+                f"ScaledResidualBlockBridge at '{name}': mlp_contribution_scale "
+                f"must be nonzero (the write path divides by it)."
+            )
+        self.mlp_contribution_scale = mlp_scale
+        self.mixer_contribution_scales = {
+            k: float(v) for k, v in (mixer_contribution_scales or {}).items()
+        }
         self.scaled_attn_submodule = scaled_attn_submodule
         self.scaled_mlp_submodule = scaled_mlp_submodule
         attn = self.submodules.get(scaled_attn_submodule) if self.submodules else None
@@ -585,10 +602,10 @@ def _maybe_wire_capture_hooks(self) -> None:
         if self._capture_hooks_wired:
             return
         super()._maybe_wire_capture_hooks()
-        targets = [(self.scaled_attn_submodule, "hook_attn_out")]
+        targets = [(self.scaled_attn_submodule, "hook_attn_out", self.residual_contribution_scale)]
         if self.scaled_mlp_submodule is not None:
-            targets.append((self.scaled_mlp_submodule, "hook_mlp_out"))
-        for sub_name, hook_name in targets:
+            targets.append((self.scaled_mlp_submodule, "hook_mlp_out", self.mlp_contribution_scale))
+        for sub_name, hook_name, scale in targets:
             hook_point = getattr(self, hook_name, None)
             sub = self.submodules.get(sub_name) if self.submodules else None
             if (
@@ -597,14 +614,15 @@ def _maybe_wire_capture_hooks(self) -> None:
                 or getattr(sub, "original_component", None) is None
             ):
                 continue
-            handle = sub.register_forward_hook(self._make_scaled_contribution_hook(hook_point))
+            handle = sub.register_forward_hook(
+                self._make_scaled_contribution_hook(hook_point, scale)
+            )
             self._capture_hook_handles.append(handle)
 
     def _make_scaled_contribution_hook(
-        self, hook_point: HookPoint
+        self, hook_point: HookPoint, scale: float
     ) -> Callable[[torch.nn.Module, tuple, Any], Any]:
         """Build a forward hook exposing ``output * scale`` through hook_point."""
-        scale = self.residual_contribution_scale
 
         def _hook(_module: torch.nn.Module, _args: tuple, output: Any) -> Any:
             if not hook_point.has_hooks(dir="both"):

transformer_lens/model_bridge/sources/_bridge_builder.py

@@ -73,6 +73,9 @@
     "mamba_d_conv",
     "mamba_chunk_size",
     "lm_head_multiplier",
+    "attention_out_multiplier",
+    "ssm_out_multiplier",
+    "mlp_multipliers",
     # Multimodal
     "vision_config",
     # Cohere