From cf2b5ae869206341cc8de0c1746ef28b97d566f4 Mon Sep 17 00:00:00 2001
From: Joseph Magly <1159087+jmagly@users.noreply.github.com>
Date: Fri, 21 Aug 2026 14:43:00 -0400
Subject: [PATCH] test: cover single-prompt flattened routing

---
 obliteratus/abliterate.py |  6 +++++-
 tests/test_abliterate.py  | 20 ++++++++++++++++++++
 2 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/obliteratus/abliterate.py b/obliteratus/abliterate.py
index 8caffb6..ca5322c 100644
--- a/obliteratus/abliterate.py
+++ b/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)
diff --git a/tests/test_abliterate.py b/tests/test_abliterate.py
index ef75506..7f802b5 100644
--- a/tests/test_abliterate.py
+++ b/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()