test: cover single-prompt flattened routing

obliteratus/abliterate.py

@@ -1842,7 +1842,11 @@ class AbliterationPipeline:
                     max_length=max_length,
                 )
                 inputs = {k: v.to(device) for k, v in inputs.items()}
-                self._routing_attention_mask = inputs.get("attention_mask")
+                routing_mask = inputs.get("attention_mask")
+                input_ids = inputs.get("input_ids")
+                if routing_mask is None and isinstance(input_ids, torch.Tensor):
+                    routing_mask = torch.ones_like(input_ids, dtype=torch.bool)
+                self._routing_attention_mask = routing_mask
                 try:
                     with torch.no_grad():
                         model(**inputs)

tests/test_abliterate.py

@@ -2273,6 +2273,26 @@ class TestRouterProfilingHooks:
                 h.remove()
             abl_module.get_ffn_module = orig_get_ffn
 
+    def test_hooks_record_single_prompt_flattened_tokens(self):
+        """A flattened single-prompt sequence should select its final token."""
+        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(1, 4, dtype=torch.long)
+            x = torch.arange(4 * 16, dtype=torch.float32).reshape(4, 16) / 100
+            expected = torch.nn.functional.linear(x, layer.mlp.gate.weight)[-1]
+
+            layer.mlp.gate(x)
+
+            assert len(pipeline._routing_harmful[0]) == 1
+            assert torch.equal(pipeline._routing_harmful[0][0], expected)
+        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()