Fix DeepSeek V2 dense MLP hook semantics

tests/integration/model_bridge/test_deepseek_v2_adapter.py

@@ -48,6 +48,14 @@ def tiny_deepseek_v2_bridge():
     return _make_bridge(q_lora_rank=64)
 
 
+@pytest.fixture(scope="module")
+def tiny_deepseek_v2_bridge_compat():
+    """V2-full with compatibility mode enabled without mutating shared state."""
+    bridge = _make_bridge(q_lora_rank=64)
+    bridge.enable_compatibility_mode(no_processing=True)
+    return bridge
+
+
 @pytest.fixture(scope="module")
 def tiny_deepseek_v2_lite_bridge():
     """V2-Lite: q_lora_rank=None — direct Q projection, no LoRA compression."""
@@ -124,6 +132,43 @@ def test_all_layers_have_mlp_hooks(self, tiny_deepseek_v2_bridge):
             assert f"blocks.{i}.mlp.hook_out" in cache
             assert not torch.isnan(cache[f"blocks.{i}.mlp.hook_out"]).any()
 
+    def test_dense_layer_compatibility_hooks_use_neuron_basis(self, tiny_deepseek_v2_bridge_compat):
+        dense_mlp = tiny_deepseek_v2_bridge_compat.original_model.model.layers[0].mlp
+        captured = {}
+        handles = [
+            dense_mlp.gate_proj.register_forward_hook(
+                lambda _module, _args, output: captured.__setitem__("pre", output.detach())
+            ),
+            dense_mlp.up_proj.register_forward_hook(
+                lambda _module, _args, output: captured.__setitem__("pre_linear", output.detach())
+            ),
+            dense_mlp.down_proj.register_forward_pre_hook(
+                lambda _module, args: captured.__setitem__("post", args[0].detach())
+            ),
+        ]
+        try:
+            _, cache = tiny_deepseek_v2_bridge_compat.run_with_cache(_tokens())
+        finally:
+            for handle in handles:
+                handle.remove()
+
+        for hook, expected in {
+            "hook_pre": "pre",
+            "hook_pre_linear": "pre_linear",
+            "hook_post": "post",
+        }.items():
+            actual = cache[f"blocks.0.mlp.{hook}"]
+            assert actual.shape[-1] == 512
+            torch.testing.assert_close(actual, captured[expected])
+
+    def test_sparse_layer_compatibility_hooks_remain_block_boundaries(
+        self, tiny_deepseek_v2_bridge_compat
+    ):
+        _, cache = tiny_deepseek_v2_bridge_compat.run_with_cache(_tokens())
+        assert torch.equal(cache["blocks.1.mlp.hook_pre"], cache["blocks.1.mlp.hook_in"])
+        assert torch.equal(cache["blocks.1.mlp.hook_post"], cache["blocks.1.mlp.hook_out"])
+        assert "blocks.1.mlp.hook_pre_linear" not in cache
+
 
 class TestDeepSeekV2AttentionHooks:
     def test_attn_hooks_fire_all_layers(self, tiny_deepseek_v2_bridge):

tests/unit/model_bridge/supported_architectures/test_deepseek_v2_adapter.py

@@ -220,7 +220,24 @@ def test_moe_submodule_keys(self, adapter: DeepSeekV2ArchitectureAdapter) -> Non
         """Gate and shared experts are bridged; DeepseekV2Moe.forward calls self.gate,
         so its routing logits are hookable."""
         mlp = adapter.component_mapping["blocks"].submodules["mlp"]
-        assert set(mlp.submodules.keys()) == {"gate", "shared_experts"}
+        assert set(mlp.submodules.keys()) == {
+            "gate",
+            "shared_experts",
+            "dense_gate",
+            "dense_in",
+            "dense_out",
+        }
+
+    def test_dense_projections_are_optional(self, adapter: DeepSeekV2ArchitectureAdapter) -> None:
+        mlp = adapter.component_mapping["blocks"].submodules["mlp"]
+        for key, path in {
+            "dense_gate": "gate_proj",
+            "dense_in": "up_proj",
+            "dense_out": "down_proj",
+        }.items():
+            assert isinstance(mlp.submodules[key], LinearBridge)
+            assert mlp.submodules[key].name == path
+            assert mlp.submodules[key].optional is True
 
     def test_shared_experts_is_optional(self, adapter: DeepSeekV2ArchitectureAdapter) -> None:
         """Dense layers (idx < first_k_dense_replace) have no shared_experts."""

transformer_lens/model_bridge/supported_architectures/deepseek_v2.py

@@ -29,6 +29,16 @@
 )
 
 
+class _DeepSeekV2MLPBridge(MoEBridge):
+    """Expose neuron-basis hooks on dense-prefix layers and MoE boundaries elsewhere."""
+
+    hook_aliases = {
+        "hook_pre": ["dense_gate.hook_out", "hook_in"],
+        "hook_pre_linear": "dense_in.hook_out",
+        "hook_post": ["dense_out.hook_in", "hook_out"],
+    }
+
+
 class DeepSeekV2ArchitectureAdapter(ArchitectureAdapter):
     """Architecture adapter for DeepSeek V2 / V2-Lite / Coder-V2 models.
 
@@ -100,12 +110,15 @@ def __init__(self, cfg: Any) -> None:
 
     def _build_mlp_bridge(self):
         """Routed MoE with optional shared experts; Youtu (all-dense) overrides."""
-        return MoEBridge(
+        return _DeepSeekV2MLPBridge(
             name="mlp",
             config=self.cfg,
             submodules={
                 # Router is a custom Module, not nn.Linear.
                 "gate": GeneralizedComponent(name="gate", optional=True),
                 "shared_experts": self._gated_mlp(name="shared_experts", optional=True),
+                "dense_gate": LinearBridge(name="gate_proj", optional=True),
+                "dense_in": LinearBridge(name="up_proj", optional=True),
+                "dense_out": LinearBridge(name="down_proj", optional=True),
             },
         )