fix: preserve per-prompt MoE router observations

tests/test_abliterate.py

@@ -2186,21 +2186,152 @@ class TestRouterProfilingHooks:
             # Simulate harmful forward pass
             pipeline._routing_is_harmful = True
             x = torch.randn(1, 5, 16)
+            harmful_expected = torch.nn.functional.linear(x, layer.mlp.gate.weight)[0, -1]
             layer.mlp.gate(x)  # triggers hook
 
             assert len(pipeline._routing_harmful[0]) == 1
             assert pipeline._routing_harmful[0][0].shape[0] == 4  # n_experts
+            assert torch.equal(pipeline._routing_harmful[0][0], harmful_expected)
 
             # Simulate harmless forward pass
             pipeline._routing_is_harmful = False
-            layer.mlp.gate(x)
+            harmless_x = x + 1
+            harmless_expected = torch.nn.functional.linear(
+                harmless_x, layer.mlp.gate.weight,
+            )[0, -1]
+            layer.mlp.gate(harmless_x)
 
             assert len(pipeline._routing_harmless[0]) == 1
+            assert torch.equal(pipeline._routing_harmless[0][0], harmless_expected)
+            assert not torch.equal(
+                pipeline._routing_harmful[0][0], pipeline._routing_harmless[0][0],
+            )
         finally:
             for h in hooks:
                 h.remove()
             abl_module.get_ffn_module = orig_get_ffn
 
+    def test_hooks_record_single_prompt_vector_output(self):
+        """A one-dimensional single-prompt router output should remain intact."""
+        pipeline, layers, layer, abl_module, orig_get_ffn = self._make_moe_pipeline_and_layers()
+        hooks = []
+        try:
+            hooks = pipeline._install_router_profiling_hooks(layers)
+            x = torch.arange(16, dtype=torch.float32) / 100
+            expected = torch.nn.functional.linear(x, layer.mlp.gate.weight)
+
+            layer.mlp.gate(x)
+
+            assert len(pipeline._routing_harmful[0]) == 1
+            assert torch.equal(pipeline._routing_harmful[0][0], expected)
+        finally:
+            for h in hooks:
+                h.remove()
+            abl_module.get_ffn_module = orig_get_ffn
+
+    def test_hooks_record_one_vector_for_each_3d_batch_prompt(self):
+        """Native 3-D router outputs should retain every batch member."""
+        pipeline, layers, layer, abl_module, orig_get_ffn = self._make_moe_pipeline_and_layers()
+        hooks = []
+        try:
+            hooks = pipeline._install_router_profiling_hooks(layers)
+            x = torch.arange(2 * 5 * 16, dtype=torch.float32).reshape(2, 5, 16) / 100
+            expected = torch.nn.functional.linear(x, layer.mlp.gate.weight)
+
+            layer.mlp.gate(x)
+
+            recorded = pipeline._routing_harmful[0]
+            assert len(recorded) == 2
+            assert torch.equal(recorded[0], expected[0, -1])
+            assert torch.equal(recorded[1], expected[1, -1])
+        finally:
+            for h in hooks:
+                h.remove()
+            abl_module.get_ffn_module = orig_get_ffn
+
+    def test_hooks_preserve_flattened_prompt_association_with_padding(self):
+        """Flattened token rows should become one vector per padded prompt."""
+        pipeline, layers, layer, abl_module, orig_get_ffn = self._make_moe_pipeline_and_layers()
+        hooks = []
+        try:
+            hooks = pipeline._install_router_profiling_hooks(layers)
+            pipeline._routing_attention_mask = torch.tensor(
+                [[1, 1, 0, 0], [1, 1, 1, 1]], dtype=torch.long,
+            )
+            x = torch.arange(8 * 16, dtype=torch.float32).reshape(8, 16) / 100
+            expected = torch.nn.functional.linear(x, layer.mlp.gate.weight).reshape(2, 4, 4)
+
+            layer.mlp.gate(x)
+
+            recorded = pipeline._routing_harmful[0]
+            assert len(recorded) == 2
+            assert torch.equal(recorded[0], expected[0, 1])
+            assert torch.equal(recorded[1], expected[1, 3])
+        finally:
+            pipeline._routing_attention_mask = None
+            for h in hooks:
+                h.remove()
+            abl_module.get_ffn_module = orig_get_ffn
+
+    def test_hooks_average_only_valid_cot_tokens_per_prompt(self):
+        """CoT aggregation should exclude padding without mixing prompts."""
+        pipeline, layers, layer, abl_module, orig_get_ffn = self._make_moe_pipeline_and_layers()
+        hooks = []
+        try:
+            hooks = pipeline._install_router_profiling_hooks(layers)
+            pipeline.cot_aware = True
+            pipeline._routing_attention_mask = torch.tensor(
+                [[1, 1, 1, 1, 1, 0], [0, 1, 1, 1, 1, 1]], dtype=torch.long,
+            )
+            x = torch.arange(12 * 16, dtype=torch.float32).reshape(12, 16) / 100
+            expected = torch.nn.functional.linear(x, layer.mlp.gate.weight).reshape(2, 6, 4)
+
+            layer.mlp.gate(x)
+
+            recorded = pipeline._routing_harmful[0]
+            assert len(recorded) == 2
+            assert torch.allclose(recorded[0], expected[0, :5].mean(dim=0))
+            assert torch.allclose(recorded[1], expected[1, 1:].mean(dim=0))
+        finally:
+            pipeline._routing_attention_mask = None
+            for h in hooks:
+                h.remove()
+            abl_module.get_ffn_module = orig_get_ffn
+
+    def test_hooks_skip_ambiguous_flattened_logits(self):
+        """Ambiguous flattened layouts should fail visibly instead of collapsing."""
+        pipeline, layers, layer, abl_module, orig_get_ffn = self._make_moe_pipeline_and_layers()
+        hooks = []
+        try:
+            hooks = pipeline._install_router_profiling_hooks(layers)
+
+            with pytest.warns(RuntimeWarning, match="multi-row 2-D router logits are ambiguous"):
+                layer.mlp.gate(torch.randn(5, 16))
+
+            assert pipeline._routing_harmful[0] == []
+        finally:
+            for h in hooks:
+                h.remove()
+            abl_module.get_ffn_module = orig_get_ffn
+
+    def test_hooks_skip_mismatched_flattened_logits(self):
+        """A stale or incompatible mask should produce an actionable warning."""
+        pipeline, layers, layer, abl_module, orig_get_ffn = self._make_moe_pipeline_and_layers()
+        hooks = []
+        try:
+            hooks = pipeline._install_router_profiling_hooks(layers)
+            pipeline._routing_attention_mask = torch.ones(2, 4, dtype=torch.long)
+
+            with pytest.warns(RuntimeWarning, match="do not match batch × sequence"):
+                layer.mlp.gate(torch.randn(7, 16))
+
+            assert pipeline._routing_harmful[0] == []
+        finally:
+            pipeline._routing_attention_mask = None
+            for h in hooks:
+                h.remove()
+            abl_module.get_ffn_module = orig_get_ffn
+
     def test_no_handle_returns_empty(self):
         """Should return empty list when handle is None."""
         pipeline = AbliterationPipeline(model_name="test", method="surgical")