diff --git a/devtools/adapter_builder/docs/adapter-specification.md b/devtools/adapter_builder/docs/adapter-specification.md
index 3fb35ae8b5..bceb7c7233 100644
--- a/devtools/adapter_builder/docs/adapter-specification.md
+++ b/devtools/adapter_builder/docs/adapter-specification.md
@@ -249,7 +249,7 @@ def setup_component_testing(self, hf_model, bridge_model=None):
 
 ### `preprocess_weights(state_dict)`
 
-Apply architecture-specific weight transformations before standard processing. Example: Gemma scales embeddings by `sqrt(d_model)`.
+Apply architecture-specific weight transformations before standard processing. Example: Cohere folds `logit_scale` into `unembed.weight`. Do not fold runtime-applied scales such as Gemma's `sqrt(d_model)` embedding scale; HF applies it itself and a fold double-scales.
 
 ### `prepare_loading(model_name, model_kwargs)`
 
diff --git a/docs/source/content/adapter_development/adapter-creation-guide.md b/docs/source/content/adapter_development/adapter-creation-guide.md
index bb5e7da772..378341d867 100644
--- a/docs/source/content/adapter_development/adapter-creation-guide.md
+++ b/docs/source/content/adapter_development/adapter-creation-guide.md
@@ -177,7 +177,7 @@ For combined-QKV models (GPT-2 style), see `gpt2.py`'s `QKVSplitRearrangeConvers
 Implement only the ones you need:
 
 - **`setup_component_testing(hf_model, bridge_model=None)`** — required for RoPE models, to wire the rotary embedding instance through to the attention bridges. Skip for models with standard positional embeddings.
