fix: preserve per-prompt MoE router observations

obliteratus/abliterate.py

@@ -1033,6 +1033,10 @@ class AbliterationPipeline:
         self._routing_harmful: dict[int, list[torch.Tensor]] = {}
         self._routing_harmless: dict[int, list[torch.Tensor]] = {}
         self._routing_is_harmful: bool = True  # flag for routing hooks
+        # Attention mask for the synchronous forward pass currently observed
+        # by router hooks. Flattened MoE routers need this shape information
+        # to recover the original (batch, sequence) prompt association.
+        self._routing_attention_mask: torch.Tensor | None = None
 
     def log(self, msg: str):
         self._on_log(msg)
@@ -1126,17 +1130,39 @@ class AbliterationPipeline:
 
             def make_hook(layer_idx: int):
                 def hook_fn(module, input, output):
-                    logits = output if isinstance(output, torch.Tensor) else output[0]
-                    # Extract router logits — use mean across positions for
-                    # CoT-aware models so we capture expert routing at reasoning
-                    # tokens, not just the final output token.
-                    if logits.dim() == 3:
-                        if getattr(self, "cot_aware", False) and logits.shape[1] > 4:
-                            logits = logits.mean(dim=1)  # (batch, num_experts)
-                        else:
-                            logits = logits[:, -1, :]  # (batch, num_experts)
-                    elif logits.dim() == 2 and logits.shape[0] > 1:
-                        logits = logits[-1:, :]
+                    if isinstance(output, torch.Tensor):
+                        logits = output
+                    elif (
+                        isinstance(output, (tuple, list))
+                        and output
+                        and isinstance(output[0], torch.Tensor)
+                    ):
+                        logits = output[0]
+                    else:
+                        warnings.warn(
+                            f"Skipping router profiling for layer {layer_idx}: "
+                            "router output does not contain a tensor",
+                            RuntimeWarning,
+                            stacklevel=2,
+                        )
+                        return
+                    router_input = (
+                        input[0] if input and isinstance(input[0], torch.Tensor) else None
+                    )
+                    try:
+                        logits = self._router_logits_per_prompt(
+                            logits,
+                            router_input=router_input,
+                            attention_mask=self._routing_attention_mask,
+                            cot_aware=getattr(self, "cot_aware", False),
+                        )
+                    except ValueError as exc:
+                        warnings.warn(
+                            f"Skipping router profiling for layer {layer_idx}: {exc}",
+                            RuntimeWarning,
+                            stacklevel=2,
+                        )
+                        return
                     target = (self._routing_harmful
                               if self._routing_is_harmful
                               else self._routing_harmless)
@@ -1156,6 +1182,86 @@ class AbliterationPipeline:
             self.log(f"  Router profiling hooks installed on {len(hooks)} MoE layers")
         return hooks
 
+    @staticmethod
+    def _router_logits_per_prompt(
+        logits: torch.Tensor,
+        *,
+        router_input: torch.Tensor | None,
+        attention_mask: torch.Tensor | None,
+        cot_aware: bool,
+    ) -> torch.Tensor:
+        """Return one router-logit vector per prompt without mixing batches.
+
+        Some MoE implementations expose router logits as ``(batch, sequence,
+        experts)`` while others flatten the first two dimensions. A flattened
+        tensor is reshaped only when the active attention mask or a 3-D router
+        input proves the original dimensions. Ambiguous layouts are rejected
+        instead of silently treating a token as an entire prompt.
+        """
+        if not isinstance(logits, torch.Tensor):
+            raise ValueError("router output does not contain a tensor")
+        if logits.dim() == 1:
+            return logits.unsqueeze(0)
+        if logits.dim() not in (2, 3):
+            raise ValueError(
+                f"unsupported router-logit shape {tuple(logits.shape)}; "
+                "expected 1-D, 2-D, or 3-D"
+            )
+
+        mask = attention_mask
+        if mask is not None:
+            if mask.dim() != 2:
+                raise ValueError(
+                    f"attention mask has shape {tuple(mask.shape)}; "
+                    "expected (batch, sequence)"
+                )
+            mask = mask.to(device=logits.device, dtype=torch.bool)
+
+        if logits.dim() == 2:
+            rows, experts = logits.shape
+            batch_seq: tuple[int, int] | None = None
+            if mask is not None:
+                batch_seq = (mask.shape[0], mask.shape[1])
+            elif router_input is not None and router_input.dim() == 3:
+                batch_seq = (router_input.shape[0], router_input.shape[1])
+            elif rows == 1:
+                return logits
+            else:
+                raise ValueError(
+                    "multi-row 2-D router logits are ambiguous without a 2-D "
+                    "attention mask or 3-D router input"
+                )
+
+            batch, sequence = batch_seq
+            if rows != batch * sequence:
+                raise ValueError(
+                    f"router-logit rows ({rows}) do not match batch × sequence "
+                    f"({batch} × {sequence})"
+                )
+            logits = logits.reshape(batch, sequence, experts)
+
+        batch, sequence, _ = logits.shape
+        if mask is None:
+            mask = torch.ones((batch, sequence), dtype=torch.bool, device=logits.device)
+        elif tuple(mask.shape) != (batch, sequence):
+            raise ValueError(
+                f"attention mask shape {tuple(mask.shape)} does not match router logits "
+                f"{(batch, sequence)}"
+            )
+        if (~mask.any(dim=1)).any():
+            raise ValueError("attention mask contains a prompt with no valid tokens")
+
+        positions = torch.arange(sequence, device=logits.device).expand(batch, -1)
+        last_positions = positions.masked_fill(~mask, -1).max(dim=1).values
+        batch_positions = torch.arange(batch, device=logits.device)
+        last_logits = logits[batch_positions, last_positions]
+        if not cot_aware:
+            return last_logits
+
+        valid_counts = mask.sum(dim=1)
+        means = (logits * mask.unsqueeze(-1)).sum(dim=1) / valid_counts.unsqueeze(-1)
+        return torch.where((valid_counts > 4).unsqueeze(-1), means, last_logits)
+
     def run(self) -> Path:
         """Execute the full abliteration pipeline. Returns path to saved model."""
         # Remove any steering hooks left from a previous run() call
@@ -1736,13 +1842,18 @@ class AbliterationPipeline:
                     max_length=max_length,
                 )
                 inputs = {k: v.to(device) for k, v in inputs.items()}
-                with torch.no_grad():
-                    model(**inputs)
+                self._routing_attention_mask = inputs.get("attention_mask")
+                try:
+                    with torch.no_grad():
+                        model(**inputs)
+                finally:
+                    self._routing_attention_mask = None
                 del inputs
                 # Free GPU memory every few batches, not every prompt
                 if (batch_end % (batch_size * 4) == 0) or batch_end == len(prompts):
                     self._free_gpu_memory()
         finally:
+            self._routing_attention_mask = None
             tokenizer.padding_side = orig_padding_side
             for h in hooks:
                 h.remove()

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")