test: cover single-prompt flattened routing

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