-- **`preprocess_weights(state_dict)`** — for arch-specific weight transforms before standard processing (e.g., Gemma scales embeddings by `sqrt(d_model)`).
+- **`preprocess_weights(state_dict)`** — for arch-specific weight transforms before standard processing (e.g., Cohere folds `logit_scale` into `unembed.weight`; not Gemma's `sqrt(d_model)` embedding scale, which HF applies at runtime).
 - **`prepare_loading(model_name, model_kwargs)`** — patch HF model classes before `from_pretrained()`.
 - **`prepare_model(hf_model)`** — post-load fixups before bridge creation.
 
diff --git a/docs/source/content/adapter_development/adapter-specification.md b/docs/source/content/adapter_development/adapter-specification.md
index 123a2137df..3deb394c6c 100644
--- a/docs/source/content/adapter_development/adapter-specification.md
+++ b/docs/source/content/adapter_development/adapter-specification.md
@@ -252,7 +252,7 @@ def setup_component_testing(self, hf_model, bridge_model=None):
 
 ### `preprocess_weights(state_dict)`
 
-Apply architecture-specific weight transformations before standard processing. Example: Gemma scales embeddings by `sqrt(d_model)`.
+Apply architecture-specific weight transformations before standard processing. Example: Cohere folds `logit_scale` into `unembed.weight`. Do not fold runtime-applied scales such as Gemma's `sqrt(d_model)` embedding scale; HF applies it itself and a fold double-scales.
 
 ### `prepare_loading(model_name, model_kwargs)`
 
diff --git a/docs/source/content/analysis_tools.md b/docs/source/content/analysis_tools.md
index df0b7b1de0..6025ac5e25 100644
--- a/docs/source/content/analysis_tools.md
+++ b/docs/source/content/analysis_tools.md
@@ -34,12 +34,12 @@ available hooks, and the meaning of the selected tensor axes.
 
 | Tool | Requirements to check | Main compute or memory cost |
 |---|---|---|
-| DLA | Targets TransformerBridge in 4.0. Enable compatibility mode and use the standard attention/MLP residual decomposition; Mamba/SSM/Mixer/LinearAttention hybrid layouts are rejected, as are bridges whose readout is not `ln_final -> unembed` on the last block (multi-stream AltUp blocks, a head transform such as BERT's `mlm_head`). When DLA runs the prompt, the model and all submodules must be in evaluation mode, or it raises `ValueError` before the forward; a supplied cache is accepted as is. | One cached forward pass, or reuse of a suitable cache. Head decomposition can require additional per-head results. |
+| DLA | Targets TransformerBridge in 4.0. Enable compatibility mode. Hybrid and SSM stacks are supported, with each sequence mixer's output as its own component (`unit="component"` raises `ValueError` before any forward for a block whose mixer output `decompose_resid` would not emit; a block with no `hook_mlp_out`, such as GraniteMoeHybrid with experts, contributes `resid_post - resid_mid` as its MLP row); `unit="layer"` raises `ValueError` before any forward when a sequential block has no `resid_mid` (mixer-only blocks such as Mamba, NemotronH and Zamba2). Bridges whose readout is not `ln_final -> unembed` on the last block (multi-stream AltUp blocks, a head transform such as BERT's `mlm_head`) are rejected. When DLA runs the prompt, the model and all submodules must be in evaluation mode, or it raises `ValueError` before the forward; a supplied cache is accepted as is. | One cached forward pass, or reuse of a suitable cache. Head decomposition can require additional per-head results. |
 | Activation Patching | Any bridge whose activation names resolve (`resid_pre`, `attn_out`, `mlp_out`, `z`, `q`/`k`/`v`, `pattern`). The metric returns a scalar by default and a `[batch]` tensor under `reduce="mean"` or `"none"`; several metrics can be passed as a mapping and share one sweep. Run the model in evaluation mode. | One forward of the corrupted batch per swept cell (plus one for `baseline=True`); per-example storage is `cells × batch` floats per metric. Statistics are computed afterwards from the stored values and add no forwards. |
 | Attribution Patching | Targets TransformerBridge. Clean/corrupt inputs must be `[batch, seq]` token tensors with matching shapes and aligned positions. The required embedding, head, and MLP hook aliases must exist. Edge granularity additionally requires attention bridges that support per-head results and a pre-Q/K/V residual fork. On GQA/MQA models, where any attention layer's K/V head count (`per_layer_num_key_value_heads` where set, else `n_key_value_heads`) differs from `n_heads`, edge granularity and faithfulness raise `ValueError` before any forward; node granularity remains available. Edge granularity and faithfulness are unsupported on post-norm blocks (such as Gemma 2/3, OLMo 2/3 and GLM-4), scaled-residual blocks (such as Granite) and parallel attention/SSM blocks (such as Falcon-H1), where per-head outputs are normalized, scaled or joined by an unmodeled mixer before the residual add; both raise there, and node granularity remains available. Edge sweeps temporarily enable the required writer and reader hooks, then restore the caller's hook flags. Run the model and all nested stochastic modules in evaluation mode. The metric must return a differentiable scalar, and the current backward-cache path requires an active autograd graph; fully frozen models are not yet supported. | Two forwards and one backward **per prompt pair**; activation and gradient caches. For a uniform-head decoder, the number of edge scores grows as `O(seq_len × (n_layers × n_heads)²)`, while the cache also gains per-head `[batch, seq, n_heads, d_model]` tensors. Pair scores are averaged across the batch. |
 | Direct Path Patching | Requires a TransformerBridge whose attention layers have the pre-Q/K/V residual fork that `use_split_qkv_input` enables (the joint-QKV and position-embedding attention bridges, e.g. GPT-2, Pythia, Llama, Mistral, Qwen2); attention without it, such as GPT-J, GPT-Neo, OPT, BERT or a bridge built with `TransformerBridge.boot_native`, raises `NotImplementedError`. The supplied caches must hold the source layer's `hook_z`. The function turns the fork flag on for the sweep and restores the caller's `use_split_qkv_input` and `use_attn_in`; LayerNorm need not be folded. Q destinations use query heads; K/V destinations use native KV heads, shared by groups of query heads on GQA/MQA. K/V sweeps require a uniform KV count across destination layers (`per_layer_num_key_value_heads` where set, else `n_key_value_heads` or `n_heads`), and raise `ValueError` before any forward if those counts differ. Every block from the source layer on must have an `attn` component; SSM, mixer, linear-attention and AltUp blocks raise `ValueError` before any forward. Source blocks must add the head outputs to the residual unchanged: post-norm blocks (such as Gemma 2/3, OLMo 2/3 and GLM-4), scaled-residual blocks (such as Granite) and parallel attention/SSM blocks (such as Falcon-H1) raise `ValueError` before any forward, since a head's residual write there is not `hook_z @ W_O`. Run the model and all submodules in evaluation mode; both entry points raise `ValueError` before any forward otherwise. | Clean/corrupt caching, then one forward of the corrupted batch per later destination query head or KV group. While the fork is on, every attention layer normalizes `n_heads` Q copies and `n_kv_heads` copies each for K and V, about `n_heads + 2 × n_kv_heads` times the usual attention-input activation memory. Sweeping all source heads adds another sweep dimension. |
-| Logit Lens | Targets TransformerBridge, raw or compatibility mode. Needs a decoder with an `ln_final` component (or no final norm) and `W_U`; encoder-decoder and encoder-only bridges are rejected, as are quantized unembeddings and bridges whose readout is not `ln_final -> unembed` on the last block (multi-stream AltUp blocks, a head transform such as BERT's `mlm_head`; a prompt run with `apply_ln=True` is also checked against the model's own output, a supplied cache only structurally). Hybrid SSM stacks work with `incl_mid=False`. Run the model in evaluation mode. | One cached forward pass over the residual hooks, or reuse of a cache. Logits are produced in chunks of `chunk_size` rows; full-vocabulary output is materialized only when no `targets`, `top_k` or `vocab` selector is given. |
-| Jacobian Lens | Requires a fresh, causal decoder-only Bridge with raw HF weights, without compatibility mode or weight processing. Validate the lens against the model. Run the model and all submodules in evaluation mode; fitting and every call that runs the model raise `ValueError` otherwise. | Loading an existing artifact avoids fitting. The ordinary fitting estimator uses one forward and `ceil(d_model / dim_batch)` backwards per prompt; larger batches increase memory. |
+| Logit Lens | Targets TransformerBridge, raw or compatibility mode. Needs a decoder with an `ln_final` component (or no final norm) and `W_U`; encoder-decoder and encoder-only bridges are rejected, as are quantized unembeddings and bridges whose readout is not `ln_final -> unembed` on the last block (multi-stream AltUp blocks, a head transform such as BERT's `mlm_head`; a prompt run with `apply_ln=True` is also checked against the model's own output, a supplied cache only structurally). `incl_mid=True` needs a `resid_mid` hook on every sequential block, so it raises `ValueError` on mixer-only blocks (Mamba, NemotronH, Zamba2) and works on hybrids whose blocks all have one (Qwen3.5, Jamba, RecurrentGemma). Run the model in evaluation mode. | One cached forward pass over the residual hooks, or reuse of a cache. Logits are produced in chunks of `chunk_size` rows; full-vocabulary output is materialized only when no `targets`, `top_k` or `vocab` selector is given. |
+| Jacobian Lens | Requires a fresh, causal decoder-only Bridge with raw HF weights, without compatibility mode or weight processing; bridges whose readout is not `ln_final -> unembed` on the last block (multi-stream blocks, a head transform such as ModernBERT-decoder's `prediction_head`) are rejected. Validate the lens against the model. Run the model and all submodules in evaluation mode; fitting and every call that runs the model raise `ValueError` otherwise. | Loading an existing artifact avoids fitting. The ordinary fitting estimator uses one forward and `ceil(d_model / dim_batch)` backwards per prompt; larger batches increase memory. |
 | Interchange Interventions | Any model with the hook API (TransformerBridge or `HookedRootModule`). Activations at the hook must be `[batch, d_model]`, or `[batch, pos, d_model]` with one position index per example; base and source rows are paired. | One forward pass per `batch_size` sources to collect activations, then one per `batch_size` bases per intervention with `intervened_logits`; both run without gradients, so memory is bounded by `batch_size`. `run_with_intervention` is a single differentiable pass over all the bases it is given. DAS adds a forward and backward pass per batch per epoch; `micro_batch_size` bounds its memory. |
 | Projection Kernel | The numerical API accepts finite, real floating-point matrices via orthonormal bases. The attention-head wrapper requires Bridge weights with compatible dimensions and ranks. | Basis extraction uses SVD. All-head comparisons allocate basis stacks and a pairwise score grid; they can be large despite requiring no forward pass. |
 
@@ -89,6 +89,12 @@ it is exact for twelve or fewer prompts and a Monte Carlo estimate above that.
 Neither says anything about prompts outside the batch, and neither changes what a
 single cell establishes: the measured effect of one intervention.
 
+On GQA/MQA models the `*_every` helpers pad the `k`/`v` sweeps to `n_heads`. The
+padded cells are zero in `values` and `mean()` (matching the tensor return),
+`NaN` in every statistic, absent from `to_dataframe()`, and `False` in
+`result.valid`; select `p[..., result.valid]` (the mask spans the swept axes, so the
+ellipsis also skips a leading metric axis) before any multiple-comparison correction.
+
 ### Attribution Patching: screen candidates, then measure interventions
 
 At node granularity, the score uses
diff --git a/docs/source/content/contributing.md b/docs/source/content/contributing.md
index 06c1b24fb2..cb3bc50ad3 100644
--- a/docs/source/content/contributing.md
+++ b/docs/source/content/contributing.md
@@ -304,7 +304,7 @@ Two test layers:
 ### Common adapter gotchas
 
 - **HF raw config attributes are invisible to TL-side consumers unless explicitly propagated to `self.cfg`.** Walk the HF `config.json` and mirror any non-standard knobs (`sliding_window`, `layer_types`, …) onto `self.cfg` so weight processing and forward passes can see them. Gemma's softcaps are already renamed by the bridge builder (`cfg.output_logits_soft_cap`, `cfg.attn_scores_soft_cap`), and `query_pre_attn_scalar` needs nothing — the wrapped HF attention applies it.
-- **Some config attrs need both surface-on-cfg AND fold-into-weight** via a `preprocess_weights()` override. The trigger: a numerical operation HF's forward applies natively must also be baked into the raw weights, or `bridge.enable_compatibility_mode()` (which calls `process_weights` on raw weights) produces wrong results. Concrete examples in-tree: Cohere `logit_scale` → `unembed.weight`; Gemma embedding scale (`√d_model`) → `embed.weight`. Skip the fold and Phase 3 / Phase 4 of `verify_models` will silently degrade.
+- **Some config attrs need both surface-on-cfg AND fold-into-weight** via a `preprocess_weights()` override. The trigger: a numerical operation HF's forward applies natively must also be baked into the raw weights, or `bridge.enable_compatibility_mode()` (which calls `process_weights` on raw weights) produces wrong results. Concrete examples in-tree: Cohere `logit_scale` → `unembed.weight`. Gemma is the counter-example: HF applies its `√d_model` embedding scale at runtime on both raw and compat paths and `bridge.embed` already captures it, so no Gemma adapter overrides `preprocess_weights`; folding would double-scale. Skip a needed fold and Phase 3 / Phase 4 of `verify_models` will silently degrade.
 - **Tokenizer policy is per-model, not per-architecture.** Sibling models in the same family routinely differ — the chat-instruct variant may prepend BOS where the base does not, padding side can flip, EOS handling can differ. It's worth re-checking `default_prepend_bos`, padding side, and EOS handling against the specific target rather than copying them from a starter adapter. `tokenizer_config.json` is not always reliable on its own — some architectures (Cohere is a notable example) declare `add_bos_token=False` but HF's `__call__` prepends BOS anyway. The most reliable check is to invoke the tokenizer directly:
 
   ```python
diff --git a/docs/source/content/jacobian_lens_fitting.md b/docs/source/content/jacobian_lens_fitting.md
index 3f617ca0ff..abdd31c233 100644
--- a/docs/source/content/jacobian_lens_fitting.md
+++ b/docs/source/content/jacobian_lens_fitting.md
@@ -21,7 +21,9 @@ manifest, and set of estimator options.
 ## Requirements
 
 `JacobianLens.fit` requires a freshly booted, causal decoder-only
-`TransformerBridge` with raw Hugging Face weights:
+`TransformerBridge` with raw Hugging Face weights, whose logits are
+`unembed(ln_final(...))` of the last block's single-stream output (a head
+transform such as ModernBERT-decoder's `prediction_head` is rejected):
 
 ```python
 import torch
diff --git a/tests/integration/model_bridge/test_forward_padding_side.py b/tests/integration/model_bridge/test_forward_padding_side.py
index 31e4221180..aa0a4e3c06 100644
--- a/tests/integration/model_bridge/test_forward_padding_side.py
+++ b/tests/integration/model_bridge/test_forward_padding_side.py
@@ -184,3 +184,19 @@ def test_right_padded_logits_still_run_unmasked(distilgpt2_bridge) -> None:
         both_logits, _ = distilgpt2_bridge(padded, return_type="both", padding_side="right")
 
     torch.testing.assert_close(both_logits, logits, rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("prepend_bos", [True, False])
+def test_right_padded_loss_keeps_a_real_trailing_eos(distilgpt2_bridge, prepend_bos) -> None:
+    """A real trailing EOS shares the pad id here, yet its transition still counts."""
+    bridge = distilgpt2_bridge
+    eos = bridge.tokenizer.eos_token
+    prompts = ["hello" + eos, "the cat sat on a mat" + eos]
+    with torch.no_grad():
+        per_token = [
+            bridge(prompt, prepend_bos=prepend_bos, return_type="loss", loss_per_token=True)[0]
+            for prompt in prompts
+        ]
+        batched = bridge(prompts, prepend_bos=prepend_bos, padding_side="right", return_type="loss")
+
+    torch.testing.assert_close(batched, torch.cat(per_token).mean())
diff --git a/tests/integration/model_bridge/test_gemma3n_bridge.py b/tests/integration/model_bridge/test_gemma3n_bridge.py
index 92b592ecb3..6ee26ebf6d 100644
--- a/tests/integration/model_bridge/test_gemma3n_bridge.py
+++ b/tests/integration/model_bridge/test_gemma3n_bridge.py
@@ -66,6 +66,21 @@ def test_run_with_cache_text_only(bridge):
     assert len(cache) > 0
 
 
+def test_svd_w_in_uses_hf_ln2_scale(bridge):
+    """W_in is scaled by the per-channel factor HF's ln2 actually applies."""
+    from transformer_lens.SVDInterpreter import SVDInterpreter
+
+    interpreter = SVDInterpreter(bridge)
+    for layer in range(bridge.cfg.n_layers):
+        hf_ln2 = bridge.blocks[layer].ln2.original_component
+        x = torch.randn(8, bridge.cfg.d_model)
+        x_hat = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + hf_ln2.eps)
+        with torch.no_grad():
+            scale = (hf_ln2(x) / x_hat).mean(0)
+        expected = interpreter.params[f"blocks.{layer}.mlp.W_in"].T * scale
+        torch.testing.assert_close(interpreter._get_w_in_matrix(layer), expected)
+
+
 def test_contribution_hooks_are_post_norm_outputs(bridge):
     """hook_attn_out / hook_mlp_out carry the post-sublayer norm outputs, the tensors the
     layer actually adds (Gemma 2/3 contract), checked against an independent HF run."""
@@ -113,3 +128,9 @@ def test_half_converts_unmapped_parameters():
     assert unconverted == []
     with torch.no_grad():
         assert bridge(IDS).dtype == torch.float16
+
+
+def test_start_at_layer_is_refused_on_altup_blocks(bridge):
+    """An AltUp block carries a stack of streams, so there is no single residual to inject."""
+    with pytest.raises(NotImplementedError, match="AltUpBlockBridge"):
+        bridge(torch.zeros(1, 3, bridge.cfg.d_model), start_at_layer=1)
diff --git a/tests/integration/model_bridge/test_post_norm_head_decomposition.py b/tests/integration/model_bridge/test_post_norm_head_decomposition.py
index 2c81bd32aa..64d51bc5d3 100644
--- a/tests/integration/model_bridge/test_post_norm_head_decomposition.py
+++ b/tests/integration/model_bridge/test_post_norm_head_decomposition.py
@@ -26,6 +26,7 @@
 )
 
 from transformer_lens.model_bridge.sources import build_bridge_from_module
+from transformer_lens.utilities.slice import Slice
 
 _TINY = dict(
     hidden_size=64,
@@ -165,6 +166,37 @@ def test_output_bias_is_mapped_through_post_norm(hf_architecture: str, config_cl
             )
 
 
+@pytest.mark.parametrize("hf_architecture,config_cls,config_kwargs", CASES)
+@pytest.mark.parametrize("neuron_slice", [5, Slice(5)], ids=["int", "Slice-int"])
+@pytest.mark.parametrize("pos_slice", [None, -1], ids=["all-pos", "last-pos"])
+@pytest.mark.parametrize("batched", [True, False], ids=["batched", "batchless"])
+@pytest.mark.parametrize("n_outs", [None, 0, 2], ids=["unprojected", "proj-1d", "proj-2d"])
+def test_integer_neuron_slice_matches_list(
+    hf_architecture: str,
+    config_cls: Any,
+    config_kwargs: dict[str, Any],
+    neuron_slice: Any,
+    pos_slice: Any,
+    batched: bool,
+    n_outs: Any,
+) -> None:
+    """An int neuron_slice is the one-element list slice with the neuron axis dropped."""
+    bridge, cache = _run(hf_architecture, config_cls, config_kwargs)
+    if not batched:
+        cache = cache.apply_slice_to_batch_dim(0)
+    d_model = bridge.cfg.d_model
+    project = None
+    if n_outs is not None:
+        shape = (d_model,) if n_outs == 0 else (d_model, n_outs)
+        project = torch.randn(shape, generator=torch.Generator().manual_seed(0))
+    kwargs = dict(layer=0, pos_slice=pos_slice, project_output_onto=project)
+    reference = cache.get_neuron_results(neuron_slice=[5], **kwargs)
+    expected = reference[..., 0] if n_outs == 0 else reference[..., 0, :]
+    result = cache.get_neuron_results(neuron_slice=neuron_slice, **kwargs)
+    assert result.shape == expected.shape
+    torch.testing.assert_close(result, expected)
+
+
 def test_half_precision_stacks_keep_model_dtype() -> None:
     bridge, cache = _run("Gemma2ForCausalLM", Gemma2Config, _TINY, dtype=torch.bfloat16)
     assert cache.stack_head_results().dtype == torch.bfloat16
diff --git a/tests/integration/model_bridge/test_ssm_hybrid_resid_decomposition.py b/tests/integration/model_bridge/test_ssm_hybrid_resid_decomposition.py
index b31720e025..52d5de7eb8 100644
--- a/tests/integration/model_bridge/test_ssm_hybrid_resid_decomposition.py
+++ b/tests/integration/model_bridge/test_ssm_hybrid_resid_decomposition.py
@@ -136,9 +136,9 @@ def _nemotronh():
     return _boot(AutoModelForCausalLM.from_config(c).eval(), "NemotronHForCausalLM")
 
 
-def _granite_dense():
-    """Dense GraniteMoeHybrid (shared MLP, no experts) with a non-unit residual_multiplier:
-    every sublayer output is scaled after its hook_out before the residual add."""
+def _granite_hybrid(num_local_experts: int):
+    """GraniteMoeHybrid with a non-unit residual_multiplier: every sublayer output is scaled
+    after its hook_out before the residual add."""
     from transformers import AutoModelForCausalLM
     from transformers.models.granitemoehybrid import GraniteMoeHybridConfig
 
@@ -151,7 +151,8 @@ def _granite_dense():
         num_hidden_layers=3,
         num_attention_heads=8,
         num_key_value_heads=4,
-        num_local_experts=0,
+        num_local_experts=num_local_experts,
+        num_experts_per_tok=1,
         max_position_embeddings=64,
         layer_types=["mamba", "attention", "mamba"],
         mamba_n_heads=8,
@@ -169,6 +170,16 @@ def _granite_dense():
     return _boot(AutoModelForCausalLM.from_config(c).eval(), "GraniteMoeHybridForCausalLM")
 
 
+def _granite_dense():
+    """Shared MLP only: the MLP-side write fires hook_mlp_out."""
+    return _granite_hybrid(num_local_experts=0)
+
+
+def _granite_moe():
+    """MoE + shared MLP summed inline before the scaled add, so no hook_mlp_out exists."""
+    return _granite_hybrid(num_local_experts=2)
+
+
 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."""
@@ -251,6 +262,7 @@ def _lfm2():
 
 _FACTORIES = {
     "granite_dense": _granite_dense,
+    "granite_moe": _granite_moe,
     "lfm2": _lfm2,
     "mamba2": _mamba2,
     "mamba2_bias": _mamba2_bias,
@@ -309,6 +321,7 @@ def test_labels_name_the_mixer(self, bridge_and_cache):
         expected = {
             "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"],
+            "granite_moe": ["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"],
             "mamba2": ["0_mixer_out", "1_mixer_out"],
             "mamba2_bias": ["0_mixer_out", "1_mixer_out"],
@@ -354,6 +367,36 @@ def test_mlp_input_sums_to_resid_mid(self, bridge_and_cache):
                 comps.sum(0), cache[f"blocks.{layer}.hook_resid_mid"], atol=ATOL, rtol=0
             )
 
+    def test_resid_mid_is_resid_pre_plus_the_attention_slot(self, bridge_and_cache):
+        """``hook_resid_mid`` must sit right after the attention-slot write, on every block
+        that has one (mixer-only and parallel blocks have none to check)."""
+        _, bridge, cache = bridge_and_cache
+        hooks = bridge.hook_dict
+        for layer in range(bridge.cfg.n_layers):
+            if f"blocks.{layer}.hook_resid_mid" not in hooks:
+                continue
+            rows, labels = cache.decompose_resid(
+                layer=layer, mlp_input=True, mode="attn", incl_embeds=False, return_labels=True
+            )
+            slot = [row for row, label in zip(rows, labels) if label.startswith(f"{layer}_")]
+            assert slot, labels
+            write = (
+                cache[f"blocks.{layer}.hook_resid_mid"] - cache[f"blocks.{layer}.hook_resid_pre"]
+            )
+            torch.testing.assert_close(write, torch.stack(slot).sum(0), atol=ATOL, rtol=0)
+
+    @pytest.mark.parametrize("mode", ["all", "attn", "mlp"])
+    @pytest.mark.parametrize("mlp_input", [False, True])
+    def test_labels_are_known_before_any_activation(self, bridge_and_cache, mode, mlp_input):
+        _, bridge, cache = bridge_and_cache
+        _, labels = cache.decompose_resid(mode=mode, mlp_input=mlp_input, return_labels=True)
+        assert cache.decompose_resid_labels(mode=mode, mlp_input=mlp_input) == labels
+        from transformer_lens import ActivationCache
+
+        assert ActivationCache({}, bridge).decompose_resid_labels(
+            mode=mode, mlp_input=mlp_input, incl_embeds=False
+        ) == [label for label in labels if label not in ("embed", "pos_embed")]
+
     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]))
@@ -389,7 +432,7 @@ class TestFullResidDecomposition:
     @pytest.mark.parametrize("expand_neurons", [False, True])
     def test_sums_to_final_residual(self, bridge_and_cache, expand_neurons):
         name, bridge, cache = bridge_and_cache
-        if name == "granite_dense" and expand_neurons:
+        if name in ("granite_dense", "granite_moe") and expand_neurons:
             # Its MLP is `shared_mlp`, outside the `mlp` slot neuron expansion reads.
             with pytest.raises(NotImplementedError, match="shared_mlp"):
                 cache.get_full_resid_decomposition(expand_neurons=True)
@@ -405,7 +448,7 @@ def test_ln_and_projection_sum_to_projected_scaled_residual(
         self, bridge_and_cache, expand_neurons
     ):
         name, bridge, cache = bridge_and_cache
-        if name == "granite_dense" and expand_neurons:
+        if name in ("granite_dense", "granite_moe") and expand_neurons:
             # Its MLP is `shared_mlp`, outside the `mlp` slot neuron expansion reads.
             with pytest.raises(NotImplementedError, match="shared_mlp"):
                 cache.get_full_resid_decomposition(expand_neurons=True)
@@ -433,6 +476,22 @@ def test_zamba2_hybrid_layers_write_only_their_mamba_output():
     assert not torch.allclose(cache["blocks.1.norm.hook_in"], cache["blocks.1.hook_resid_pre"])
 
 
+def test_granite_moe_mlp_row_is_the_inline_moe_plus_shared_mlp_write():
+    """With experts, no single module feeds the MLP-side add, so the block has no hook_mlp_out;
+    decompose_resid derives the row from the residual stream instead of raising."""
+    bridge = _granite_moe()
+    assert not any(name.endswith("hook_mlp_out") for name in bridge.hook_dict)
+    with torch.no_grad():
+        _, cache = bridge.run_with_cache(TOKENS)
+    rows, labels = cache.decompose_resid(mode="mlp", incl_embeds=False, return_labels=True)
+    assert labels == [f"{layer}_mlp_out" for layer in range(bridge.cfg.n_layers)]
+    for layer, row in enumerate(rows):
+        write = (
+            cache[f"blocks.{layer}.moe.hook_out"] + cache[f"blocks.{layer}.shared_mlp.hook_out"]
+        ) * bridge.blocks[layer].residual_contribution_scale
+        torch.testing.assert_close(row, write, atol=ATOL, rtol=0)
+
+
 def test_recurrent_gemma_hook_embed_is_the_scaled_embedding():
     """HF scales the embedding by its bfloat16 normalizer after embed_tokens."""
     bridge = _recurrent_gemma()
diff --git a/tests/integration/test_hybrid_dla_logit_lens.py b/tests/integration/test_hybrid_dla_logit_lens.py
new file mode 100644
index 0000000000..ba32935c82
--- /dev/null
+++ b/tests/integration/test_hybrid_dla_logit_lens.py
@@ -0,0 +1,184 @@
+"""DLA and logit-lens contracts on SSM / hybrid bridges.
+
+Component and head attributions must sum to the logit difference on every hybrid,
+since neither reads ``resid_mid``. Only the ``resid_mid`` readers (DLA
+``unit="layer"``, logit-lens ``incl_mid=True``) refuse, and only on blocks without one.
+"""
+
+from typing import Any, Callable, Dict
+
+import pytest
+import torch
+
+from tests.integration.model_bridge.test_ssm_hybrid_resid_decomposition import (
+    _QWEN3_5_AVAILABLE,
+    TOKENS,
+    _boot,
+    _granite_dense,
+    _granite_moe,
+    _mamba2,
+    _nemotronh,
+    _qwen35,
+)
+from transformer_lens.tools.analysis import direct_logit_attribution, logit_lens
+from transformer_lens.tools.analysis.logit_lens import logit_readout
+
+ATOL = 1e-5
+ANSWER = torch.tensor([3, 3])
+WRONG = torch.tensor([5, 5])
+
+needs_qwen35 = pytest.mark.skipif(not _QWEN3_5_AVAILABLE, reason="needs Qwen3.5 in transformers")
+
+
+def _falcon_h1():
+    """Falcon-H1 runs Mamba beside attention in every block (MHA, so compat mode folds)."""
+    from transformers import AutoModelForCausalLM
+    from transformers.models.falcon_h1 import FalconH1Config
+
+    torch.manual_seed(0)
+    c = FalconH1Config(
+        vocab_size=128,
+        hidden_size=64,
+        intermediate_size=128,
+        num_hidden_layers=2,
+        num_attention_heads=4,
+        num_key_value_heads=4,
+        mamba_d_ssm=64,
+        mamba_n_heads=4,
+        mamba_d_head=16,
+        mamba_n_groups=1,
+        mamba_d_state=16,
+        mamba_d_conv=4,
+        mamba_expand=2,
+        mamba_chunk_size=8,
+    )
+    c.architectures = ["FalconH1ForCausalLM"]
+    return _boot(AutoModelForCausalLM.from_config(c).eval(), "FalconH1ForCausalLM")
+
+
+_FACTORIES: Dict[str, Callable[[], Any]] = {
+    "qwen3_5": _qwen35,
+    "mamba2": _mamba2,
+    "nemotron_h": _nemotronh,
+    "granite_dense": _granite_dense,
+    "granite_moe": _granite_moe,
+    "falcon_h1": _falcon_h1,
+}
+
+
+def _compat(name):
+    bridge = _FACTORIES[name]()
+    bridge.enable_compatibility_mode(disable_warnings=True)
+    return bridge.eval()
+
+
+def _logit_diff(bridge):
+    """Residual-stream share of the answer-minus-wrong logit at the last position."""
+    with torch.no_grad():
+        logits = bridge(TOKENS)
+    diff = logits[:, -1, 3] - logits[:, -1, 5]
+    b_U = getattr(bridge, "b_U", None)
+    return diff if b_U is None else diff - (b_U[3] - b_U[5])
+
+
+def _dla(bridge, unit):
+    return direct_logit_attribution(
+        bridge, TOKENS, answer_tokens=ANSWER, incorrect_tokens=WRONG, unit=unit
+    )
+
+
+@pytest.mark.parametrize(
+    "name", [pytest.param("qwen3_5", marks=needs_qwen35), "mamba2", "nemotron_h"]
+)
+def test_component_attribution_sums_to_logit_diff(name):
+    bridge = _compat(name)
+    result = _dla(bridge, "component")
+    torch.testing.assert_close(result.attribution.sum(0), _logit_diff(bridge), atol=ATOL, rtol=0)
+
+
+@pytest.mark.parametrize(
+    "name", [pytest.param("qwen3_5", marks=needs_qwen35), "granite_dense", "granite_moe"]
+)
+def test_head_attribution_sums_to_logit_diff(name):
+    bridge = _compat(name)
+    result = _dla(bridge, "head")
+    assert result.labels[-1] == "remainder"
+    torch.testing.assert_close(result.attribution.sum(0), _logit_diff(bridge), atol=ATOL, rtol=0)
+
+
+def test_component_attribution_covers_blocks_without_hook_mlp_out():
+    """GraniteMoeHybrid with experts (the shipped Granite-4.0-H layout) sums its MoE and shared
+    MLP inline, so no hook_mlp_out exists; the MLP-side row must still be attributed."""
+    bridge = _compat("granite_moe")
+    assert not any(name.endswith("hook_mlp_out") for name in bridge.hook_dict)
+    result = _dla(bridge, "component")
+    assert [f"{l}_mlp_out" for l in range(bridge.cfg.n_layers)] == [
+        label for label in result.labels if label.endswith("_mlp_out")
+    ]
+    torch.testing.assert_close(result.attribution.sum(0), _logit_diff(bridge), atol=ATOL, rtol=0)
+
+
+_WITH_RESID_MID = [pytest.param("qwen3_5", marks=needs_qwen35), "granite_dense", "granite_moe"]
+
+
+@pytest.mark.parametrize("name", _WITH_RESID_MID)
+def test_layer_attribution_mid_rows_add_the_attention_slot(name):
+    """``{l}_mid`` must read the residual right after block l's attention-slot write."""
+    bridge = _compat(name)
+    by_layer = _dla(bridge, "layer")
+    by_component = _dla(bridge, "component")
+    for layer in range(bridge.cfg.n_layers):
+        step = (
+            by_layer.attribution[by_layer.labels.index(f"{layer}_mid")]
+            - by_layer.attribution[by_layer.labels.index(f"{layer}_pre")]
+        )
+        slot = [
+            row
+            for row, label in zip(by_component.attribution, by_component.labels)
+            if label.startswith(f"{layer}_") and not label.endswith("_mlp_out")
+        ]
+        assert slot, by_component.labels
+        torch.testing.assert_close(step, torch.stack(slot).sum(0), atol=ATOL, rtol=0)
+
+
+def test_layer_attribution_refuses_mixer_only_blocks_before_any_forward(monkeypatch):
+    bridge = _compat("mamba2")
+
+    def no_forward(*args, **kwargs):
+        raise AssertionError("the forward pass ran before the resid_mid check")
+
+    monkeypatch.setattr(bridge, "run_with_cache", no_forward)
+    with pytest.raises(ValueError, match=r"blocks \[0, 1\]"):
+        _dla(bridge, "layer")
+
+
+def test_falcon_h1_component_attribution_is_complete_or_refused():
+    """A parallel Mamba branch must be attributed, or DLA must refuse; never a short sum."""
+    bridge = _compat("falcon_h1")
+    try:
+        result = _dla(bridge, "component")
+    except ValueError as err:
+        assert "mamba output is not decomposed yet" in str(err)
+        return
+    torch.testing.assert_close(result.attribution.sum(0), _logit_diff(bridge), atol=ATOL, rtol=0)
+
+
+@pytest.mark.parametrize("name", _WITH_RESID_MID)
+def test_logit_lens_incl_mid_on_hybrid_with_resid_mid(name):
+    bridge = _FACTORIES[name]().eval()
+    n_layers = bridge.cfg.n_layers
+    with torch.no_grad():
+        ref = bridge(TOKENS)
+    _, cache = bridge.run_with_cache(TOKENS)
+    for result in (
+        logit_lens(bridge, TOKENS, incl_mid=True, positions=None),
+        logit_lens(bridge, cache=cache, incl_mid=True, positions=None),
+    ):
+        assert len(result.labels) == 2 * n_layers + 1
+        torch.testing.assert_close(result.values[-1], ref, atol=1e-4, rtol=1e-4)
+        for layer in range(n_layers):
+            # The mid row must be the readout of resid_mid itself, not of a neighbouring hook.
+            mid = logit_readout(bridge, cache[f"blocks.{layer}.hook_resid_mid"]).values
+            torch.testing.assert_close(
+                result.values[result.labels.index(f"{layer}_mid")], mid, atol=1e-4, rtol=1e-4
+            )
diff --git a/tests/integration/test_jacobian_lens.py b/tests/integration/test_jacobian_lens.py
index b097e00949..f43f1c5469 100644
--- a/tests/integration/test_jacobian_lens.py
+++ b/tests/integration/test_jacobian_lens.py
@@ -351,6 +351,52 @@ def test_bidirectional_bert_bridge_is_rejected_before_readout():
         lens.validate_model(model)
 
 
+def test_modernbert_decoder_prediction_head_is_rejected_before_fit():
+    """J^T W_U skips the dense+GELU+norm head that ModernBERT-decoder applies before lm_head."""
+    from transformers import ModernBertDecoderConfig, ModernBertDecoderForCausalLM
+
+    from transformer_lens.model_bridge.sources import build_bridge_from_module
+    from transformer_lens.tools.analysis import JacobianLens
+
+    config = ModernBertDecoderConfig(
+        vocab_size=32,
+        hidden_size=16,
+        intermediate_size=32,
+        num_hidden_layers=2,
+        num_attention_heads=4,
+        max_position_embeddings=16,
+        pad_token_id=0,
+        bos_token_id=0,
+        eos_token_id=0,
+    )
+    model = build_bridge_from_module(
+        ModernBertDecoderForCausalLM(config).eval(),
+        "ModernBertDecoderForCausalLM",
+        hf_config=config,
+        dtype=torch.float32,
+        device="cpu",
+        model_name="tiny-random-modernbert-decoder-jacobian-lens-contract",
+    )
+    lens = JacobianLens(
+        {0: torch.eye(config.hidden_size)},
+        n_prompts=1,
+        d_model=config.hidden_size,
+        metadata={"target_layer": 1},
+    )
+
+    with pytest.raises(ValueError, match=r"mounts \['prediction_head'\]"):
+        lens.validate_model(model)
+    # No tokenizer is attached, so reaching tokenization would raise a different error.
+    with pytest.raises(ValueError, match=r"mounts \['prediction_head'\]"):
+        JacobianLens.fit(
+            model,
+            ["a b c"],
+            corpus="c",
+            skip_first_positions=0,
+            show_progress=False,
+        )
+
+
 @pytest.mark.slow
 def test_gemma_published_artifact_parity_topk_and_softcap():
     """Gemma readout matches the artifact formula and its soft-capped output."""
diff --git a/tests/unit/model_bridge/generalized_components/test_layernorm_weight_offset.py b/tests/unit/model_bridge/generalized_components/test_layernorm_weight_offset.py
index e6c2e77895..60968ac9b1 100644
--- a/tests/unit/model_bridge/generalized_components/test_layernorm_weight_offset.py
+++ b/tests/unit/model_bridge/generalized_components/test_layernorm_weight_offset.py
@@ -8,6 +8,7 @@
 
 from transformer_lens.model_bridge.generalized_components.normalization import (
     NormalizationBridge,
+    norm_uses_weight_offset,
 )
 
 
@@ -60,3 +61,33 @@ def test_python_layernorm1p_preserves_offset():
     bridge, component = _make_norm(native=False)
     x = torch.randn(2, 4, 16)
     torch.testing.assert_close(bridge(x), component(x), atol=1e-6, rtol=1e-5)
+
+
+@pytest.mark.parametrize("cfg_offset", [True, False])
+def test_norm_uses_weight_offset_prefers_mounted_norm(cfg_offset):
+    """A mounted norm's own declaration wins over the config; others fall back to it."""
+    cfg = SimpleNamespace(uses_rms_norm=True, rmsnorm_uses_offset=cfg_offset)
+    for declared in (True, False):
+        bridge = NormalizationBridge("ln2", cfg, uses_weight_offset=declared)
+        assert norm_uses_weight_offset(bridge, cfg) is declared
+    assert norm_uses_weight_offset(torch.nn.Identity(), cfg) is cfg_offset
+    assert norm_uses_weight_offset(None, cfg) is cfg_offset
+
+
+def test_python_fallback_keeps_the_wrapped_norms_eps():
+    """A backward hook routes the native path through the python norm, which must keep the
+    module's eps (Muse Glimmer's post_norm_eps) rather than the config's rms_norm_eps."""
+    component = NemotronLayerNorm1P(16, eps=1e-2)
+    bridge = NormalizationBridge(
+        "ln1_post",
+        SimpleNamespace(uses_rms_norm=False, rmsnorm_uses_offset=True, eps=1e-5),
+        use_native_layernorm_autograd=True,
+        uses_weight_offset=True,
+    )
+    bridge.set_original_component(component)
+    bridge.hook_scale.add_hook(lambda grad, hook: None, dir="bwd")
+    torch.manual_seed(0)
+    x = torch.randn(2, 3, 16) * 1e-2
+    with pytest.warns(UserWarning, match="python-norm path"):
+        out = bridge(x)
+    torch.testing.assert_close(out, component(x), atol=1e-6, rtol=0)
diff --git a/tests/unit/model_bridge/supported_architectures/test_gemma3n_adapter.py b/tests/unit/model_bridge/supported_architectures/test_gemma3n_adapter.py
index 1a68b9028e..4e957a8fee 100644
--- a/tests/unit/model_bridge/supported_architectures/test_gemma3n_adapter.py
+++ b/tests/unit/model_bridge/supported_architectures/test_gemma3n_adapter.py
@@ -71,7 +71,22 @@ def test_config_flags():
     assert a.supports_fold_ln is False
     assert a.weight_processing_conversions == {}
     assert a.cfg.normalization_type == "RMS"
-    assert a.cfg.rmsnorm_uses_offset is True
+
+
+def test_rmsnorm_offset_flag_matches_hf_norm():
+    """cfg.rmsnorm_uses_offset describes how Gemma3nRMSNorm applies its weight."""
+    from transformers.models.gemma3n.modeling_gemma3n import Gemma3nRMSNorm
+
+    torch.manual_seed(0)
+    d_model = 16
+    norm = Gemma3nRMSNorm(d_model, eps=1e-6)
+    with torch.no_grad():
+        norm.weight.copy_(torch.rand(d_model) + 0.5)
+        x = torch.randn(3, d_model)
+        x_hat = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + norm.eps)
+        w = norm.weight
+        expected = x_hat * ((1.0 + w) if _adapter().cfg.rmsnorm_uses_offset else w)
+        torch.testing.assert_close(norm(x), expected)
 
 
 def test_text_path_nested_under_language_model():
diff --git a/tests/unit/model_bridge/supported_architectures/test_muse_glimmer_adapter.py b/tests/unit/model_bridge/supported_architectures/test_muse_glimmer_adapter.py
index 17a0ca89b2..0a4903d9d7 100644
--- a/tests/unit/model_bridge/supported_architectures/test_muse_glimmer_adapter.py
+++ b/tests/unit/model_bridge/supported_architectures/test_muse_glimmer_adapter.py
@@ -94,3 +94,97 @@ def test_vision_projector_hook_out_is_the_merged_embedding(models):
     projector = bridge.vision_projector
     assert projector.name == "model.perception_emb_norm"
     assert type(projector.original_component).__name__ == "MuseGlimmerRMSNorm"
+
+
+def _hf_norm_scale(norm: torch.nn.Module, d_model: int) -> torch.Tensor:
+    """Effective per-channel scale of an HF RMSNorm, measured from its output."""
+    x = torch.randn(8, d_model)
+    x_hat = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + norm.eps)
+    with torch.no_grad():
+        return (norm(x) / x_hat).mean(0)
+
+
+@pytest.fixture(scope="module")
+def perturbed_models():
+    """Bridge + HF reference with non-identity norms (HF zero-inits the block norms)."""
+    torch.manual_seed(1)
+    cfg = MuseGlimmerConfig(
+        text_config=dict(
+            vocab_size=64,
+            hidden_size=64,
+            intermediate_size=96,
+            num_hidden_layers=2,
+            num_attention_heads=4,
+            num_key_value_heads=2,
+            head_dim=16,
+            max_position_embeddings=64,
+            sliding_window=4,
+            post_norm_eps=1e-3,
+        ),
+        vision_config=dict(
+            hidden_size=32,
+            intermediate_size=48,
+            num_hidden_layers=1,
+            num_attention_heads=2,
+            pos_emb_height=4,
+            pos_emb_width=4,
+        ),
+        out_hidden_size=128,
+        projector_hidden_size=32,
+    )
+    cfg.architectures = [ARCH]
+    hf = AutoModelForImageTextToText.from_config(cfg, attn_implementation="eager")
+    hf = hf.to(torch.float32).eval()
+    with torch.no_grad():
+        for name, param in hf.named_parameters():
+            if "norm" in name and param.ndim == 1:
+                param.add_(0.3 * torch.randn_like(param))
+    reference = copy.deepcopy(hf)
+    bridge_cfg = build_bridge_config_from_hf(hf.config, ARCH, "muse-glimmer-tiny", torch.float32)
+    bridge = TransformerBridge(hf, MuseGlimmerArchitectureAdapter(bridge_cfg), tokenizer=_Tok())
+    return bridge, reference
+
+
+def test_svd_w_in_uses_hf_ln2_scale(perturbed_models):
+    """Block norms scale by 1 + w while the final norm (and cfg) use raw w."""
+    from transformer_lens.SVDInterpreter import SVDInterpreter
+
+    bridge, _ = perturbed_models
+    interpreter = SVDInterpreter(bridge)
+    for layer in range(bridge.cfg.n_layers):
+        hf_ln2 = bridge.blocks[layer].ln2.original_component
+        scale = _hf_norm_scale(hf_ln2, bridge.cfg.d_model)
+        expected = interpreter.params[f"blocks.{layer}.mlp.W_in"].T * scale
+        torch.testing.assert_close(interpreter._get_w_in_matrix(layer), expected)
+
+
+@pytest.mark.parametrize(
+    "hook_name", ["blocks.0.ln2.hook_normalized", "blocks.1.ln1_post.hook_scale"]
+)
+def test_identity_norm_hook_edit_keeps_hf_logits(perturbed_models, hook_name):
+    """A replaced norm intermediate is rebuilt with the module's 1 + w scale and own eps."""
+    bridge, reference = perturbed_models
+    with torch.no_grad():
+        hf_logits = reference(input_ids=TOKENS).logits
+        with pytest.warns(UserWarning, match="edited hook_scale/hook_normalized"):
+            logits = bridge.run_with_hooks(
+                TOKENS, fwd_hooks=[(hook_name, lambda t, hook: t.clone())]
+            )
+    torch.testing.assert_close(logits, hf_logits, atol=1e-5, rtol=0)
+
+
+def test_post_norm_head_and_neuron_decompositions_sum_to_sublayer_outputs(perturbed_models):
+    """Per-head and per-neuron stacks route through ln*_post and reproduce attn/mlp out."""
+    bridge, _ = perturbed_models
+    n_heads, d_mlp = bridge.cfg.n_heads, bridge.cfg.d_mlp
+    with torch.no_grad():
+        _, cache = bridge.run_with_cache(TOKENS)
+        for layer in range(bridge.cfg.n_layers):
+            heads = cache.stack_head_results(layer=layer + 1)[-n_heads:].sum(0)
+            torch.testing.assert_close(
+                heads, cache[f"blocks.{layer}.hook_attn_out"], atol=1e-4, rtol=0
+            )
+            neurons = cache.stack_neuron_results(layer=layer + 1)[-d_mlp:].sum(0)
+            torch.testing.assert_close(
+                neurons, cache[f"blocks.{layer}.hook_mlp_out"], atol=1e-4, rtol=0
+            )
diff --git a/tests/unit/model_bridge/test_batched_string_padding.py b/tests/unit/model_bridge/test_batched_string_padding.py
index f945416a84..f2490d2d36 100644
--- a/tests/unit/model_bridge/test_batched_string_padding.py
+++ b/tests/unit/model_bridge/test_batched_string_padding.py
@@ -71,15 +71,19 @@ def test_ragged_forward_matches_public_layout_and_single_sequences(
         expected_tokens = compatibility_bridge.to_tokens(PROMPTS)
         attention_mask = get_attention_mask(tokenizer, expected_tokens, prepend_bos=True).bool()
         tokenization_calls: list[torch.Tensor] = []
-        original_to_tokens: Callable[..., torch.Tensor] = compatibility_bridge.to_tokens
-
-        def recording_to_tokens(input: Any, *args: Any, **kwargs: Any) -> torch.Tensor:
-            tokens = original_to_tokens(input, *args, **kwargs)
+        original_tokenize: Callable[
+            ..., tuple[torch.Tensor, torch.Tensor]
+        ] = compatibility_bridge._to_tokens_and_mask
+
+        def recording_tokenize(
+            input: Any, *args: Any, **kwargs: Any
+        ) -> tuple[torch.Tensor, torch.Tensor]:
+            tokens, mask = original_tokenize(input, *args, **kwargs)
             if isinstance(input, list):
                 tokenization_calls.append(tokens.detach().clone())
-            return tokens
+            return tokens, mask
 
-        monkeypatch.setattr(compatibility_bridge, "to_tokens", recording_to_tokens)
+        monkeypatch.setattr(compatibility_bridge, "_to_tokens_and_mask", recording_tokenize)
 
         with torch.no_grad():
             batch_logits, batch_cache = compatibility_bridge.run_with_cache(
@@ -148,7 +152,7 @@ def _explode(*args: Any, **kwargs: Any) -> Any:
 
     try:
         tokenizer.padding_side = "right"
-        monkeypatch.setattr(compatibility_bridge, "to_tokens", _explode)
+        monkeypatch.setattr(compatibility_bridge, "_to_tokens_and_mask", _explode)
 
         with pytest.raises(RuntimeError, match="tokenization failed"):
             compatibility_bridge.generate(PROMPTS, max_new_tokens=1, do_sample=False, verbose=False)
@@ -156,3 +160,103 @@ def _explode(*args: Any, **kwargs: Any) -> Any:
         assert tokenizer.padding_side == "right"
     finally:
         tokenizer.padding_side = original_side
+
+
+@pytest.mark.parametrize("prepend_bos", [True, False])
+def test_leading_literal_pad_token_matches_single_prompt(
+    compatibility_bridge, prepend_bos: bool
+) -> None:
+    """A real leading token with the pad id must stay attended, as the tokenizer marks it.
+
+    distilgpt2's pad, BOS and EOS share one id, so the token ids alone can't show
+    where a left-padded row's padding ends.
+    """
+    eos = compatibility_bridge.tokenizer.eos_token
+    prompts = [eos + " hello", "the cat sat on a mat"]
+    with torch.no_grad():
+        alone = compatibility_bridge(prompts[0], prepend_bos=prepend_bos)[0, -1]
+        logits = compatibility_bridge(prompts, prepend_bos=prepend_bos, padding_side="left")
+        cached_logits, _ = compatibility_bridge.run_with_cache(
+            prompts, prepend_bos=prepend_bos, padding_side="left", names_filter=[RESID_PRE]
+        )
+
+    torch.testing.assert_close(logits[0, -1], alone, rtol=0, atol=1e-4)
+    torch.testing.assert_close(cached_logits[0, -1], alone, rtol=0, atol=1e-4)
+
+
+def test_input_to_embed_honors_padding_side_for_strings(compatibility_bridge) -> None:
+    """A left-padded string batch must return the mask that resuming the pass needs."""
+    tokenizer = compatibility_bridge.tokenizer
+    original_side = tokenizer.padding_side
+    tokenizer.padding_side = "right"
+    try:
+        with torch.no_grad():
+            alone = compatibility_bridge(PROMPTS[0])[0, -1]
+            residual, _, _, mask = compatibility_bridge.input_to_embed(PROMPTS, padding_side="left")
+            resumed = compatibility_bridge(residual, start_at_layer=0, attention_mask=mask)
+    finally:
+        tokenizer.padding_side = original_side
+
+    torch.testing.assert_close(resumed[0, -1], alone, rtol=0, atol=1e-4)
+
+
+LLAMA = "hf-internal-testing/tiny-random-LlamaForCausalLM"
+
+
+@pytest.fixture(scope="module")
+def llama_bridge():
+    """A BOS-prepending tokenizer whose pad id differs from BOS, unlike distilgpt2's."""
+    from transformer_lens.model_bridge import TransformerBridge
+
+    bridge = TransformerBridge.boot_transformers(LLAMA, device="cpu", dtype=torch.float32)
+    assert bridge.cfg.tokenizer_prepends_bos and not bridge.cfg.tokenizer_appends_eos
+    return bridge.eval()
+
+
+@pytest.fixture(scope="module")
+def llama_eos_bridge():
+    """The same model with an EOS-appending tokenizer (OLMo-style), detected at boot."""
+    from transformers import AutoTokenizer
+
+    from transformer_lens.model_bridge import TransformerBridge
+
+    tokenizer = AutoTokenizer.from_pretrained(LLAMA, add_eos_token=True)
+    bridge = TransformerBridge.boot_transformers(
+        LLAMA, device="cpu", dtype=torch.float32, tokenizer=tokenizer
+    )
+    assert bridge.cfg.tokenizer_prepends_bos and bridge.cfg.tokenizer_appends_eos
+    return bridge.eval()
+
+
+@pytest.mark.parametrize("padding_side", ["left", "right"])
+@pytest.mark.parametrize("prepend_bos", [False, True])
+def test_bos_removal_keeps_every_pad_masked(llama_bridge, prepend_bos, padding_side) -> None:
+    """Dropping the tokenizer's BOS shifts each row by one token; the mask must shift with it."""
+    bos = llama_bridge.tokenizer.bos_token_id
+    _, tokens, _, mask = llama_bridge.input_to_embed(
+        PROMPTS, prepend_bos=prepend_bos, padding_side=padding_side
+    )
+    with torch.no_grad():
+        logits = llama_bridge(PROMPTS, prepend_bos=prepend_bos, padding_side=padding_side)
+    for row, prompt in enumerate(PROMPTS):
+        alone_tokens = llama_bridge.to_tokens(prompt, prepend_bos=prepend_bos)[0]
+        assert (alone_tokens == bos).sum().item() == int(prepend_bos)
+        keep = mask[row].bool()
+        assert torch.equal(tokens[row][keep], alone_tokens)
+        with torch.no_grad():
+            alone = llama_bridge(prompt, prepend_bos=prepend_bos)[0]
+        torch.testing.assert_close(logits[row][keep], alone, rtol=0, atol=1e-5)
+
+
+@pytest.mark.parametrize("prepend_bos", [False, True])
+def test_eos_strip_keeps_every_pad_masked(llama_eos_bridge, prepend_bos) -> None:
+    """A left-padded batch ends in EOS on every row, so the strip drops a column from both."""
+    eos = llama_eos_bridge.tokenizer.eos_token_id
+    _, tokens, _, mask = llama_eos_bridge.input_to_embed(
+        PROMPTS, prepend_bos=prepend_bos, padding_side="left"
+    )
+    assert not (tokens[:, -1] == eos).any()
+    for row, prompt in enumerate(PROMPTS):
+        alone_tokens = llama_eos_bridge.to_tokens(prompt, prepend_bos=prepend_bos)[0]
+        assert alone_tokens[-1].item() != eos
+        assert torch.equal(tokens[row][mask[row].bool()], alone_tokens)
diff --git a/tests/unit/model_bridge/test_bridge_generate_bos.py b/tests/unit/model_bridge/test_bridge_generate_bos.py
index 8af1fb8bf2..bab13c2eeb 100644
--- a/tests/unit/model_bridge/test_bridge_generate_bos.py
+++ b/tests/unit/model_bridge/test_bridge_generate_bos.py
@@ -321,20 +321,16 @@ def test_literal_bos_stays_attended(self, distilgpt2_bridge, stream):
         prompts = [bos + prompt for prompt in self.PROMPTS]
         self._assert_rows_match_single_prompts(distilgpt2_bridge, stream, prompts, False)
 
-    def test_explicit_mask_skips_derived_mask(self, distilgpt2_bridge, stream, monkeypatch):
-        """The loop ignores the derived mask when one is passed, so it must not be built."""
+    def test_explicit_mask_wins(self, distilgpt2_bridge, stream):
+        """An all-ones mask attends the pads, so the batch runs as its bare tensor does."""
         bridge = distilgpt2_bridge
-        expected = self._run(bridge, stream, self.PROMPTS, False)
         tokens = bridge.to_tokens(
             self.PROMPTS, prepend_bos=False, padding_side="left", truncate=False
         )
-        mask = bridge._left_padded_prompt_mask(self.PROMPTS, tokens, False)
+        unmasked = self._run(bridge, stream, tokens, False)
+        assert not torch.equal(unmasked, self._run(bridge, stream, self.PROMPTS, False))
 
-        def _fail(*args, **kwargs):
-            raise AssertionError("_left_padded_prompt_mask ran under an explicit attention_mask")
-
-        monkeypatch.setattr(bridge, "_left_padded_prompt_mask", _fail)
-        kwargs = {**self.KWARGS, "attention_mask": mask}
+        kwargs = {**self.KWARGS, "attention_mask": torch.ones_like(tokens)}
         if stream:
             chunks = bridge.generate_stream(
                 self.PROMPTS, prepend_bos=False, max_tokens_per_yield=100, **kwargs
@@ -342,15 +338,48 @@ def _fail(*args, **kwargs):
             out = torch.cat(list(chunks), dim=-1)
         else:
             out = bridge.generate(self.PROMPTS, prepend_bos=False, **kwargs)
-        assert torch.equal(out, expected)
+        assert torch.equal(out, unmasked)
 
 
-def test_empty_prompt_attends_only_its_last_pad(distilgpt2_bridge):
-    prompts = ["", "Hello there friend"]
-    tokens = distilgpt2_bridge.to_tokens(
-        prompts, prepend_bos=False, padding_side="left", truncate=False
+@pytest.mark.parametrize(
+    "padding_side, expected",
+    [("left", [[0, 0, 1], [1, 1, 1]]), ("right", [[1, 0, 0], [1, 1, 1]])],
+)
+def test_empty_prompt_attends_one_pad(distilgpt2_bridge, padding_side, expected):
+    """A row with no real token keeps one slot attended; an all-masked row softmaxes to NaN."""
+    _, mask = distilgpt2_bridge._to_tokens_and_mask(
+        ["", "Hello there friend"], prepend_bos=False, padding_side=padding_side
     )
 
-    mask = distilgpt2_bridge._left_padded_prompt_mask(prompts, tokens, False)
+    assert mask.tolist() == expected
+
+
+@pytest.mark.parametrize("prepend_bos", [True, False])
+def test_forward_and_generate_keep_a_literal_pad_prefix_as_context(distilgpt2_bridge, prepend_bos):
+    """A leading literal <|endoftext|> is real context to both forward and generate.
 
-    assert mask.tolist() == [[0, 0, 1], [1, 1, 1]]
+    Each batched path is compared with the prompt run alone: a mask rebuilt from the ids
+    would break forward and generate in the same way, so comparing them with each other
+    could not see it.
+    """
+    bridge = distilgpt2_bridge
+    prompts = [
+        bridge.tokenizer.eos_token + " The capital of France is",
+        "the cat sat on a mat and then went",
+    ]
+    generate_kwargs = {
+        "max_new_tokens": 3,
+        "do_sample": False,
+        "stop_at_eos": False,
+        "prepend_bos": prepend_bos,
+        "return_type": "tokens",
+        "verbose": False,
+    }
+    with torch.no_grad():
+        alone_logits = bridge(prompts[0], prepend_bos=prepend_bos)[0, -1]
+        alone_tokens = bridge.generate(prompts[0], **generate_kwargs)[0]
+        logits = bridge(prompts, prepend_bos=prepend_bos, padding_side="left")
+        generated = bridge.generate(prompts, **generate_kwargs)
+
+    torch.testing.assert_close(logits[0, -1], alone_logits, rtol=0, atol=1e-4)
+    assert torch.equal(generated[0, -alone_tokens.shape[-1] :], alone_tokens)
diff --git a/tests/unit/model_bridge/test_execution_kwargs.py b/tests/unit/model_bridge/test_execution_kwargs.py
index d6a306787a..d8bf739ea3 100644
--- a/tests/unit/model_bridge/test_execution_kwargs.py
+++ b/tests/unit/model_bridge/test_execution_kwargs.py
@@ -86,8 +86,10 @@ def unexpected_forward(**kwargs):
 
     monkeypatch.setattr(bridge._driver, "forward", unexpected_forward)
     with pytest.raises(TypeError, match="fwd_hooks"):
-        bridge(_tokens(), fwd_hooks=[], stop_at_layer=1, start_at_layer=0)
+        bridge(_tokens(), fwd_hooks=[], stop_at_layer=0, start_at_layer=0)
+    # Arming either leaves the attribute behind (teardown resets it to None, not absent).
     assert not hasattr(bridge.blocks[0], "_stop_at_layer_idx")
+    assert not hasattr(bridge.blocks[0], "_start_residual")
 
 
 @torch.no_grad()
diff --git a/tests/unit/model_bridge/test_stop_at_layer_guard.py b/tests/unit/model_bridge/test_stop_at_layer_guard.py
index e0e5db99a0..d4e5007ed7 100644
--- a/tests/unit/model_bridge/test_stop_at_layer_guard.py
+++ b/tests/unit/model_bridge/test_stop_at_layer_guard.py
@@ -1,11 +1,21 @@
+import copy
 from types import SimpleNamespace
+from typing import Any, Callable, Dict
 
 import pytest
 import torch
 import torch.nn as nn
+import transformers
 
 from transformer_lens.config import TransformerBridgeConfig
 from transformer_lens.model_bridge import TransformerBridge
+from transformer_lens.model_bridge.sources._bridge_builder import (
+    build_bridge_from_module,
+)
+
+N_LAYERS = 3
+D_MODEL = 16
+TOKENS = torch.tensor([[3, 4, 5, 6, 7]])
 
 
 def _bare_bridge(**block_lists: nn.Module) -> TransformerBridge:
@@ -48,3 +58,172 @@ def test_blocks_guard_ignores_unregistered_blocks_attribute() -> None:
         bridge.forward(torch.zeros(1, 3, dtype=torch.long), stop_at_layer=0)
     with pytest.raises(NotImplementedError, match="start_at_layer requires a 'blocks' stack"):
         bridge.forward(torch.zeros(1, 3, 8), start_at_layer=0)
+
+
+def _tiny_bridge(architecture: str, cfg: Any, model_cls: type) -> TransformerBridge:
+    cfg._attn_implementation = "eager"
+    cfg.architectures = [architecture]
+    torch.manual_seed(0)
+    model = model_cls(cfg).eval()
+    bridge = build_bridge_from_module(
+        model, architecture, hf_config=copy.deepcopy(cfg), tokenizer=None, device="cpu"
+    )
+    return bridge.eval()
+
+
+_BUILDERS: Dict[str, Callable[[], TransformerBridge]] = {
+    "gpt2": lambda: _tiny_bridge(
+        "GPT2LMHeadModel",
+        transformers.GPT2Config(
+            vocab_size=64, n_embd=D_MODEL, n_layer=N_LAYERS, n_head=2, n_positions=32
+        ),
+        transformers.GPT2LMHeadModel,
+    ),
+    "mamba": lambda: _tiny_bridge(
+        "MambaForCausalLM",
+        transformers.MambaConfig(
+            vocab_size=64,
+            hidden_size=D_MODEL,
+            state_size=8,
+            num_hidden_layers=N_LAYERS,
+            expand=2,
+            time_step_rank=4,
+            conv_kernel=4,
+        ),
+        transformers.MambaForCausalLM,
+    ),
+    "bloom": lambda: _tiny_bridge(
+        "BloomForCausalLM",
+        transformers.BloomConfig(vocab_size=64, hidden_size=D_MODEL, n_layer=N_LAYERS, n_head=2),
+        transformers.BloomForCausalLM,
+    ),
+    "bert": lambda: _tiny_bridge(
+        "BertForMaskedLM",
+        transformers.BertConfig(
+            vocab_size=64,
+            hidden_size=D_MODEL,
+            num_hidden_layers=N_LAYERS,
+            num_attention_heads=2,
+            intermediate_size=32,
+            max_position_embeddings=32,
+        ),
+        transformers.BertForMaskedLM,
+    ),
+    "roberta": lambda: _tiny_bridge(
+        "RobertaForMaskedLM",
+        transformers.RobertaConfig(
+            vocab_size=64,
+            hidden_size=D_MODEL,
+            num_hidden_layers=N_LAYERS,
+            num_attention_heads=2,
+            intermediate_size=32,
+            max_position_embeddings=32,
+        ),
+        transformers.RobertaForMaskedLM,
+    ),
+    "t5": lambda: _tiny_bridge(
+        "T5ForConditionalGeneration",
+        transformers.T5Config(
+            vocab_size=64, d_model=D_MODEL, d_kv=8, d_ff=32, num_layers=2, num_heads=2
+        ),
+        transformers.T5ForConditionalGeneration,
+    ),
+}
+
+
+@pytest.fixture(scope="module")
+def tiny_bridge(request: pytest.FixtureRequest) -> TransformerBridge:
+    return _BUILDERS[request.param]()
+
+
+def _stop_with(bridge: TransformerBridge, api: str, k: int) -> torch.Tensor:
+    if api == "forward":
+        return bridge(TOKENS, stop_at_layer=k)
+    if api == "run_with_hooks":
+        return bridge.run_with_hooks(TOKENS, stop_at_layer=k)
+    out, _ = bridge.run_with_cache(TOKENS, stop_at_layer=k)
+    return out
+
+
+@pytest.mark.parametrize("tiny_bridge", ["t5"], indirect=True)
+def test_run_with_cache_raises_without_blocks_stack(tiny_bridge: TransformerBridge) -> None:
+    """Encoder/decoder stacks must not let run_with_cache ignore stop_at_layer and return logits."""
+    with pytest.raises(NotImplementedError, match="stop_at_layer requires a 'blocks' stack"):
+        tiny_bridge.run_with_cache(TOKENS, stop_at_layer=1, decoder_input=TOKENS)
+    leaked = [name for name, hp in tiny_bridge.hook_dict.items() if hp.fwd_hooks or hp.bwd_hooks]
+    assert leaked == []
+    assert getattr(tiny_bridge, "context_level", 0) == 0
+
+
+@pytest.mark.parametrize("k", [0, 1, -1])
+@pytest.mark.parametrize("api", ["forward", "run_with_hooks", "run_with_cache"])
+@pytest.mark.parametrize("tiny_bridge", ["bert", "roberta"], indirect=True)
+def test_stop_at_layer_stops_on_encoder_layer_names(
+    tiny_bridge: TransformerBridge, api: str, k: int
+) -> None:
+    """``*.encoder.layer.N`` blocks stop at the residual entering block k, as on decoder stacks."""
+    hook_name = f"blocks.{k % N_LAYERS}.hook_in"
+    with torch.no_grad():
+        _, cache = tiny_bridge.run_with_cache(TOKENS, names_filter=hook_name)
+        out = _stop_with(tiny_bridge, api, k)
+    assert out.shape == (1, TOKENS.shape[1], D_MODEL)
+    torch.testing.assert_close(out, cache[hook_name])
+
+
+@pytest.mark.parametrize("tiny_bridge", ["bert", "roberta"], indirect=True)
+def test_input_to_embed_on_bert_returns_residual(tiny_bridge: TransformerBridge) -> None:
+    """input_to_embed must return resid_pre[0], not MLM logits."""
+    with torch.no_grad():
+        residual, _, _, _ = tiny_bridge.input_to_embed(TOKENS)
+        _, cache = tiny_bridge.run_with_cache(TOKENS, names_filter="blocks.0.hook_in")
+    assert residual.shape == (1, TOKENS.shape[1], D_MODEL)
+    torch.testing.assert_close(residual, cache["blocks.0.hook_in"])
+
+
+@pytest.mark.parametrize("start", [1, 2])
+@pytest.mark.parametrize("tiny_bridge", ["bert", "roberta", "bloom", "mamba"], indirect=True)
+def test_start_at_layer_resumes_on_every_block_class(
+    tiny_bridge: TransformerBridge, start: int
+) -> None:
+    """Resuming from a cached ``hook_in`` reproduces the full pass on encoder-layer names,
+    Bloom's alibi block and the opaque SSM block, each of which runs its own forward."""
+    with torch.no_grad():
+        full, cache = tiny_bridge.run_with_cache(TOKENS)
+        resumed = tiny_bridge(cache[f"blocks.{start}.hook_in"], start_at_layer=start)
+    torch.testing.assert_close(resumed, full, atol=1e-5, rtol=0)
+
+
+@pytest.mark.parametrize("k", [0, 1, N_LAYERS, -1])
+@pytest.mark.parametrize("tiny_bridge", ["gpt2", "mamba"], indirect=True)
+def test_run_with_cache_stop_parity_gpt2_and_mamba(tiny_bridge: TransformerBridge, k: int) -> None:
+    """run_with_cache's stop returns the stopped residual, caches nothing past it, fires nothing later."""
+    stop = k % N_LAYERS if k < 0 else k
+    expected_name = (
+        f"blocks.{stop}.hook_in" if stop < N_LAYERS else f"blocks.{N_LAYERS - 1}.hook_out"
+    )
+    fired: list = []
+
+    def record(tensor: torch.Tensor, hook: Any) -> None:
+        fired.append(hook.name)
+
+    def block_idx(name: str) -> int:
+        return int(name.split(".")[1]) if name.startswith("blocks.") else -1
+
+    with torch.no_grad():
+        _, full_cache = tiny_bridge.run_with_cache(TOKENS)
+        # Component boundaries only: gated hooks (hook_result, ...) are off and would warn.
+        with tiny_bridge.hooks(
+            fwd_hooks=[(lambda n: n.endswith((".hook_in", ".hook_out")), record)]
+        ):
+            out, cache = tiny_bridge.run_with_cache(TOKENS, stop_at_layer=k)
+
+    torch.testing.assert_close(out, full_cache[expected_name])
+    assert [name for name in cache.keys() if block_idx(name) >= stop] == []
+    late = [
+        name
+        for name in fired
+        if (block_idx(name) >= stop and name != f"blocks.{stop}.hook_in")
+        or name.startswith(("unembed.", "ln_final."))
+    ]
+    assert late == []
+    assert expected_name in fired
diff --git a/tests/unit/test_activation_cache.py b/tests/unit/test_activation_cache.py
index 3bf73a4157..9c977188a3 100644
--- a/tests/unit/test_activation_cache.py
+++ b/tests/unit/test_activation_cache.py
@@ -890,3 +890,36 @@ def test_layer_zero_ln_projected_stack_follows_float32_scales_of_a_bf16_model(
     ]
     assert projected[0].shape[0] == 0
     assert projected[0].dtype == projected[1].dtype
+
+
+def _single_block_cache(hook_mlp_out: bool) -> ActivationCache:
+    """One sequential attn + MoE block; ``hook_mlp_out`` registered and cached only on request."""
+    torch.manual_seed(0)
+    cache_dict = {
+        "hook_embed": torch.randn(2, 4, 8),
+        "blocks.0.hook_attn_out": torch.randn(2, 4, 8),
+        "blocks.0.hook_resid_mid": torch.randn(2, 4, 8),
+        "blocks.0.hook_resid_post": torch.randn(2, 4, 8),
+    }
+    if hook_mlp_out:
+        cache_dict["blocks.0.hook_mlp_out"] = torch.randn(2, 4, 8)
+    block = SimpleNamespace(_modules={"attn": object(), "moe": object()})
+    model = SimpleNamespace(
+        cfg=SimpleNamespace(n_layers=1, attn_only=False),
+        blocks=[block],
+        hook_dict={name: object() for name in cache_dict},
+    )
+    return ActivationCache(cache_dict, model)
+
+
+def test_decompose_resid_derives_mlp_out_from_the_residual_when_the_block_has_no_hook():
+    cache = _single_block_cache(hook_mlp_out=False)
+    components, labels = cache.decompose_resid(return_labels=True)
+    assert labels == ["embed", "0_attn_out", "0_mlp_out"]
+    expected = cache["blocks.0.hook_resid_post"] - cache["blocks.0.hook_resid_mid"]
+    assert torch.equal(components[2], expected)
+
+
+def test_decompose_resid_reads_hook_mlp_out_when_the_block_has_one():
+    cache = _single_block_cache(hook_mlp_out=True)
+    assert torch.equal(cache.decompose_resid()[2], cache["blocks.0.hook_mlp_out"])
diff --git a/tests/unit/test_patching.py b/tests/unit/test_patching.py
index c7af54d3c4..37bf670be6 100644
--- a/tests/unit/test_patching.py
+++ b/tests/unit/test_patching.py
@@ -359,24 +359,123 @@ def test_stacked_by_pos_dataframe_labels_match_the_direct_pattern_sweep(batched)
         assert row.value.item() == pytest.approx(direct.values[layer, head, pos, example].item())
 
 
-def test_stacked_gqa_dataframe_has_zero_rows_for_padded_heads():
-    from transformer_lens.patching import get_act_patch_attn_head_all_pos_every
-
+@pytest.fixture(scope="module")
+def gqa():
+    """A batch-3 clean/corrupt pair on a tiny model with two key/value heads for four queries."""
     model = _tiny_model(n_key_value_heads=2)
     torch.manual_seed(0)
     clean = torch.randint(0, model.cfg.d_vocab, (3, 5))
     torch.manual_seed(1)
     corrupt = torch.randint(0, model.cfg.d_vocab, (3, 5))
-    _, cache = model.run_with_cache(clean)
+    clean_logits, cache = model.run_with_cache(clean)
+    return model, clean_logits, corrupt, cache
+
+
+@pytest.fixture(scope="module")
+def gqa_result(gqa):
+    from transformer_lens.patching import get_act_patch_attn_head_all_pos_every
+
+    model, _, corrupt, cache = gqa
+    return get_act_patch_attn_head_all_pos_every(
+        model, corrupt, cache, _per_example, reduce="none", baseline=True
+    )
+
+
+def _assert_nan_only_on_padded(t, kv):
+    padded = torch.zeros(t.shape[:3], dtype=torch.bool)
+    padded[2:4, :, kv:] = True
+    assert t[padded].isnan().all()
+    assert t[~padded].isfinite().all()
+
+
+def test_gqa_padded_cells_are_nan_in_effect_and_normalized(gqa, gqa_result):
+    _, clean_logits, _, _ = gqa
+    effect = gqa_result.effect()
+    _assert_nan_only_on_padded(effect, kv=2)
+    live = gqa_result.values - gqa_result.baseline
+    torch.testing.assert_close(effect[:2], live[:2])
+    torch.testing.assert_close(effect[2:4, :, :2], live[2:4, :, :2])
+    normalized = gqa_result.normalized(_per_example(clean_logits).detach())
+    _assert_nan_only_on_padded(normalized, kv=2)
+
+
+def test_gqa_padded_cells_are_nan_in_every_statistic(gqa_result):
+    _assert_nan_only_on_padded(gqa_result.permutation_pvalue(), kv=2)
+    _assert_nan_only_on_padded(gqa_result.stderr(), kv=2)
+    for of_effect in (False, True):
+        low, high = gqa_result.bootstrap_ci(of_effect=of_effect, n_resamples=50)
+        _assert_nan_only_on_padded(low, kv=2)
+        _assert_nan_only_on_padded(high, kv=2)
+
+
+def test_gqa_valid_mask_marks_exactly_the_padded_cells(gqa, gqa_result, batched):
+    from transformer_lens.patching import (
+        get_act_patch_attn_head_all_pos_every,
+        get_act_patch_attn_head_by_pos_every,
+        get_act_patch_block_every,
+    )
+
+    model = gqa[0]
+    expected = torch.ones(5, model.cfg.n_layers, model.cfg.n_heads, dtype=torch.bool)
+    expected[2:4, :, 2:] = False
+    assert gqa_result.valid is not None and torch.equal(gqa_result.valid, expected)
+
+    mqa = _tiny_model(n_key_value_heads=1)
+    torch.manual_seed(0)
+    clean = torch.randint(0, mqa.cfg.d_vocab, (2, 4))
+    torch.manual_seed(1)
+    corrupt = torch.randint(0, mqa.cfg.d_vocab, (2, 4))
+    _, cache = mqa.run_with_cache(clean)
+    by_pos = get_act_patch_attn_head_by_pos_every(mqa, corrupt, cache, _per_example, reduce="none")
+    expected_pos = torch.ones(5, mqa.cfg.n_layers, 4, mqa.cfg.n_heads, dtype=torch.bool)
+    expected_pos[2:4, :, :, 1:] = False
+    assert by_pos.valid is not None and torch.equal(by_pos.valid, expected_pos)
+
+    model, _, corrupt, cache = batched
+    full = get_act_patch_attn_head_all_pos_every(model, corrupt, cache, _per_example, reduce="none")
+    assert full.valid is None
+    block = get_act_patch_block_every(model, corrupt, cache, _per_example, reduce="none")
+    assert block.valid is None
+
+
+def test_gqa_mask_survives_metric_selection_and_reduced_baseline(gqa):
+    from transformer_lens.patching import get_act_patch_attn_head_all_pos_every
+
+    model, _, corrupt, cache = gqa
+    multi = get_act_patch_attn_head_all_pos_every(
+        model,
+        corrupt,
+        cache,
+        {"a": _per_example, "b": lambda logits: logits[:, -1, 5]},
+        reduce="none",
+        baseline=True,
+    )
+    _assert_nan_only_on_padded(multi["b"].effect(), kv=2)
+    # The documented recipe: the ellipsis skips the metric axis the mask does not have.
+    p = multi.permutation_pvalue()
+    assert multi.valid is not None
+    live = p[..., multi.valid]
+    assert live.shape == (2, int(multi.valid.sum()))
+    assert live.isfinite().all()
+    reduced = get_act_patch_attn_head_all_pos_every(model, corrupt, cache, _scalar, baseline=True)
+    assert reduced.reduced
+    _assert_nan_only_on_padded(reduced.effect(), kv=2)
+
+
+def test_stacked_gqa_dataframe_omits_padded_heads(gqa):
+    from transformer_lens.patching import get_act_patch_attn_head_all_pos_every
+
+    model, _, corrupt, cache = gqa
     stacked = get_act_patch_attn_head_all_pos_every(
         model, corrupt, cache, _per_example, reduce="none"
     )
     frame = stacked.to_dataframe()
-    assert len(frame) == 5 * model.cfg.n_layers * model.cfg.n_heads * 3
+    n_layers, n_heads = model.cfg.n_layers, model.cfg.n_heads
+    assert len(frame) == (5 * n_layers * n_heads - 2 * n_layers * (n_heads - 2)) * 3
     padded = frame[(frame.patch_type_name.isin(["k", "v"])) & (frame["head"] >= 2)]
-    assert len(padded) == 2 * model.cfg.n_layers * 2 * 3 and (padded.value == 0).all()
+    assert len(padded) == 0
     live = frame[(frame.patch_type_name == "k") & (frame["head"] < 2)]
-    assert (live.value != 0).any()
+    assert len(live) == n_layers * 2 * 3 and (live.value != 0).any()
 
 
 def test_every_helpers_compute_the_baseline_once(batched, monkeypatch):
diff --git a/tests/unit/tools/test_direct_logit_attribution.py b/tests/unit/tools/test_direct_logit_attribution.py
index 81fc6a024e..de869fefea 100644
--- a/tests/unit/tools/test_direct_logit_attribution.py
+++ b/tests/unit/tools/test_direct_logit_attribution.py
@@ -1,11 +1,12 @@
 """Unit tests for Direct Logit Attribution guards and argument validation.
 
 These exercise the fast-failing checks (argument validation, the Bridge
-compatibility-mode requirement, and the hybrid-architecture refusal) without
+compatibility-mode requirement, and the structural refusals) without
 loading a real model — using a ``spec``-ed mock TransformerBridge so the checks
 fire before any forward pass.
 """
 
+from types import SimpleNamespace
 from unittest.mock import MagicMock
 
 import pytest
@@ -18,15 +19,17 @@
     LinearBridge,
 )
 from transformer_lens.tools.analysis import direct_logit_attribution
+from transformer_lens.tools.analysis.direct_logit_attribution import (
+    _require_every_writer_decomposed,
+)
 
 TOKENS = torch.tensor([[1, 2, 3, 4]])
 
 
-def _mock_bridge(compatibility_mode=True, layer_types=("attn+mlp",)):
+def _mock_bridge(compatibility_mode=True):
     """A mock TransformerBridge that satisfies isinstance checks."""
     bridge = MagicMock(spec=TransformerBridge)
     bridge.compatibility_mode = compatibility_mode
-    bridge.layer_types.return_value = list(layer_types)
     return bridge
 
 
@@ -48,10 +51,41 @@ def test_requires_compatibility_mode():
         direct_logit_attribution(bridge, "hi", answer_tokens=" world")
 
 
-def test_rejects_hybrid_architecture():
-    bridge = _mock_bridge(layer_types=("attn+mlp", "mamba+mlp"))
-    with pytest.raises(NotImplementedError, match="hybrid"):
-        direct_logit_attribution(bridge, "hi", answer_tokens=" world")
+def test_layer_unit_refuses_blocks_without_resid_mid_before_any_forward():
+    bridge = _mock_bridge()
+    bridge.blocks = [SimpleNamespace(), SimpleNamespace()]
+    bridge.hook_dict = {"blocks.0.hook_resid_mid": object()}
+    with pytest.raises(ValueError, match=r"unit='layer' .* blocks \[1\]"):
+        direct_logit_attribution(bridge, "hi", answer_tokens=" world", unit="layer")
+    bridge.run_with_cache.assert_not_called()
+
+
+def _model_over_blocks(*slots_per_block):
+    blocks = [
+        SimpleNamespace(_modules={slot: object() for slot in slots}) for slots in slots_per_block
+    ]
+    return SimpleNamespace(
+        blocks=blocks, cfg=SimpleNamespace(n_layers=len(blocks), attn_only=False)
+    )
+
+
+def test_component_unit_refuses_a_mixer_decompose_resid_would_not_emit():
+    with pytest.raises(ValueError, match="block 1: its mamba output is not decomposed yet"):
+        _require_every_writer_decomposed(_model_over_blocks(("attn",), ("attn", "mamba")))
+    _require_every_writer_decomposed(_model_over_blocks(("attn",), ("mamba",)))
+    _require_every_writer_decomposed(_model_over_blocks(("linear_attn",), ("attn",)))
+
+
+def test_component_unit_refuses_a_parallel_mixer_before_any_forward():
+    bridge = _mock_bridge()
+    bridge.cfg = SimpleNamespace(n_layers=2, attn_only=False)
+    bridge.blocks = [
+        SimpleNamespace(_modules={"attn": object()}),
+        SimpleNamespace(_modules={"attn": object(), "mamba": object()}),
+    ]
+    with pytest.raises(ValueError, match="block 1: its mamba output is not decomposed yet"):
+        direct_logit_attribution(bridge, "hi", answer_tokens=" world", unit="component")
+    bridge.run_with_cache.assert_not_called()
 
 
 def test_direct_logit_attribution_is_exported_from_analysis_package():
diff --git a/tests/unit/tools/test_jacobian_lens.py b/tests/unit/tools/test_jacobian_lens.py
index dd382d24ce..432c5fa54a 100644
--- a/tests/unit/tools/test_jacobian_lens.py
+++ b/tests/unit/tools/test_jacobian_lens.py
@@ -46,6 +46,7 @@
     JSpaceOccupancy,
     JSpaceVarianceProfile,
 )
+from transformer_lens.tools.analysis.logit_lens import _READOUT_TRANSFORM_COMPONENTS
 from transformer_lens.utilities.activation_functions import apply_softcap
 from transformer_lens.utilities.parameter_swap import temporarily_swap_parameter
 
@@ -695,34 +696,39 @@ def test_validate_model_rejects_unsupported_final_output_paths() -> None:
         _lens().validate_model(projected_unembed)
 
     projected_output = _ToyBridge()
-    projected_output.adapter.get_component_mapping = lambda: {
-        "blocks": SimpleNamespace(hook_out_is_single_residual_stream=True),
-        "ln_final": object(),
-        "project_out": object(),
-        "unembed": object(),
-    }
-    with pytest.raises(ValueError, match="final output projection"):
+    projected_output.add_module("project_out", nn.Identity())
+    with pytest.raises(ValueError, match=r"mounts \['project_out'\]"):
         _lens().validate_model(projected_output)
 
     altup_model = _ToyBridge()
-    altup_model.adapter.get_component_mapping = lambda: {
-        "blocks": AltUpBlockBridge(name="blocks"),
-        "ln_final": object(),
-        "unembed": object(),
-    }
-    with pytest.raises(ValueError, match="single-stream.*AltUpBlockBridge"):
+    altup_model.blocks[0] = AltUpBlockBridge(name="blocks.0")
+    with pytest.raises(ValueError, match=r"single-stream.*blocks\.0 \(AltUpBlockBridge\)"):
         _lens().validate_model(altup_model)
 
     mhc_model = _ToyBridge()
-    mhc_model.adapter.get_component_mapping = lambda: {
-        "blocks": DeepseekV4BlockBridge(name="blocks"),
-        "ln_final": object(),
-        "unembed": object(),
-    }
-    with pytest.raises(ValueError, match="single-stream.*DeepseekV4BlockBridge"):
+    mhc_model.blocks[0] = DeepseekV4BlockBridge(name="blocks.0")
+    with pytest.raises(ValueError, match=r"single-stream.*blocks\.0 \(DeepseekV4BlockBridge\)"):
         _lens().validate_model(mhc_model)
 
 
+@pytest.mark.parametrize("component", _READOUT_TRANSFORM_COMPONENTS)
+def test_mounted_readout_transform_is_rejected_before_fit(component: str) -> None:
+    model = _ToyBridge()
+    model.add_module(component, nn.Identity().eval())
+    match = rf"JacobianLens .* mounts \['{component}'\]"
+
+    with pytest.raises(ValueError, match=match):
+        _lens().validate_model(model)
+    with pytest.raises(ValueError, match=match):
+        JacobianLens.fit(
+            model,
+            ["a toy prompt"],
+            corpus=CORPUS,
+            skip_first_positions=SKIP_FIRST,
+            show_progress=False,
+        )
+
+
 def test_residual_runtime_shape_guard() -> None:
     with pytest.raises(ValueError, match=r"blocks.0.hook_out.*\(2, 1, 3, 6\)"):
         jacobian_lens_module._validate_residual_activation(
diff --git a/tests/unit/tools/test_logit_lens.py b/tests/unit/tools/test_logit_lens.py
index 4c65abdf91..12d43e31d7 100644
--- a/tests/unit/tools/test_logit_lens.py
+++ b/tests/unit/tools/test_logit_lens.py
@@ -150,6 +150,31 @@ def test_incl_mid_adds_mid_entries(tiny):
     torch.testing.assert_close(result.values[-1], ref, atol=1e-5, rtol=1e-5)
 
 
+def test_incl_mid_skips_parallel_blocks_without_resid_mid():
+    """Parallel blocks have no resid_mid by design, so the resid_mid check must not refuse them."""
+    from transformers import GPTNeoXConfig, GPTNeoXForCausalLM
+
+    torch.manual_seed(0)
+    cfg = GPTNeoXConfig(
+        vocab_size=D_VOCAB,
+        hidden_size=D_MODEL,
+        intermediate_size=64,
+        num_hidden_layers=2,
+        num_attention_heads=4,
+        max_position_embeddings=32,
+        use_parallel_residual=True,
+    )
+    hf = GPTNeoXForCausalLM(cfg).eval()
+    model = build_bridge_from_module(
+        hf, "GPTNeoXForCausalLM", hf_config=cfg, dtype=torch.float32, device="cpu"
+    )
+    with torch.no_grad():
+        ref = model(TOKENS)
+    result = logit_lens(model, TOKENS, incl_mid=True, positions=None)
+    assert result.labels == ["0_pre", "1_pre", "final_post"]
+    torch.testing.assert_close(result.values[-1], ref, atol=1e-5, rtol=1e-5)
+
+
 def test_positions_and_batch_slices(tiny):
     model, ref, cache = tiny
     last = logit_lens(cache)  # default positions=-1
diff --git a/tests/unit/tools/test_model_state.py b/tests/unit/tools/test_model_state.py
index d63c06f64d..048a2c9414 100644
--- a/tests/unit/tools/test_model_state.py
+++ b/tests/unit/tools/test_model_state.py
@@ -8,8 +8,9 @@
 import torch
 
 from transformer_lens.tools.analysis._model_state import (
-    hybrid_block_types,
+    blocks_without_resid_mid,
     require_int,
+    require_resid_mid,
     to_token_ids,
 )
 
@@ -49,12 +50,21 @@ def test_require_int_message_replaces_both_failures(value: object) -> None:
         require_int(value, "k", message="ks must hold integers")
 
 
-def test_hybrid_block_types_are_sorted_distinct_and_default_to_none() -> None:
+def test_blocks_without_resid_mid_exempts_parallel_blocks() -> None:
     model = SimpleNamespace(
-        layer_types=lambda: ["mamba+mlp", "attn+mlp", "linear_attn", "mamba+mlp"]
+        blocks=[
+            SimpleNamespace(),
+            SimpleNamespace(mlp_input_includes_attn=False),
+            SimpleNamespace(),
+        ],
+        hook_dict={"blocks.0.hook_resid_mid": object()},
     )
-    assert hybrid_block_types(model) == ["linear_attn", "mamba+mlp"]
-    assert hybrid_block_types(SimpleNamespace()) == []
+    assert blocks_without_resid_mid(model) == [2]
+    assert blocks_without_resid_mid(SimpleNamespace()) == []
+    with pytest.raises(ValueError, match=r"^op reads hook_resid_mid .* blocks \[2\] .*; do x$"):
+        require_resid_mid(model, operation="op", alternative="do x")
+    model.hook_dict["blocks.2.hook_resid_mid"] = object()
+    require_resid_mid(model, operation="op", alternative="do x")
 
 
 def _vocab_model(d_vocab: int = 10) -> SimpleNamespace:
diff --git a/tests/unit/utilities/test_statistics.py b/tests/unit/utilities/test_statistics.py
index 08a7a44fca..e65ef93347 100644
--- a/tests/unit/utilities/test_statistics.py
+++ b/tests/unit/utilities/test_statistics.py
@@ -12,6 +12,7 @@
     derive_generator,
     sign_flip_permutation_pvalue,
     standard_error,
+    statistics,
     success_rate_ci,
 )
 from transformer_lens.utilities.statistics import _binomial_cdf, _count_hits
@@ -183,6 +184,65 @@ def test_sign_flip_exact_enumeration_matches_hand_count():
     assert greater < less
 
 
+def test_sign_flip_exact_branch_evaluates_bounded_sign_chunks(monkeypatch):
+    null_widths = []
+
+    def spy(null, observed, alternative):
+        null_widths.append(null.shape[-1])
+        return _count_hits(null, observed, alternative)
+
+    monkeypatch.setattr(statistics, "_count_hits", spy)
+    sign_flip_permutation_pvalue(torch.randn(3, 12))
+    assert sum(null_widths) == 2**12
+    assert max(null_widths) <= 256
+
+
+def _unchunked_sign_flip_pvalue(effects: torch.Tensor, alternative: str) -> torch.Tensor:
+    n = effects.shape[-1]
+    signs = torch.tensor(list(itertools.product((1.0, -1.0), repeat=n)), dtype=effects.dtype)
+    null = (effects.unsqueeze(-2) * signs).mean(dim=-1)
+    observed = effects.mean(dim=-1, keepdim=True)
+    return _count_hits(null, observed, alternative) / 2**n
+
+
+@pytest.mark.parametrize("chunk", [1, 7])
+@pytest.mark.parametrize("n", [1, 5, 9, 12])
+@pytest.mark.parametrize("kind", ["tied", "randn"])
+@pytest.mark.parametrize("alternative", ["two-sided", "greater", "less"])
+def test_sign_flip_exact_chunking_matches_unchunked_enumeration(
+    monkeypatch, chunk: int, n: int, kind: str, alternative: str
+):
+    generator = torch.Generator().manual_seed(n)
+    if kind == "tied":
+        effects = torch.randint(-3, 4, (4, n), generator=generator).float()
+    else:
+        effects = torch.randn(4, n, generator=generator)
+    monkeypatch.setattr(statistics, "_SIGN_FLIP_CHUNK", chunk)
+    p = sign_flip_permutation_pvalue(effects, alternative=alternative)
+    assert torch.equal(p, _unchunked_sign_flip_pvalue(effects, alternative))
+
+
+@pytest.mark.parametrize("exact_max_n", [12, 3])
+@pytest.mark.parametrize("bad", [float("nan"), float("inf"), float("-inf")])
+def test_sign_flip_returns_nan_for_cells_with_non_finite_effects(exact_max_n: int, bad: float):
+    effects = torch.tensor([[1.0, 2.0, bad, 3.0, 0.5], [1.0, 2.0, -1.5, 3.0, 0.5]])
+    p = sign_flip_permutation_pvalue(
+        effects, exact_max_n=exact_max_n, n_permutations=200, generator=derive_generator(5)
+    )
+    assert torch.isnan(p[0])
+    expected = sign_flip_permutation_pvalue(
+        effects[1:], exact_max_n=exact_max_n, n_permutations=200, generator=derive_generator(5)
+    )
+    assert torch.equal(p[1:], expected)
+
+
+@pytest.mark.parametrize("alternative", ["two-sided", "greater", "less"])
+def test_sign_flip_mixed_infinite_effects_are_nan(alternative: str):
+    """+inf and -inf average to a NaN observed mean that an isnan() check on the inputs misses."""
+    effects = torch.tensor([[float("inf"), float("-inf"), 1.0, 2.0]])
+    assert torch.isnan(sign_flip_permutation_pvalue(effects, alternative=alternative)).all()
+
+
 def test_sign_flip_monte_carlo_separates_effect_from_null():
     torch.manual_seed(0)
     strong = torch.randn(1, 40) * 0.1 + 2.0
diff --git a/transformer_lens/ActivationCache.py b/transformer_lens/ActivationCache.py
index 08c743317d..dffcfd7288 100644
--- a/transformer_lens/ActivationCache.py
+++ b/transformer_lens/ActivationCache.py
@@ -475,7 +475,6 @@ def _mixer_stack(
         ``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)
@@ -491,8 +490,93 @@ def _mixer_stack(
             if scale != 1.0:
                 out = out * scale
             components.append(out)
-            labels.append(self._mixer_label(layer, name))
-        return components, labels
+        return components, self._mixer_labels(layer)
+
+    def _mixer_labels(self, layer: int) -> List[str]:
+        """Labels of :meth:`_mixer_stack`'s rows for block ``layer``."""
+        return [self._mixer_label(layer, name) for name in self._mixer_out_names(layer)]
+
+    def _decomposition_layer(self, layer: Optional[int]) -> int:
+        """``None`` / ``-1`` mean the full stack: the residual entering the unembedding."""
+        if layer is None or layer == -1:
+            return int(self.model.cfg.n_layers)
+        return layer
+
+    def _resid_component_plan(
+        self,
+        layer: int,
+        mode: Literal["all", "mlp", "attn"] = "all",
+        mlp_input: bool = False,
+        incl_embeds: bool = True,
+    ) -> List[Tuple[List[str], Callable[[], List[torch.Tensor]]]]:
+        """:meth:`decompose_resid`'s rows as ``(labels, read)`` groups.
+
+        The labels depend only on the block layout, so they exist before any activation
+        does; ``read`` fetches the group's tensors from the cache.
+        """
+        incl_attn = mode != "mlp"
+        incl_mlp = mode != "attn" and not self.model.cfg.attn_only
+        hook_dict = getattr(self.model, "hook_dict", None)
+        registered = hook_dict if isinstance(hook_dict, dict) else None
+        plan: List[Tuple[List[str], Callable[[], List[torch.Tensor]]]] = []
+
+        def read(key: Any) -> Callable[[], List[torch.Tensor]]:
+            return lambda: [self[key]]
+
+        def mlp_out(l: int) -> List[torch.Tensor]:
+            # Several modules can feed one MLP-side add (GraniteMoeHybrid sums its MoE and
+            # shared MLP inline), leaving no hook_mlp_out; the residual difference is that write.
+            if registered is not None and f"blocks.{l}.hook_mlp_out" not in registered:
+                resid = (f"blocks.{l}.hook_resid_mid", f"blocks.{l}.hook_resid_post")
+                if all(name in registered for name in resid):
+                    return [self[("resid_post", l)] - self[("resid_mid", l)]]
+            return [self[("mlp_out", l)]]
+
+        def add_mixer(l: int) -> None:
+            plan.append((self._mixer_labels(l), lambda: self._mixer_stack(l, Slice(None))[0]))
+
+        def add_attn_slot(l: int) -> None:
+            if self._has_attn(l):
+                plan.append(([f"{l}_attn_out"], read(("attn_out", l))))
+            elif self._mixer_role(l) == "attn":
+                add_mixer(l)
+
+        def add_mlp_slot(l: int) -> None:
+            if self._has_mlp(l):
+                plan.append(([f"{l}_mlp_out"], lambda: mlp_out(l)))
+            elif not self._has_attn(l) and self._mixer_role(l) == "mlp":
+                add_mixer(l)
+
+        if incl_embeds:
+            if self.has_embed:
+                plan.append((["embed"], read("hook_embed")))
+            if self.has_pos_embed:
+                plan.append((["pos_embed"], read("hook_pos_embed")))
+        for l in range(layer):
+            if incl_attn:
+                add_attn_slot(l)
+            if incl_mlp:
+                add_mlp_slot(l)
+        if mlp_input and incl_attn and self._mlp_input_includes_attn(layer):
+            add_attn_slot(layer)
+        return plan
+
+    def decompose_resid_labels(
+        self,
+        layer: Optional[int] = None,
+        mlp_input: bool = False,
+        mode: Literal["all", "mlp", "attn"] = "all",
+        incl_embeds: bool = True,
+    ) -> List[str]:
+        """The labels :meth:`decompose_resid` returns for these arguments, read from the
+        block layout alone, so a tool can check a decomposition before running the model."""
+        plan = self._resid_component_plan(
+            self._decomposition_layer(layer),
+            mode=mode,
+            mlp_input=mlp_input,
+            incl_embeds=incl_embeds,
+        )
+        return [label for labels, _ in plan for label in labels]
 
     def accumulated_resid(
         self,
@@ -845,7 +929,9 @@ 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``. A sequential block with no ``hook_mlp_out`` (GraniteMoeHybrid with
+                experts sums its MoE and shared MLP inline) contributes
+                ``resid_post - resid_mid`` as its ``mlp_out``.
             apply_ln:
                 Whether to apply LayerNorm to the stack.
             pos_slice:
@@ -865,47 +951,12 @@ def decompose_resid(
         if not isinstance(pos_slice, Slice):
             pos_slice = Slice(pos_slice)
         pos_slice = cast(Slice, pos_slice)  # mypy can't seem to infer this
-        if layer is None or layer == -1:
-            # Default to the residual stream immediately pre unembed
-            layer = self.model.cfg.n_layers
-        assert isinstance(layer, int)
-
-        incl_attn = mode != "mlp"
-        incl_mlp = mode != "attn" and not self.model.cfg.attn_only
-        components_list = []
-        labels = []
-        if incl_embeds:
-            if self.has_embed:
-                components_list = [self["hook_embed"]]
-                labels.append("embed")
-            if self.has_pos_embed:
-                components_list.append(self["hook_pos_embed"])
-                labels.append("pos_embed")
-
-        def add_mixer(l: int) -> None:
-            mixers, mixer_labels = self._mixer_stack(l, Slice(None))
-            components_list.extend(mixers)
-            labels.extend(mixer_labels)
-
-        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")
-            elif self._mixer_role(l) == "attn":
-                add_mixer(l)
-
-        for l in range(layer):
-            if incl_attn:
-                add_attn_slot(l)
-            if incl_mlp:
-                if self._has_mlp(l):
-                    components_list.append(self[("mlp_out", l)])
-                    labels.append(f"{l}_mlp_out")
-                elif not self._has_attn(l) and self._mixer_role(l) == "mlp":
-                    add_mixer(l)
-        if mlp_input and incl_attn and self._mlp_input_includes_attn(layer):
-            add_attn_slot(layer)
-        components_list = [pos_slice.apply(c, dim=-2) for c in components_list]
+        layer = self._decomposition_layer(layer)
+        plan = self._resid_component_plan(
+            layer, mode=mode, mlp_input=mlp_input, incl_embeds=incl_embeds
+        )
+        labels = [label for entry_labels, _ in plan for label in entry_labels]
+        components_list = [pos_slice.apply(c, dim=-2) for _, read in plan for c in read()]
         if components_list:
             components = torch.stack(components_list, dim=0)
         else:
@@ -1394,7 +1445,7 @@ def get_neuron_results(
             layer:
                 Layer index.
             neuron_slice:
-                Slice of the neuron.
+                Slice of the neuron. An int drops the neuron axis from the result.
             pos_slice:
                 Slice of the positions.
             project_output_onto:
@@ -1419,17 +1470,18 @@ def get_neuron_results(
             # Note - order is important, as Slice.apply *may* collapse a dimension, so this ensures
             # that position dimension is -2 when we apply position slice
             neuron_acts = pos_slice.apply(neuron_acts, dim=-2)
-        if neuron_slice is not None:
-            neuron_acts = neuron_slice.apply(neuron_acts, dim=-1)
-            W_out = neuron_slice.apply(W_out, dim=0)
-        if contribution_scale != 1.0:
-            neuron_acts = neuron_acts * contribution_scale
         post = self._post_norm_map(layer, "ln2_post", pos_slice)
         if post is not None:
             scale, weight, centers, _ = post
             # center(a_n W_out_n) / s * w = (a_n / s) * center(W_out_n) * w
             W_out = self._apply_post_norm_map(W_out, torch.ones(()), weight, centers)
+            # Divide before slicing: an int neuron_slice drops the axis scale's trailing 1 broadcasts over.
             neuron_acts = (neuron_acts / scale).to(neuron_acts.dtype)
+        if neuron_slice is not None:
+            neuron_acts = neuron_slice.apply(neuron_acts, dim=-1)
+            W_out = neuron_slice.apply(W_out, dim=0)
+        if contribution_scale != 1.0:
+            neuron_acts = neuron_acts * contribution_scale
         if project_output_onto is None:
             return neuron_acts[..., None] * W_out
         # W_out: [d_mlp, d_model]; project: [d_model] or [d_model, n_outs]
diff --git a/transformer_lens/SVDInterpreter.py b/transformer_lens/SVDInterpreter.py
index ce914e29ed..681415f096 100644
--- a/transformer_lens/SVDInterpreter.py
+++ b/transformer_lens/SVDInterpreter.py
@@ -162,10 +162,17 @@ def _get_w_in_matrix(self, layer_index: int) -> torch.Tensor:
         w_in = self.params[key].T
 
         if f"blocks.{layer_index}.ln2.w" in self.params:  # If fold_ln == False
+            # Lazy: keeps `import transformer_lens` from loading the model_bridge package.
+            from transformer_lens.model_bridge.generalized_components.normalization import (
+                norm_uses_weight_offset,
+            )
+
             ln_2 = self.params[f"blocks.{layer_index}.ln2.w"]
-            if getattr(self.cfg, "rmsnorm_uses_offset", False):
-                # Offset-RMS models (Gemma family) store w as an offset from 1:
-                # the effective scale is (1 + w), and post-fold w = 0 is identity.
+            blocks = getattr(self.model, "blocks", None)
+            norm = getattr(blocks[layer_index], "ln2", None) if blocks is not None else None
+            if norm_uses_weight_offset(norm, self.cfg):
+                # Offset-RMS norms store w as an offset from 1: the effective scale is
+                # (1 + w), and post-fold w = 0 is identity.
                 ln_2 = 1.0 + ln_2
             return w_in * ln_2
 
diff --git a/transformer_lens/model_bridge/architecture_adapter.py b/transformer_lens/model_bridge/architecture_adapter.py
index e7fd9033c0..ab03d873c8 100644
--- a/transformer_lens/model_bridge/architecture_adapter.py
+++ b/transformer_lens/model_bridge/architecture_adapter.py
@@ -370,7 +370,9 @@ def preprocess_weights(self, state_dict: dict[str, torch.Tensor]) -> dict[str, t
 
         This method allows architectures to apply custom transformations to weights
         before standard weight processing (fold_layer_norm, center_writing_weights, etc.).
-        For example, Gemma models scale embeddings by sqrt(d_model).
+        For example, Cohere folds ``logit_scale`` into the unembedding here. Gemma's
+        ``sqrt(d_model)`` embedding scale is not a fold: HF applies it at runtime, so
+        folding it as well would double-scale.
 
         Args:
             state_dict: The state dictionary, keyed in TransformerLens format. The caller
diff --git a/transformer_lens/model_bridge/bridge_core.py b/transformer_lens/model_bridge/bridge_core.py
index c40ccba0c8..57292c77b1 100644
--- a/transformer_lens/model_bridge/bridge_core.py
+++ b/transformer_lens/model_bridge/bridge_core.py
@@ -1174,6 +1174,22 @@ def cache_some(
             remove_batch_dim=remove_batch_dim,
         )
 
+    def _has_registered_blocks(self) -> bool:
+        """Whether a ``blocks`` stack is registered as a submodule on this bridge.
+
+        Checks ``_modules`` directly rather than ``hasattr``: ``__getattr__`` falls
+        through to the wrapped HF model, so ``hasattr(self, "blocks")`` can be True
+        for a model that merely exposes its own ``.blocks`` attribute.
+        """
+        modules = self.__dict__.get("_modules") or {}
+        return "blocks" in modules
+
+    def _registered_block_count(self) -> Optional[int]:
+        """Length of the registered ``blocks`` stack, or None when there is none."""
+        if not self._has_registered_blocks():
+            return None
+        return len(self.__dict__["_modules"]["blocks"])
+
     @staticmethod
     def _resolve_stop_at_layer(stop_at_layer: int, n_blocks: int) -> int:
         """Map ``stop_at_layer`` to ``[0, n_blocks]``, with ``n_blocks`` past the last block.
@@ -1317,13 +1333,14 @@ def run_with_hooks(
         """
         _validate_execution_kwargs(kwargs, "run_with_hooks")
         added_hooks: List[Tuple[HookPoint, Literal["fwd", "bwd"], LensHandle]] = []
+        n_blocks = self._registered_block_count()
         effective_stop_layer = None
-        if stop_at_layer is not None and hasattr(self, "blocks"):
-            effective_stop_layer = self._resolve_stop_at_layer(stop_at_layer, len(self.blocks))
+        if stop_at_layer is not None and n_blocks is not None:
+            effective_stop_layer = self._resolve_stop_at_layer(stop_at_layer, n_blocks)
         effective_start_layer = None
-        if start_at_layer is not None and hasattr(self, "blocks"):
+        if start_at_layer is not None and n_blocks is not None:
             effective_start_layer = (
-                len(self.blocks) + start_at_layer if start_at_layer < 0 else start_at_layer
+                n_blocks + start_at_layer if start_at_layer < 0 else start_at_layer
             )
 
         def add_hook_to_point(
@@ -1478,7 +1495,9 @@ def run_with_cache(
     ) -> Tuple[Any, Union[ActivationCache, Dict[str, torch.Tensor]]]:
         """Run the model and cache activations. Returns ``(output, cache)``.
 
-        ``stop_at_layer`` raises :class:`StopAtLayerException` to stop early.
+        ``stop_at_layer`` raises :class:`StopAtLayerException` to stop early; on an
+        architecture with no ``blocks`` stack it raises ``NotImplementedError``, as
+        :meth:`forward` does.
         ``start_at_layer`` treats ``input`` as the residual entering block ``k``
         (see :meth:`forward`); blocks below ``k`` are excluded from the cache to
         match HookedTransformer. ``pos_slice`` slices each cached activation along
@@ -1541,8 +1560,6 @@ def create_names_filter_fn(filter_input):
         cache: Dict[str, torch.Tensor] = {}
         hooks: List[Tuple[HookPoint, str]] = []
         visited: set[int] = set()
-        stop_hook_point: Optional[HookPoint] = None
-        stop_hook_fn: Optional[Callable] = None
 
         # None → no-op .to(None), tensors stay on their current device.
         cache_device = kwargs.pop("device", None)
@@ -1577,26 +1594,17 @@ def cache_hook(tensor: torch.Tensor, *, hook: Any) -> torch.Tensor:
             return cache_hook
 
         hook_dict = self.hook_dict
+        n_blocks = self._registered_block_count()
         effective_stop_layer = None
-        if stop_at_layer is not None and hasattr(self, "blocks"):
-            effective_stop_layer = self._resolve_stop_at_layer(stop_at_layer, len(self.blocks))
-            if effective_stop_layer == len(self.blocks):
-                # No block follows the last for a hook_in to stop on; forward stops on its exit.
-                kwargs["stop_at_layer"] = effective_stop_layer
-            else:
-                # Stop at the beginning of the specified block, not at the end of the previous block
-                block_hook_name = f"blocks.{effective_stop_layer}.hook_in"
-                if block_hook_name in hook_dict:
-
-                    def stop_hook(tensor: torch.Tensor, *, hook: Any) -> torch.Tensor:
-                        raise StopAtLayerException(tensor)
-
-                    stop_hook_point = hook_dict[block_hook_name]
-                    stop_hook_fn = stop_hook
+        if stop_at_layer is not None:
+            # forward owns the guard and the stop; resolving here only bounds the cache.
+            kwargs["stop_at_layer"] = stop_at_layer
+            if n_blocks is not None:
+                effective_stop_layer = self._resolve_stop_at_layer(stop_at_layer, n_blocks)
         effective_start_layer = None
-        if start_at_layer is not None and hasattr(self, "blocks"):
+        if start_at_layer is not None and n_blocks is not None:
             effective_start_layer = (
-                len(self.blocks) + start_at_layer if start_at_layer < 0 else start_at_layer
+                n_blocks + start_at_layer if start_at_layer < 0 else start_at_layer
             )
         matched_any = False
         for hook_name, hook in hook_dict.items():
@@ -1692,7 +1700,6 @@ def grad_hook(tensor: torch.Tensor, *, hook: Any) -> None:
         context_level = getattr(self, "context_level", 0) + 1
         self.context_level = context_level
         added_hooks: List[Tuple[HookPoint, Literal["fwd", "bwd"], LensHandle]] = []
-        stop_handle: Optional[LensHandle] = None
         try:
             for hp, name in hooks:
                 fwd_handle = hp.add_hook(make_cache_hook(name), level=context_level)
@@ -1702,9 +1709,6 @@ def grad_hook(tensor: torch.Tensor, *, hook: Any) -> None:
                         make_grad_cache_hook(name), dir="bwd", level=context_level
                     )
                     added_hooks.append((hp, "bwd", bwd_handle))
-            if stop_hook_point is not None and stop_hook_fn is not None:
-                stop_handle = stop_hook_point.add_hook(stop_hook_fn, level=context_level)
-                added_hooks.append((stop_hook_point, "fwd", stop_handle))
         except Exception:
             try:
                 for hp, direction, handle in added_hooks:
@@ -1764,14 +1768,11 @@ def grad_hook(tensor: torch.Tensor, *, hook: Any) -> None:
         finally:
             try:
                 for hp, direction, handle in added_hooks:
-                    # Retain requested cache hooks, but never internal early-stop hooks.
-                    if reset_hooks_end or handle is stop_handle:
+                    if reset_hooks_end:
                         self._remove_temporary_hook(hp, direction, handle)
                 if clear_contexts:
                     for hp, _ in hooks:
                         hp.clear_context()
-                    if stop_hook_point is not None:
-                        stop_hook_point.clear_context()
             finally:
                 self.context_level -= 1
         if self.compatibility_mode == True:
diff --git a/transformer_lens/model_bridge/generalized_components/altup_block.py b/transformer_lens/model_bridge/generalized_components/altup_block.py
index 28c1ccd680..3dbf369419 100644
--- a/transformer_lens/model_bridge/generalized_components/altup_block.py
+++ b/transformer_lens/model_bridge/generalized_components/altup_block.py
@@ -1,7 +1,6 @@
 """Block bridge for AltUp (Alternating Updates) decoder layers."""
 from __future__ import annotations
 
-import re
 from typing import Any, Dict, Optional
 
 import torch
@@ -109,11 +108,8 @@ def _hook_output(self, output: Any) -> Any:
         return primary
 
     def _check_stop_at_layer(self, *args: Any, **kwargs: Any) -> None:
-        """Raise StopAtLayerException (carrying the AltUp stack) at the configured stop index."""
-        if getattr(self, "_stop_at_layer_idx", None) is None or self.name is None:
-            return
-        match = re.search(r"\.layers\.(\d+)", self.name) or re.search(r"blocks\.(\d+)", self.name)
-        if not match or int(match.group(1)) != self._stop_at_layer_idx:
+        """Raise StopAtLayerException (carrying the AltUp stack) if the bridge armed this block."""
+        if getattr(self, "_stop_at_layer_idx", None) is None:
             return
         if len(args) > 0 and isinstance(args[0], torch.Tensor):
             tensor = args[0]
diff --git a/transformer_lens/model_bridge/generalized_components/block.py b/transformer_lens/model_bridge/generalized_components/block.py
index a383bd661e..0da73e082b 100644
--- a/transformer_lens/model_bridge/generalized_components/block.py
+++ b/transformer_lens/model_bridge/generalized_components/block.py
@@ -5,7 +5,6 @@
 from __future__ import annotations
 
 import inspect
-import re
 import weakref
 from typing import Any, Callable, Dict, Optional, cast
 
@@ -321,45 +320,31 @@ def _is_standalone_hidden_state_call(args: tuple, kwargs: dict) -> bool:
             and isinstance(kwargs["hidden_states"], torch.Tensor)
         )
 
-    def _extract_layer_idx(self) -> Optional[int]:
-        """Parse this block's layer index from its name (TL/GPT-2/LLaMA patterns)."""
-        if self.name is None:
-            return None
-        match = re.search(r"(?:^|\.)(?:blocks|h|layers)\.(\d+)", self.name)
-        return int(match.group(1)) if match else None
-
     def _check_stop_at_layer(self, *args: Any, **kwargs: Any) -> None:
-        """Check if execution should stop before this block. Raises StopAtLayerException.
-
-        The _stop_at_layer_idx attribute is set by the bridge's forward method.
-        Supports TL/GPT-2/LLaMA naming patterns for layer index extraction.
-        """
-        if getattr(self, "_stop_at_layer_idx", None) is None:
+        """Raise StopAtLayerException with the hooked input if the bridge armed this block."""
+        layer_idx = getattr(self, "_stop_at_layer_idx", None)
+        if layer_idx is None:
             return
-        layer_idx = self._extract_layer_idx()
-        if layer_idx is not None and layer_idx == self._stop_at_layer_idx:
-            if len(args) > 0 and isinstance(args[0], torch.Tensor):
-                input_tensor = args[0]
-            elif "hidden_states" in kwargs and isinstance(kwargs["hidden_states"], torch.Tensor):
-                input_tensor = kwargs["hidden_states"]
-            else:
-                raise ValueError(f"Cannot find input tensor to stop at layer {layer_idx}")
-            input_tensor = self.hook_in(input_tensor)
-            raise StopAtLayerException(input_tensor)
+        if len(args) > 0 and isinstance(args[0], torch.Tensor):
+            input_tensor = args[0]
+        elif "hidden_states" in kwargs and isinstance(kwargs["hidden_states"], torch.Tensor):
+            input_tensor = kwargs["hidden_states"]
+        else:
+            raise ValueError(f"Cannot find input tensor to stop at layer {layer_idx}")
+        input_tensor = self.hook_in(input_tensor)
+        raise StopAtLayerException(input_tensor)
 
     def _maybe_inject_start_residual(self, args: tuple, kwargs: dict) -> tuple[tuple, dict]:
         """If this is the start_at_layer block, swap in the caller's residual.
 
-        Mirror of ``_check_stop_at_layer``: the bridge's forward stashes the
-        residual-stream input on the block via ``_start_residual`` and sets
+        Mirror of ``_check_stop_at_layer``: the bridge's forward arms only the start
+        block, stashing the residual-stream input via ``_start_residual`` and setting
         ``_start_at_layer_idx``. This block replaces its incoming hidden states
         with that residual; ``_hook_input_hidden_states`` then fires ``hook_in``
         on it, so ``hook_resid_pre`` reflects the injected value.
         """
         if getattr(self, "_start_at_layer_idx", None) is None:
             return args, kwargs
-        if self._extract_layer_idx() != self._start_at_layer_idx:
-            return args, kwargs
         residual = getattr(self, "_start_residual", None)
         if residual is None:
             return args, kwargs
diff --git a/transformer_lens/model_bridge/generalized_components/bloom_block.py b/transformer_lens/model_bridge/generalized_components/bloom_block.py
index d3a2f9a020..e463ddf3ec 100644
--- a/transformer_lens/model_bridge/generalized_components/bloom_block.py
+++ b/transformer_lens/model_bridge/generalized_components/bloom_block.py
@@ -79,6 +79,7 @@ def forward(self, *args: Any, **kwargs: Any) -> Any:
                 f"Original component not set for {self.name}. Call set_original_component() first."
             )
 
+        args, kwargs = self._maybe_inject_start_residual(args, kwargs)
         self._check_stop_at_layer(*args, **kwargs)
         args, kwargs = self._hook_input_hidden_states(args, kwargs)
 
diff --git a/transformer_lens/model_bridge/generalized_components/normalization.py b/transformer_lens/model_bridge/generalized_components/normalization.py
index 5920445c4c..3c51df0330 100644
--- a/transformer_lens/model_bridge/generalized_components/normalization.py
+++ b/transformer_lens/model_bridge/generalized_components/normalization.py
@@ -104,6 +104,14 @@ def uses_weight_offset(self) -> bool:
             return self._uses_weight_offset_override
         return self.uses_rms_norm and bool(getattr(self.config, "rmsnorm_uses_offset", False))
 
+    def _eps(self) -> float | torch.Tensor:
+        """The wrapped norm's epsilon (``eps`` or ``variance_epsilon``), else the config's."""
+        for attr in ("eps", "variance_epsilon"):
+            value = getattr(self.original_component, attr, None)
+            if value is not None:
+                return value
+        return getattr(self.config, "eps", 1e-05)
+
     def forward(self, hidden_states: torch.Tensor, **kwargs: Any) -> torch.Tensor:
         """Forward pass through the normalization bridge.
 
@@ -139,9 +147,7 @@ def _python_norm_forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
             hidden_states = hidden_states - hidden_states.mean(-1, keepdim=True)
         # sqrt saves its output for backward; the clone keeps in-place hook edits off it.
         scale = self.hook_scale(
-            (hidden_states.pow(2).mean(-1, keepdim=True) + getattr(self.config, "eps", 1e-05))
-            .sqrt()
-            .clone()
+            (hidden_states.pow(2).mean(-1, keepdim=True) + self._eps()).sqrt().clone()
         )
         hidden_states = self.hook_normalized(hidden_states / scale)
         return self._apply_weight_and_bias(hidden_states, input_dtype)
@@ -209,13 +215,7 @@ def _hf_autograd_forward_with_hooks(self, x: torch.Tensor) -> torch.Tensor:
                 x_centered = x_float - x_float.mean(-1, keepdim=True)
             else:
                 x_centered = x_float
-            eps_tensor = getattr(self.original_component, "eps", None)
-            if eps_tensor is None:
-                eps_tensor = getattr(self.original_component, "variance_epsilon", None)
-            if eps_tensor is None:
-                eps_value: float | torch.Tensor = getattr(self.config, "eps", 1e-05)
-            else:
-                eps_value = eps_tensor
+            eps_value = self._eps()
             variance = x_centered.pow(2).mean(-1, keepdim=True)
             if isinstance(eps_value, torch.Tensor):
                 inv_rms = torch.rsqrt(variance + eps_value)
@@ -301,3 +301,11 @@ def _apply_weight_and_bias(
     ) -> torch.Tensor:
         """No learnable scale to apply — only restore the input dtype."""
         return hidden_states.to(input_dtype)
+
+
+def norm_uses_weight_offset(norm: Optional[torch.nn.Module], cfg: Any) -> bool:
+    """Whether ``norm`` scales by ``1 + w``: the mounted bridge's own flag, else the config's."""
+    # Mixed adapters declare offset block norms per-norm while cfg describes ln_final.
+    if isinstance(norm, NormalizationBridge):
+        return norm.uses_weight_offset
+    return bool(getattr(cfg, "rmsnorm_uses_offset", False))
diff --git a/transformer_lens/model_bridge/generalized_components/opaque_block.py b/transformer_lens/model_bridge/generalized_components/opaque_block.py
index 2bfb990fc6..60ef6685d2 100644
--- a/transformer_lens/model_bridge/generalized_components/opaque_block.py
+++ b/transformer_lens/model_bridge/generalized_components/opaque_block.py
@@ -2,7 +2,6 @@
 
 from __future__ import annotations
 
-import re
 from typing import Any, Dict, Optional
 
 import torch
@@ -63,6 +62,7 @@ def forward(self, *args: Any, **kwargs: Any) -> Any:
                 "Call set_original_component() first."
             )
 
+        args, kwargs = self._maybe_inject_start_residual(args, kwargs)
         self._check_stop_at_layer(*args, **kwargs)
         args, kwargs = self._hook_input_hidden_states(args, kwargs)
         output = self.original_component(*args, **kwargs)
@@ -89,18 +89,23 @@ def _hook_input_hidden_states(self, args: tuple, kwargs: dict) -> tuple[tuple, d
             kwargs["hidden_states"] = self.hook_in(kwargs["hidden_states"])
         return args, kwargs
 
+    def _maybe_inject_start_residual(self, args: tuple, kwargs: dict) -> tuple[tuple, dict]:
+        """Swap in the caller's residual when the bridge armed this block as the start layer."""
+        if getattr(self, "_start_at_layer_idx", None) is None:
+            return args, kwargs
+        residual = getattr(self, "_start_residual", None)
+        if residual is None:
+            return args, kwargs
+        if len(args) > 0 and isinstance(args[0], torch.Tensor):
+            return (residual,) + args[1:], kwargs
+        if "hidden_states" in kwargs:
+            kwargs = {**kwargs, "hidden_states": residual}
+        return args, kwargs
+
     def _check_stop_at_layer(self, *args: Any, **kwargs: Any) -> None:
-        """Raise StopAtLayerException when the configured stop index matches this block."""
-        if not (hasattr(self, "_stop_at_layer_idx") and self._stop_at_layer_idx is not None):
-            return
-        if self.name is None:
-            return
-        # Mamba uses `.layers.{i}`; `blocks.{i}` is the fallback TL convention.
-        match = re.search(r"\.layers\.(\d+)", self.name) or re.search(r"blocks\.(\d+)", self.name)
-        if not match:
-            return
-        layer_idx = int(match.group(1))
-        if layer_idx != self._stop_at_layer_idx:
+        """Raise StopAtLayerException with the hooked input if the bridge armed this block."""
+        layer_idx = getattr(self, "_stop_at_layer_idx", None)
+        if layer_idx is None:
             return
         if len(args) > 0 and isinstance(args[0], torch.Tensor):
             input_tensor = args[0]
diff --git a/transformer_lens/model_bridge/supported_architectures/AGENTS.md b/transformer_lens/model_bridge/supported_architectures/AGENTS.md
index 8e7a205a65..94a4f06f89 100644
--- a/transformer_lens/model_bridge/supported_architectures/AGENTS.md
+++ b/transformer_lens/model_bridge/supported_architectures/AGENTS.md
@@ -78,7 +78,7 @@ Optional overrides on the class:
 | Decoder-only causal LM, RMSNorm + RoPE | [`llama.py`](llama.py) | Canonical modern shape; already handles GQA via `n_key_value_heads` |
 | Decoder-only causal LM, LayerNorm + combined QKV | [`gpt2.py`](gpt2.py) | Older GPT-style; demonstrates QKV split/rearrange |
 | GQA / MQA | [`llama.py`](llama.py) | Pass `n_kv_heads=` to `_qkvo_weight_conversions()` |
-| RMSNorm with offset | [`gemma1.py`](gemma1.py) | Sets `cfg.rmsnorm_uses_offset = True` only — weight processing and the runtime norm read the flag and apply the +1 themselves; conversion entries on norm weights are never read |
+| RMSNorm with offset | [`gemma1.py`](gemma1.py) | Sets `cfg.rmsnorm_uses_offset = True` only — weight processing and the runtime norm read the flag and apply the +1 themselves; conversion entries on norm weights are never read. Check the HF norm's forward, not the family name. Mixed layouts: see [`muse_glimmer.py`](muse_glimmer.py) |
 | Mixture-of-experts | [`mixtral.py`](mixtral.py) | Uses `MoEBridge` with batched experts |
 | Vision-language | [`llava.py`](llava.py) | Dual encoder + projection pathways |
 | Encoder-decoder (T5-style) | [`t5.py`](t5.py) | Sets `supports_generation = False`, separate encoder/decoder block lists |
@@ -277,6 +277,7 @@ Failure message names the missing set. (`INTENTIONAL_EXCLUDES` in the test handl
 | RoPE (rotary positional embeddings) | `llama.py`, `mistral.py`, `qwen2.py`+ | `RotaryEmbeddingBridge(name="model.rotary_emb")` + `cfg.positional_embedding_type = "rotary"` |
 | GQA / MQA (`n_key_value_heads < n_heads`) | `llama.py`, `mistral.py`, `falcon.py`, `cohere.py` | Set `cfg.n_key_value_heads`; pass `n_kv_heads=` to `_qkvo_weight_conversions()` |
 | RMSNorm with offset | `gemma1.py`, `gemma2.py`, `gemma3.py` | `cfg.rmsnorm_uses_offset = True` only (no conversion entry: processing reads norm weights directly and applies the offset from the flag; conversion entries on norm weights are never read) |
+| RMSNorm with offset on some norms only | `muse_glimmer.py` | `cfg.rmsnorm_uses_offset` describes `ln_final`; block norms declare `NormalizationBridge(..., uses_weight_offset=True)`. Analysis readers holding a module (SVD `w_in`, DLA `ln_final`) consult the mounted norm first via `norm_uses_weight_offset`, falling back to cfg. Weight-processing folds read cfg only, so a mixed adapter must keep `supports_fold_ln = False` |
 | Standard LayerNorm | `gpt2.py`, `bloom.py` | `cfg.normalization_type = "LN"` |
 | Gated MLP (`gate_proj`, `up_proj`, `down_proj`) | `llama.py`, `mistral.py`, `gemma1.py`, `qwen2.py`+ | `GatedMLPBridge` with submodules `{gate, in, out}` |
 | Combined QKV (`c_attn`) | `gpt2.py`, `bloom.py` | `QKVSplitRearrangeConversion` to split + rearrange |
diff --git a/transformer_lens/model_bridge/supported_architectures/gemma3n.py b/transformer_lens/model_bridge/supported_architectures/gemma3n.py
index e151d9189d..f551649c96 100644
--- a/transformer_lens/model_bridge/supported_architectures/gemma3n.py
+++ b/transformer_lens/model_bridge/supported_architectures/gemma3n.py
@@ -64,7 +64,8 @@ def __init__(self, cfg: Any) -> None:
         self.cfg.gated_mlp = True
         self.cfg.uses_rms_norm = True
         self.cfg.normalization_type = "RMS"
-        self.cfg.rmsnorm_uses_offset = True  # Gemma RMSNorm uses (1.0 + weight)
+        # Unlike Gemma 1-3, Gemma3nRMSNorm scales by the raw weight (initialized to ones).
+        self.cfg.rmsnorm_uses_offset = False
         self.cfg.positional_embedding_type = "rotary"
         self.cfg.attn_implementation = "eager"
         # AltUp + per-layer-embedding residual topology isn't fold-safe.
diff --git a/transformer_lens/model_bridge/supported_architectures/muse_glimmer.py b/transformer_lens/model_bridge/supported_architectures/muse_glimmer.py
index 313ec6896a..91bef1b200 100644
--- a/transformer_lens/model_bridge/supported_architectures/muse_glimmer.py
+++ b/transformer_lens/model_bridge/supported_architectures/muse_glimmer.py
@@ -10,6 +10,7 @@
     BlockBridge,
     EmbeddingBridge,
     LinearBridge,
+    NormalizationBridge,
     RotaryEmbeddingBridge,
     UnembeddingBridge,
     VisionProjectionBridge,
@@ -40,6 +41,18 @@ def __init__(self, cfg: Any) -> None:
         self.cfg.attn_implementation = "eager"
         self.weight_processing_conversions: dict = {}
 
+        def sandwich_norm(name: str) -> NormalizationBridge:
+            # Block norms scale by (1 + w) but the final norm by raw w, so the offset is
+            # declared per-norm and cfg.rmsnorm_uses_offset stays False for ln_final.
+            # Both forward paths read eps from the HF norm, so the post norms keep
+            # post_norm_eps while cfg.eps stays rms_norm_eps.
+            return NormalizationBridge(
+                name=name,
+                config=self.cfg,
+                use_native_layernorm_autograd=True,
+                uses_weight_offset=True,
+            )
+
         self.component_mapping = {
             "vision_encoder": GeneralizedComponent(name="model.vision_tower"),
             # HF runs ``vision_adapter -> vision_projection -> perception_emb_norm`` and
@@ -53,10 +66,10 @@ def __init__(self, cfg: Any) -> None:
             "blocks": BlockBridge(
                 name="model.language_model.layers",
                 submodules={
-                    "ln1": GeneralizedComponent(name="input_layernorm"),
-                    "ln1_post": GeneralizedComponent(name="post_attention_layernorm"),
-                    "ln2": GeneralizedComponent(name="pre_feedforward_layernorm"),
-                    "ln2_post": GeneralizedComponent(name="post_feedforward_layernorm"),
+                    "ln1": sandwich_norm("input_layernorm"),
+                    "ln1_post": sandwich_norm("post_attention_layernorm"),
+                    "ln2": sandwich_norm("pre_feedforward_layernorm"),
+                    "ln2_post": sandwich_norm("post_feedforward_layernorm"),
                     "attn": AttentionBridge(
                         name="self_attn",
                         config=self.cfg,
diff --git a/transformer_lens/model_bridge/transformer_bridge.py b/transformer_lens/model_bridge/transformer_bridge.py
index 3e19706d98..28456fece9 100644
--- a/transformer_lens/model_bridge/transformer_bridge.py
+++ b/transformer_lens/model_bridge/transformer_bridge.py
@@ -530,16 +530,6 @@ def n_params_total(self) -> int:
         """
         return self._n_params_total
 
-    def _has_registered_blocks(self) -> bool:
-        """Whether a ``blocks`` stack is registered as a submodule on this bridge.
-
-        Checks ``_modules`` directly rather than ``hasattr``: ``__getattr__`` falls
-        through to the wrapped HF model, so ``hasattr(self, "blocks")`` can be True
-        for a model that merely exposes its own ``.blocks`` attribute.
-        """
-        modules = self.__dict__.get("_modules") or {}
-        return "blocks" in modules
-
     def __getattr__(self, name: str) -> Any:
         """Provide a clear error message for missing attributes."""
         # Re-invoke original_model's property so its descriptive AttributeError
@@ -1008,6 +998,27 @@ def to_tokens(
         Returns:
             Token tensor of shape ``[batch, pos]``.
         """
+        return self._to_tokens_and_mask(
+            input,
+            prepend_bos=prepend_bos,
+            padding_side=padding_side,
+            move_to_device=move_to_device,
+            truncate=truncate,
+        )[0]
+
+    def _to_tokens_and_mask(
+        self,
+        input: Union[str, List[str]],
+        prepend_bos: Optional[bool] = None,
+        padding_side: Optional[str] = None,
+        move_to_device: bool = True,
+        truncate: bool = True,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """``to_tokens`` plus the tokenizer's attention mask, sliced in lockstep with the tokens.
+
+        Rebuilding the mask from the token ids can't tell a pad from a real token with
+        the pad's id (a literal ``<|endoftext|>`` on GPT-2), so string callers keep this one.
+        """
         assert self.tokenizer is not None, "Cannot use to_tokens without a tokenizer"
         if prepend_bos is None:
             prepend_bos = getattr(self.cfg, "default_prepend_bos", True)
@@ -1027,14 +1038,17 @@ def to_tokens(
             # tokenize to subword garbage, so skip rather than pollute the input.
         if isinstance(input, str):
             input = [input]
-        tokens = self.tokenizer(
+        encoded = self.tokenizer(
             input,
             return_tensors="pt",
             padding=True,
             padding_side=padding_side,
             truncation=truncate,
             max_length=self.cfg.n_ctx if truncate else None,
-        )["input_ids"]
+            return_attention_mask=True,
+        )
+        tokens = encoded["input_ids"]
+        mask: Optional[torch.Tensor] = encoded.get("attention_mask")
         # Strip auto-appended EOS tokens (e.g., OLMo)
         if (
             getattr(self.cfg, "tokenizer_appends_eos", False)
@@ -1043,13 +1057,38 @@ def to_tokens(
             # Remove trailing EOS, keep at least 1 token
             while tokens.shape[-1] > 1 and (tokens[:, -1] == self.tokenizer.eos_token_id).all():
                 tokens = tokens[:, :-1]
+                if mask is not None:
+                    mask = mask[:, :-1]
         if not prepend_bos and tokenizer_prepends_bos:
+            n_cols = tokens.shape[-1]
             tokens = utils.get_tokens_with_bos_removed(
                 self.tokenizer, tokens, padding_side=padding_side
             )
+            if mask is not None and tokens.shape[-1] < n_cols:
+                # Each row loses one real token, so its pads stay where they were.
+                mask = mask[..., :-1] if padding_side == "left" else mask[..., 1:]
+        if mask is None:
+            # Remote-code tokenizers may omit the mask; fall back to reading the pads,
+            # which only a multi-row batch can have.
+            if (
+                tokens.shape[0] > 1
+                and isinstance(self.tokenizer, PreTrainedTokenizerBase)
+                and self.tokenizer.pad_token_id is not None
+            ):
+                mask = utils.get_attention_mask(
+                    self.tokenizer, tokens, prepend_bos, padding_side=padding_side
+                )
+            else:
+                mask = torch.ones_like(tokens)
+        mask = mask.long()
+        if mask.shape[-1] > 0:
+            # An empty prompt keeps one pad as context; an all-masked row softmaxes to NaN.
+            empty = ~mask.bool().any(dim=-1)
+            mask[empty, -1 if padding_side == "left" else 0] = 1
         if move_to_device:
             tokens = tokens.to(self.cfg.device)
-        return tokens
+            mask = mask.to(self.cfg.device)
+        return tokens, mask
 
     def encoder_output(
         self,
@@ -2057,6 +2096,8 @@ def input_to_embed(
         Bridge analog of :meth:`HookedTransformer.input_to_embed`. Returns
         ``(residual, tokens, shortformer_pos_embed, attention_mask)``; feed the
         residual to ``forward(..., start_at_layer=0)`` to resume the pass.
+        ``attention_mask`` is the caller's mask, else the tokenizer's for a padded
+        string batch on either padding side, else ``None``.
 
         ``shortformer_pos_embed`` is always ``None``: for the models the bridge
         supports the residual already carries positional information, so there is
@@ -2064,18 +2105,22 @@ def input_to_embed(
         """
         if isinstance(input, (str, list)):
             assert self.tokenizer is not None, "Must provide a tokenizer for string input."
-            tokens = self.to_tokens(input, prepend_bos=prepend_bos, padding_side=padding_side)
+            tokens, string_mask = self._to_tokens_and_mask(
+                input, prepend_bos=prepend_bos, padding_side=padding_side
+            )
+            if attention_mask is None and not bool(string_mask.all()):
+                attention_mask = string_mask
         else:
             tokens = input
             if isinstance(tokens, torch.Tensor) and tokens.ndim == 1:
                 tokens = tokens.unsqueeze(0)
-        if (
-            attention_mask is None
-            and self.tokenizer is not None
-            and self.tokenizer.padding_side == "left"
-        ):
-            _prepend = self.cfg.default_prepend_bos if prepend_bos is None else prepend_bos
-            attention_mask = utils.get_attention_mask(self.tokenizer, tokens, _prepend)
+            if (
+                attention_mask is None
+                and self.tokenizer is not None
+                and self.tokenizer.padding_side == "left"
+            ):
+                _prepend = self.cfg.default_prepend_bos if prepend_bos is None else prepend_bos
+                attention_mask = utils.get_attention_mask(self.tokenizer, tokens, _prepend)
         residual = self.forward(tokens, stop_at_layer=0, attention_mask=attention_mask)
         return residual, tokens, None, attention_mask
 
@@ -2189,10 +2234,10 @@ def forward(
                 feed back on the next call for incremental decoding.
             loss_per_token: Whether to return loss per token
             prepend_bos: Whether to prepend BOS token
-            padding_side: Which side to pad on when tokenizing a list of strings. For a
-                token tensor, the side it was already padded on: the attention mask is
-                then read off the pad token. Pass ``attention_mask`` instead when the pad
-                id also occurs as a real token.
+            padding_side: Which side to pad on when tokenizing a list of strings, whose
+                mask is the tokenizer's own. For a token tensor, the side it was already
+                padded on: the attention mask is then read off the pad token. Pass
+                ``attention_mask`` instead when the pad id also occurs as a real token.
             labels: Explicit language-model targets. Encoder-decoder models require
                 labels for loss; decoder-only models fall back to input IDs when omitted.
             past_key_values: HuggingFace KV cache from a prior ``use_cache=True`` step
@@ -2268,9 +2313,10 @@ def forward(
             effective_stop_at_layer = self._resolve_stop_at_layer(stop_at_layer, len(self.blocks))
         if start_at_layer is not None:
             input = self._setup_start_at_layer(input, start_at_layer)
-        if effective_stop_at_layer is not None:
-            for block in self.blocks:
-                block._stop_at_layer_idx = effective_stop_at_layer
+        # Armed by position in the bridge's own stack: HF module names (h.N, layers.N,
+        # encoder.layer.N, ...) are too varied to parse back into an index.
+        if effective_stop_at_layer is not None and effective_stop_at_layer < len(self.blocks):
+            self.blocks[effective_stop_at_layer]._stop_at_layer_idx = effective_stop_at_layer
 
         # Map HookedEncoderDecoder-style kwargs to HF-compatible names
         if "decoder_input" in kwargs:
@@ -2289,6 +2335,7 @@ def forward(
             and not getattr(self.cfg, "is_audio_model", False)
             and not getattr(self.cfg, "is_visual_model", False)
         )
+        _string_mask: Optional[torch.Tensor] = None
         _resolved_padding_side = padding_side
         if _resolved_padding_side is None and self.tokenizer is not None:
             _resolved_padding_side = getattr(self.tokenizer, "padding_side", "right")
@@ -2317,9 +2364,14 @@ def forward(
                         "Visual models require tensor input (pixel values), not text. "
                         "Pass a torch.Tensor or use the pixel_values parameter."
                     )
-                input_ids = self.to_tokens(
-                    input, prepend_bos=prepend_bos, padding_side=padding_side
-                )
+                if _is_batched_list:
+                    input_ids, _string_mask = self._to_tokens_and_mask(
+                        input, prepend_bos=prepend_bos, padding_side=_resolved_padding_side
+                    )
+                else:
+                    input_ids = self.to_tokens(
+                        input, prepend_bos=prepend_bos, padding_side=padding_side
+                    )
             else:
                 input_ids = input
                 # Promote 1D integer token tensors to 2D [batch=1, seq] to match
@@ -2377,15 +2429,19 @@ def forward(
                             "Pass attention_mask=... to state the padding directly instead."
                         )
                 else:
-                    derived_mask = utils.get_attention_mask(
-                        self.tokenizer,
-                        input_ids,
-                        prepend_bos=(
-                            getattr(self.cfg, "default_prepend_bos", True)
-                            if prepend_bos is None
-                            else prepend_bos
-                        ),
-                        padding_side=_resolved_padding_side,
+                    derived_mask = (
+                        _string_mask
+                        if _string_mask is not None
+                        else utils.get_attention_mask(
+                            self.tokenizer,
+                            input_ids,
+                            prepend_bos=(
+                                getattr(self.cfg, "default_prepend_bos", True)
+                                if prepend_bos is None
+                                else prepend_bos
+                            ),
+                            padding_side=_resolved_padding_side,
+                        )
                     ).to(self.cfg.device)
                     # A tensor with no pads stays on the unmasked path, so its logits are unchanged.
                     _is_unpadded = _is_padded_tokens and bool(derived_mask.all())
@@ -2838,13 +2894,18 @@ def _setup_start_at_layer(self, input: Any, start_at_layer: int) -> torch.Tensor
             start_at_layer += n_blocks
         if not 0 <= start_at_layer < n_blocks:
             raise ValueError(f"start_at_layer={start_at_layer} out of range [0, {n_blocks}).")
+        start_block = self.blocks[start_at_layer]
+        if not callable(getattr(start_block, "_maybe_inject_start_residual", None)):
+            raise NotImplementedError(
+                f"start_at_layer is not supported on {type(start_block).__name__} blocks: "
+                "they carry no single residual stream to swap the caller's tensor in for."
+            )
         # The returned tensor is fed to HF as inputs_embeds; stash a clone so an
         # architecture that mutates inputs_embeds in place can't corrupt the value
         # block ``start_at_layer`` swaps back in.
         injected = residual.clone()
-        for block in self.blocks:
-            block._start_at_layer_idx = start_at_layer
-            block._start_residual = injected
+        start_block._start_at_layer_idx = start_at_layer
+        start_block._start_residual = injected
         return residual
 
     def _teardown_start_at_layer(self) -> None:
@@ -3027,23 +3088,6 @@ def _encoder_decoder_setup(
             seed = torch.cat([seed, forced], dim=1)
         return seed, min_decoder_length
 
-    def _left_padded_prompt_mask(
-        self, prompts: List[str], tokens: torch.Tensor, prepend_bos: Optional[bool]
-    ) -> torch.Tensor:
-        """Attention mask for ``prompts`` left-padded into ``tokens``, from each one's length.
-
-        The tokens alone can't give this when the pad id is also a real token, as BOS is on
-        GPT-2-style tokenizers: a row's last pad and a leading BOS look the same.
-        """
-        mask = torch.zeros_like(tokens)
-        for row, prompt in enumerate(prompts):
-            n_real = self.to_tokens(
-                prompt, prepend_bos=prepend_bos, move_to_device=False, truncate=False
-            ).shape[-1]
-            # An empty prompt keeps its last pad as context, and -0 would select the whole row.
-            mask[row, -max(n_real, 1) :] = 1
-        return mask
-
     def _generate_tokens(
         self,
         current_tokens: torch.Tensor,
@@ -3590,6 +3634,7 @@ def generate(
         )
         audio_decoder_prompt: Optional[torch.Tensor] = None
         audio_timestamps = False
+        _list_prompt_mask: Optional[torch.Tensor] = None
         if getattr(self.cfg, "is_audio_seq2seq", False):
             (
                 input_tokens,
@@ -3626,7 +3671,7 @@ def generate(
                 # flush-right and logits[:, -1, :] is always the last real token.
                 # Passed as a kwarg rather than assigned: a raise between assignment
                 # and restore would pin the shared tokenizer left for the session.
-                input_tokens = self.to_tokens(
+                input_tokens, _list_prompt_mask = self._to_tokens_and_mask(
                     input,
                     prepend_bos=prepend_bos,
                     padding_side="left" if _is_batched_list else None,
@@ -3981,7 +4026,7 @@ def generate(
                 mamba_conv_kernel=mamba_conv_kernel,
                 is_encoder_decoder=is_encoder_decoder,
                 batched_prompt_mask=(
-                    self._left_padded_prompt_mask(input, input_tokens, prepend_bos)
+                    _list_prompt_mask
                     if _is_batched_list
                     and isinstance(input, list)
                     and not is_encoder_decoder
@@ -4248,6 +4293,7 @@ def generate_stream(
             self.original_model.config, "is_encoder_decoder", False
         )
         encoder_attention_mask = None
+        _list_prompt_mask: Optional[torch.Tensor] = None
         if _encdec_early:
             use_past_kv_cache = False
         if isinstance(input, str):
@@ -4268,7 +4314,7 @@ def generate_stream(
                 encoder_attention_mask = encoded.get("attention_mask")
             else:
                 # Pass the side as a kwarg, as generate() does, instead of mutating the tokenizer.
-                input_tokens = self.to_tokens(
+                input_tokens, _list_prompt_mask = self._to_tokens_and_mask(
                     input,
                     prepend_bos=prepend_bos,
                     padding_side="left" if _is_batched_list else None,
@@ -4442,7 +4488,7 @@ def _flush_held() -> Optional[Union[str, List[str]]]:
                     mamba_conv_kernel=0,
                     is_encoder_decoder=_encdec_early,
                     batched_prompt_mask=(
-                        self._left_padded_prompt_mask(input, input_tokens, prepend_bos)
+                        _list_prompt_mask
                         if _is_batched_list
                         and isinstance(input, list)
                         and not _encdec_early
diff --git a/transformer_lens/patching.py b/transformer_lens/patching.py
index 8c0dcaafac..c55e1f4507 100644
--- a/transformer_lens/patching.py
+++ b/transformer_lens/patching.py
@@ -138,9 +138,14 @@ class PatchingResult:
         has_metric_axis: Whether ``values`` carries a leading metric axis.
         patch_type_names: Names of the ``patch_type`` axis entries when the
             result came from a ``*_every`` helper (``"out"``, ``"q"``, ...).
-            GQA ``k``/``v`` sweeps are zero-padded to ``n_heads`` there, so the
-            padded cells appear as zero-valued rows in ``index_df`` and
-            :meth:`to_dataframe`.
+            GQA ``k``/``v`` sweeps are zero-padded to ``n_heads`` there.
+        valid: Bool mask over the swept cells (``values`` without the metric and
+            batch axes), ``False`` on GQA-padded ``k``/``v`` cells that no sweep
+            ran; ``None`` when nothing was padded. Padded cells stay zero in
+            ``values`` and :meth:`mean` (matching the tensor return), are ``NaN``
+            in every statistic, and are left out of :meth:`to_dataframe`;
+            ``index_df`` keeps the full grid, aligned with ``valid.flatten()``;
+            ``stat[..., valid]`` selects the live cells under any metric axis.
     """
 
     values: torch.Tensor
@@ -151,6 +156,7 @@ class PatchingResult:
     baseline: Optional[torch.Tensor] = None
     has_metric_axis: bool = False
     patch_type_names: Optional[List[str]] = None
+    valid: Optional[torch.Tensor] = None
 
     def __getitem__(self, metric_name: str) -> "PatchingResult":
         """One metric's result, with the metric axis dropped."""
@@ -168,8 +174,19 @@ def __getitem__(self, metric_name: str) -> "PatchingResult":
             baseline=None if self.baseline is None else self.baseline[i],
             has_metric_axis=False,
             patch_type_names=self.patch_type_names,
+            valid=self.valid,
         )
 
+    def _mask_invalid(self, t: torch.Tensor) -> torch.Tensor:
+        """``NaN`` on padded cells, so no statistic reports a patch that never ran."""
+        if self.valid is None:
+            return t
+        lead = 1 if self.has_metric_axis else 0
+        cells = tuple(self.valid.shape)
+        assert tuple(t.shape[lead : lead + len(cells)]) == cells
+        mask = self.valid.reshape((1,) * lead + cells + (1,) * (t.ndim - lead - len(cells)))
+        return torch.where(mask.to(t.device), t, torch.nan)
+
     def _require_per_example(self, what: str) -> None:
         if self.reduced:
             raise ValueError(f"{what} needs per-example values; run the sweep with reduce='none'")
@@ -181,7 +198,7 @@ def mean(self) -> torch.Tensor:
     def stderr(self) -> torch.Tensor:
         """Standard error of the batch mean per cell."""
         self._require_per_example("stderr")
-        return standard_error(self.values, dim=-1)
+        return self._mask_invalid(standard_error(self.values, dim=-1))
 
     def _broadcast_baseline(self) -> torch.Tensor:
         if self.baseline is None:
@@ -195,7 +212,7 @@ def _broadcast_baseline(self) -> torch.Tensor:
 
     def effect(self) -> torch.Tensor:
         """``values - baseline``: the change each patch makes to the corrupted run, per example."""
-        return self.values - self._broadcast_baseline()
+        return self._mask_invalid(self.values - self._broadcast_baseline())
 
     def normalized(self, clean_baseline: torch.Tensor) -> torch.Tensor:
         """``(values - corrupted) / (clean - corrupted)`` per example.
@@ -216,7 +233,7 @@ def normalized(self, clean_baseline: torch.Tensor) -> torch.Tensor:
                 f"got {tuple(clean.shape)}"
             )
         clean = clean.reshape(corrupted.shape)
-        return (self.values - corrupted) / (clean - corrupted)
+        return self._mask_invalid((self.values - corrupted) / (clean - corrupted))
 
     def bootstrap_ci(
         self,
@@ -234,7 +251,7 @@ def bootstrap_ci(
         """
         self._require_per_example("bootstrap_ci")
         data = self.effect() if of_effect else self.values
-        return bootstrap_ci(
+        low, high = bootstrap_ci(
             data,
             dim=-1,
             statistic=statistic,
@@ -242,6 +259,7 @@ def bootstrap_ci(
             n_resamples=n_resamples,
             generator=derive_generator(seed, "bootstrap", statistic),
         )
+        return self._mask_invalid(low), self._mask_invalid(high)
 
     def permutation_pvalue(
         self,
@@ -256,22 +274,26 @@ def permutation_pvalue(
         examples, Monte Carlo above that.
         """
         self._require_per_example("permutation_pvalue")
-        return sign_flip_permutation_pvalue(
+        pvalue = sign_flip_permutation_pvalue(
             self.effect(),
             dim=-1,
             n_permutations=n_permutations,
             generator=derive_generator(seed, "permutation"),
             alternative=alternative,
         )
+        return self._mask_invalid(pvalue)
 
     def to_dataframe(self) -> pd.DataFrame:
-        """Long form: one row per (metric, swept cell, example) with the value."""
+        """Long form: one row per (metric, swept cell, example); padded cells are left out."""
         rows = []
         n_metrics = len(self.metric_names)
         values = self.values if self.has_metric_axis else self.values.unsqueeze(0)
         flat_cells = values.reshape(n_metrics, len(self.index_df), -1).detach().cpu()
+        live = None if self.valid is None else self.valid.flatten().tolist()
         for m, name in enumerate(self.metric_names):
             for c, (_, index_row) in enumerate(self.index_df.iterrows()):
+                if live is not None and not live[c]:
+                    continue
                 for b in range(flat_cells.shape[-1]):
                     row = {"metric": name, **index_row.to_dict()}
                     if self.patch_type_names is not None and "patch_type" in row:
@@ -1044,8 +1066,8 @@ def _stack_patch_types(
     ``axis_order`` (``head_index`` / ``dest_pos`` renamed to ``head`` / ``pos``)
     and the head axis zero-padded to ``n_heads`` for GQA ``k``/``v`` sweeps,
     keeping the metric axis first and the batch axis last. ``index_df`` is
-    rebuilt from the padded grid so it stays one row per cell of ``values``;
-    padded GQA cells therefore appear as zero-valued rows.
+    rebuilt from the padded grid so it stays one row per cell of ``values``, and
+    ``valid`` marks the padded cells (``None`` when no sweep was padded).
     """
     if all(isinstance(r, torch.Tensor) for r in results):
         return torch.stack([r for r in results if isinstance(r, torch.Tensor)], dim=0)
@@ -1058,6 +1080,7 @@ def _stack_patch_types(
             raise ValueError("every sweep must use the same metrics and reduce mode to be stacked")
     head_axis = list(axis_order).index("head") if "head" in axis_order else None
     values = []
+    widths = []
     for r in typed:
         names = [_RESULT_AXIS_RENAMES.get(n, n) for n in r.axis_names]
         if sorted(names) != sorted(axis_order):
@@ -1067,6 +1090,7 @@ def _stack_patch_types(
         tail = list(range(lead + len(names), r.values.ndim))
         v = r.values.permute(*perm, *tail)
         pad = 0 if head_axis is None else n_heads - v.shape[lead + head_axis]
+        widths.append(n_heads - pad)
         if head_axis is not None and pad > 0:
             pad_shape = list(v.shape)
             pad_shape[lead + head_axis] = pad
@@ -1075,6 +1099,11 @@ def _stack_patch_types(
     stacked = torch.stack(values, dim=1 if first.has_metric_axis else 0)
     lead = 1 if first.has_metric_axis else 0
     cell_shape = list(stacked.shape[lead : lead + 1 + len(axis_order)])
+    valid = None
+    if head_axis is not None and min(widths) < n_heads:
+        valid = torch.ones(cell_shape, dtype=torch.bool, device=stacked.device)
+        for type_i, width in enumerate(widths):
+            valid[type_i].narrow(head_axis, width, n_heads - width).fill_(False)
     return PatchingResult(
         values=stacked,
         axis_names=["patch_type", *axis_order],
@@ -1084,6 +1113,7 @@ def _stack_patch_types(
         baseline=first.baseline,
         has_metric_axis=first.has_metric_axis,
         patch_type_names=list(type_names),
+        valid=valid,
     )
 
 
diff --git a/transformer_lens/tools/analysis/_model_state.py b/transformer_lens/tools/analysis/_model_state.py
index 4644ce1233..b4640b8d6d 100644
--- a/transformer_lens/tools/analysis/_model_state.py
+++ b/transformer_lens/tools/analysis/_model_state.py
@@ -5,8 +5,6 @@
 
 import torch
 
-# Block variants with no attn_out + mlp_out split and no resid_mid hook.
-_HYBRID_VARIANT_NAMES = ("mamba", "ssm", "mixer", "linear_attn")
 _ID_DTYPES = (torch.int64, torch.int32, torch.int16, torch.int8, torch.uint8)
 
 
@@ -36,12 +34,29 @@ def require_eval_mode(model: Any, *, operation: str) -> None:
     )
 
 
-def hybrid_block_types(model: Any) -> List[str]:
-    """Sorted distinct ``model.layer_types()`` entries naming an SSM or linear-attention block."""
-    layer_types = model.layer_types() if hasattr(model, "layer_types") else []
-    return sorted(
-        {t for t in layer_types if any(part in _HYBRID_VARIANT_NAMES for part in t.split("+"))}
-    )
+def blocks_without_resid_mid(model: Any) -> List[int]:
+    """Indices of sequential blocks that expose no ``hook_resid_mid``."""
+    blocks = getattr(model, "blocks", None)
+    hook_dict = getattr(model, "hook_dict", None)
+    if blocks is None or not isinstance(hook_dict, dict):
+        return []
+    return [
+        l
+        for l, block in enumerate(blocks)
+        # Parallel blocks have no resid_mid by design and accumulated_resid skips them.
+        if getattr(block, "mlp_input_includes_attn", True)
+        and f"blocks.{l}.hook_resid_mid" not in hook_dict
+    ]
+
+
+def require_resid_mid(model: Any, *, operation: str, alternative: str) -> None:
+    """Refuse ``operation`` when a sequential block has no ``hook_resid_mid`` to read."""
+    missing = blocks_without_resid_mid(model)
+    if missing:
+        raise ValueError(
+            f"{operation} reads hook_resid_mid on every sequential block, but blocks "
+            f"{missing} have none (the sequence mixer is their only sublayer); {alternative}"
+        )
 
 
 def compute_dtype(dtype: torch.dtype) -> torch.dtype:
diff --git a/transformer_lens/tools/analysis/direct_logit_attribution.py b/transformer_lens/tools/analysis/direct_logit_attribution.py
index d1cfc92a02..b5a43aa08a 100644
--- a/transformer_lens/tools/analysis/direct_logit_attribution.py
+++ b/transformer_lens/tools/analysis/direct_logit_attribution.py
@@ -39,15 +39,15 @@
 """
 
 from dataclasses import dataclass
-from typing import List, Optional, Union
+from typing import Any, List, Optional, Union
 
 import torch
 from jaxtyping import Float
 
 from transformer_lens.ActivationCache import ActivationCache
 from transformer_lens.tools.analysis._model_state import (
-    hybrid_block_types,
     require_eval_mode,
+    require_resid_mid,
 )
 from transformer_lens.tools.analysis.logit_lens import (
     _final_norm,
@@ -135,15 +135,54 @@ def _residual_stack_and_labels(
     raise ValueError(f"unit must be one of {_VALID_UNITS}, got {unit!r}")
 
 
+def _require_every_writer_decomposed(model: Any) -> None:
+    """Refuse DLA while ``decompose_resid`` would omit a sequence-mixer slot a block writes through.
+
+    A block that runs a mixer beside attention adds both outputs to the residual stream;
+    if ``decompose_resid`` emits a row for only one, the attribution would not sum to
+    the logit and nothing else would notice. Structural, so it runs before any forward.
+    """
+    from transformer_lens.model_bridge.generalized_components.block import (
+        VARIANT_SUBMODULE_NAMES,
+    )
+
+    has_blocks = getattr(model, "_has_registered_blocks", None)
+    if callable(has_blocks) and not has_blocks():
+        return
+    blocks = getattr(model, "blocks", None)
+    if blocks is None:
+        return
+    # An empty cache reads the labels from the block layout alone.
+    layout = ActivationCache({}, model)
+    emitted = set(layout.decompose_resid_labels(incl_embeds=False))
+    for layer, block in enumerate(blocks):
+        modules = getattr(block, "_modules", None)
+        if not isinstance(modules, dict):
+            continue
+        for slot in VARIANT_SUBMODULE_NAMES:
+            if slot not in modules:
+                continue
+            label = f"{layer}_attn_out" if slot == "attn" else layout._mixer_label(layer, slot)
+            if label not in emitted:
+                raise ValueError(
+                    f"DLA unit='component' cannot attribute block {layer}: its {slot} output "
+                    f"is not decomposed yet (decompose_resid emits no {label!r}), so the "
+                    "attribution would not sum to the logit. Use unit='head' or unit='layer'."
+                )
+
+
 def _validate_bridge_compatibility(model) -> bool:
     """Reject Bridges DLA can't produce correct numbers for; return whether ``model`` is one.
 
     Compatibility mode alone does not guarantee the final norm is folded:
     ``fold_ln=False`` and ``no_processing=True`` leave its affine parameters
-    active. The hybrid-arch check catches Mamba/SSM blocks before decomposition.
+    active.
     """
     # Lazy import — keeps the module importable without dragging in the bridge.
     from transformer_lens.model_bridge import TransformerBridge
+    from transformer_lens.model_bridge.generalized_components.normalization import (
+        norm_uses_weight_offset,
+    )
 
     if not isinstance(model, TransformerBridge):
         return False
@@ -164,7 +203,7 @@ def _validate_bridge_compatibility(model) -> bool:
         if isinstance(weight, torch.Tensor):
             identity = (
                 torch.zeros_like(weight)
-                if getattr(model.cfg, "rmsnorm_uses_offset", False)
+                if norm_uses_weight_offset(final_norm, model.cfg)
                 else torch.ones_like(weight)
             )
             if not torch.equal(weight, identity):
@@ -181,14 +220,6 @@ def _validate_bridge_compatibility(model) -> bool:
                 "with fold_ln=True and no_processing=False."
             )
 
-    # Refused up front: decompose_resid would otherwise fail with a confusing KeyError.
-    hybrid = hybrid_block_types(model)
-    if hybrid:
-        raise NotImplementedError(
-            f"DLA does not yet support hybrid architectures (found block types {hybrid}). "
-            f"Only standard attention + MLP transformers (e.g. GPT-2, LLaMA, Pythia) are "
-            f"supported; hybrid support requires extending ActivationCache.decompose_resid."
-        )
     return True
 
 
@@ -221,13 +252,12 @@ def direct_logit_attribution(
     normalization's affine parameters must be folded into ``W_U`` and ``b_U``
     (or already be identity);
     otherwise the projection direction is wrong and DLA would return silently
-    incorrect numbers. Hybrid architectures
-    (Mamba/SSM/Mixer/LinearAttention) are not yet supported because
-    ``decompose_resid`` only understands the ``attn_out + mlp_out`` block layout;
-    both conditions raise an explicit error at call time, as does a Bridge whose
-    readout is not ``ln_final -> unembed`` on the last block's output; a prompt run
-    here is also checked against the model's own output, a precomputed ``cache``
-    only structurally.
+    incorrect numbers; that raises an explicit error at call time, as does a Bridge
+    whose readout is not ``ln_final -> unembed`` on the last block's output; a prompt
+    run here is also checked against the model's own output, a precomputed ``cache``
+    only structurally. On hybrid and SSM stacks, sequence-mixer blocks contribute their
+    mixer output (e.g. ``0_linear_attn_out``); ``unit="layer"`` needs a ``resid_mid``
+    hook on every sequential block, which mixer-only blocks (Mamba, NemotronH) lack.
 
     Args:
         model:
@@ -270,9 +300,10 @@ def direct_logit_attribution(
             ``TransformerBridge`` is passed without compatibility mode enabled,
             with an active learned final-norm weight or bias, or with a readout
             that is not ``ln_final -> unembed``; also if the model must run while
-            any module is in training mode (a supplied ``cache`` is not checked).
-        NotImplementedError: If a ``TransformerBridge`` reports a hybrid block
-            layout (Mamba/SSM/Mixer/LinearAttention).
+            any module is in training mode (a supplied ``cache`` is not checked);
+            also if ``unit="layer"`` and a sequential block has no ``hook_resid_mid``,
+            or ``unit="component"`` and ``decompose_resid`` would omit a block's
+            sequence-mixer output (both checked before any forward).
     """
     if unit not in _VALID_UNITS:
         raise ValueError(f"unit must be one of {_VALID_UNITS}, got {unit!r}")
@@ -280,6 +311,17 @@ def direct_logit_attribution(
         raise ValueError("answer_tokens is required")
 
     is_bridge = _validate_bridge_compatibility(model)
+    if unit == "layer":
+        require_resid_mid(
+            model,
+            operation="DLA unit='layer'",
+            alternative=(
+                "use unit='component', or logit_lens(..., incl_mid=False) for the per-block "
+                "trajectory"
+            ),
+        )
+    elif unit == "component":
+        _require_every_writer_decomposed(model)
 
     if cache is None:
         if input is None:
diff --git a/transformer_lens/tools/analysis/jacobian_lens.py b/transformer_lens/tools/analysis/jacobian_lens.py
index 5d1d21912a..eacf8f3519 100644
--- a/transformer_lens/tools/analysis/jacobian_lens.py
+++ b/transformer_lens/tools/analysis/jacobian_lens.py
@@ -98,7 +98,10 @@
     estimate_occupancy,
     get_sparse_decomposition,
 )
-from transformer_lens.tools.analysis.logit_lens import logit_readout
+from transformer_lens.tools.analysis.logit_lens import (
+    _require_direct_readout,
+    logit_readout,
+)
 from transformer_lens.tools.analysis.projection_kernel import orthonormal_subspace
 from transformer_lens.utilities.hf_utils import call_hf_with_retry
 
@@ -1900,18 +1903,7 @@ def _require_raw_bridge(model: Any, *, estimator: str) -> None:
             f"{estimator} requires the standard direct ln_final -> unembed output path; "
             f"this Bridge is missing {missing_components}."
         )
-    blocks_component = component_mapping["blocks"]
-    if not getattr(blocks_component, "hook_out_is_single_residual_stream", False):
-        raise ValueError(
-            f"{estimator} requires single-stream [batch, position, d_model] block "
-            f"outputs; {type(blocks_component).__name__} does not provide that "
-            "contract."
-        )
-    if "project_out" in component_mapping:
-        raise ValueError(
-            f"{estimator} does not yet support a final output projection between "
-            "the residual stream and unembedding."
-        )
+    _require_direct_readout(model, operation=estimator)
     unembed_width = model.W_U.shape[0]
     if unembed_width != model.cfg.d_model:
         raise ValueError(
diff --git a/transformer_lens/tools/analysis/logit_lens.py b/transformer_lens/tools/analysis/logit_lens.py
index b11de9d17b..bd65ec33f3 100644
--- a/transformer_lens/tools/analysis/logit_lens.py
+++ b/transformer_lens/tools/analysis/logit_lens.py
@@ -41,8 +41,8 @@
 
 from transformer_lens.ActivationCache import ActivationCache
 from transformer_lens.tools.analysis._model_state import (
-    hybrid_block_types,
     require_eval_mode,
+    require_resid_mid,
     to_token_ids,
 )
 from transformer_lens.utilities import Slice, SliceInput
@@ -686,8 +686,9 @@ def logit_lens(
         cache: An existing cache to read instead of running the model.
         layers: Indices into the accumulated stack (negative allowed); ``None``
             keeps every entry.
-        incl_mid: Include ``resid_mid`` entries. Raises on hybrid stacks whose
-            SSM blocks have no ``resid_mid``.
+        incl_mid: Include ``resid_mid`` entries. Raises when a sequential block has
+            no ``resid_mid`` (Mamba and other mixer-only blocks); parallel blocks are
+            skipped.
         positions: Position selection (``Slice`` semantics); default last token.
         batch_slice: Batch selection (``Slice`` semantics).
         lens: Optional ``f(entry, index) -> entry`` applied to each stack entry
@@ -721,12 +722,9 @@ def logit_lens(
         chunk_size=chunk_size,
     )
     if incl_mid:
-        hybrid = hybrid_block_types(model)
-        if hybrid:
-            raise ValueError(
-                f"incl_mid=True needs a resid_mid hook on every block, but this stack has "
-                f"block types {hybrid} without one; use incl_mid=False."
-            )
+        require_resid_mid(
+            model, operation="logit_lens(incl_mid=True)", alternative="use incl_mid=False"
+        )
     pos_slice = positions if isinstance(positions, Slice) else Slice(positions)
     if cache is None:
         if input is None:
diff --git a/transformer_lens/utilities/statistics.py b/transformer_lens/utilities/statistics.py
index d387cd7460..66a736bee4 100644
--- a/transformer_lens/utilities/statistics.py
+++ b/transformer_lens/utilities/statistics.py
@@ -9,15 +9,15 @@
 from __future__ import annotations
 
 import hashlib
-import itertools
 import math
 from dataclasses import dataclass
-from typing import Optional, Sequence, Tuple
+from typing import Iterator, Optional, Sequence, Tuple
 
 import torch
 
 _ALTERNATIVES = ("two-sided", "greater", "less")
 _STATISTICS = ("mean", "median")
+_SIGN_FLIP_CHUNK = 256
 
 
 def derive_generator(seed: int, *labels: object) -> torch.Generator:
@@ -114,8 +114,10 @@ def sign_flip_permutation_pvalue(
     ``alternative`` is ``"two-sided"`` (``|null| >= |observed|``), ``"greater"``
     or ``"less"``. A null mean within a relative ``100 * eps`` of the observed one
     counts as a tie, absorbing summation-order error; there is no absolute floor, so
-    near-zero observed effects get essentially no tolerance. Returns a tensor with
-    ``dim`` removed.
+    near-zero observed effects get essentially no tolerance. Sign patterns are
+    evaluated 256 at a time in both modes, so memory stays at ``256 × n`` per
+    leading cell. A leading cell with any non-finite effect gets a NaN p-value. Returns a
+    tensor with ``dim`` removed.
     """
     if alternative not in _ALTERNATIVES:
         raise ValueError(f"alternative must be one of {_ALTERNATIVES}, got {alternative!r}")
@@ -127,23 +129,47 @@ def sign_flip_permutation_pvalue(
         raise ValueError("cannot permute an empty axis")
     observed = moved.mean(dim=-1, keepdim=True)  # [..., 1]
     if n <= exact_max_n:
-        signs = torch.tensor(
-            list(itertools.product((1.0, -1.0), repeat=n)), dtype=moved.dtype, device=moved.device
-        )  # [2^n, n]
-        null = (moved.unsqueeze(-2) * signs).mean(dim=-1)  # [..., 2^n]
-        hits = _count_hits(null, observed, alternative)
-        return hits / signs.shape[0]
-    gen = generator if generator is not None else derive_generator(0)
-    hits_total = torch.zeros(moved.shape[:-1], dtype=moved.dtype, device=moved.device)
-    chunk = 256
-    for start in range(0, n_permutations, chunk):
-        count = min(chunk, n_permutations - start)
-        signs = (
-            torch.randint(0, 2, (count, n), generator=gen).to(moved.device).to(moved.dtype) * 2 - 1
-        )
-        null = (moved.unsqueeze(-2) * signs).mean(dim=-1)
-        hits_total = hits_total + _count_hits(null, observed, alternative)
-    return (hits_total + 1.0) / (n_permutations + 1.0)
+        hits = _sign_flip_hits(moved, observed, alternative, _exact_sign_chunks(n, moved))
+        pvalue = hits / 2**n
+    else:
+        gen = generator if generator is not None else derive_generator(0)
+        signs = _random_sign_chunks(n, n_permutations, gen, moved)
+        hits = _sign_flip_hits(moved, observed, alternative, signs)
+        pvalue = (hits + 1.0) / (n_permutations + 1.0)
+    # NaN compares false against every null mean and an infinite effect makes the tie
+    # tolerance infinite, so either would read as a maximally significant 0.
+    return pvalue.masked_fill((~moved.isfinite()).any(dim=-1), math.nan)
+
+
+def _sign_flip_hits(
+    moved: torch.Tensor,
+    observed: torch.Tensor,
+    alternative: str,
+    sign_chunks: Iterator[torch.Tensor],
+) -> torch.Tensor:
+    hits = torch.zeros(moved.shape[:-1], dtype=moved.dtype, device=moved.device)
+    for signs in sign_chunks:
+        null = (moved.unsqueeze(-2) * signs).mean(dim=-1)  # [..., chunk]
+        hits = hits + _count_hits(null, observed, alternative)
+    return hits
+
+
+def _exact_sign_chunks(n: int, like: torch.Tensor) -> Iterator[torch.Tensor]:
+    # Bit i of the row index, most significant first, reproduces itertools.product((1, -1)) order.
+    shifts = torch.arange(n - 1, -1, -1, device=like.device)
+    for start in range(0, 2**n, _SIGN_FLIP_CHUNK):
+        rows = torch.arange(start, min(start + _SIGN_FLIP_CHUNK, 2**n), device=like.device)
+        bits = (rows.unsqueeze(-1) >> shifts) & 1
+        yield (1 - 2 * bits).to(like.dtype)
+
+
+def _random_sign_chunks(
+    n: int, n_permutations: int, gen: torch.Generator, like: torch.Tensor
+) -> Iterator[torch.Tensor]:
+    for start in range(0, n_permutations, _SIGN_FLIP_CHUNK):
+        count = min(_SIGN_FLIP_CHUNK, n_permutations - start)
+        signs = torch.randint(0, 2, (count, n), generator=gen).to(like.device).to(like.dtype)
+        yield signs * 2 - 1
 
 
 def _count_hits(null: torch.Tensor, observed: torch.Tensor, alternative: str) -> torch.Tensor: