[Spec] Fix return_hidden_states under spec V2 (issue #26163) (#28496)

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

python/sglang/srt/managers/scheduler_components/batch_result_processor.py

@@ -668,9 +668,21 @@ class SchedulerBatchResultProcessor:
                 )
 
             if req.return_hidden_states and logits_output.hidden_states is not None:
-                req.hidden_states.append(
+                if batch.spec_algorithm.is_none():
-                    logits_output.hidden_states[i].cpu().clone().tolist()
+                    req.hidden_states.append(
-                )
+                        logits_output.hidden_states[i].cpu().clone().tolist()
+                    )
+                else:
+                    # Spec V2: hidden_states is [bs * speculative_num_draft_tokens, hidden_dim].
+                    stride = result.speculative_num_draft_tokens
+                    accept_len = result.num_correct_drafts_per_req_cpu[i] + 1
+                    start = i * stride
+                    req.hidden_states.extend(
+                        logits_output.hidden_states[start : start + accept_len]
+                        .cpu()
+                        .clone()
+                        .tolist()
+                    )
 
             if req.grammar is not None:
                 self._apply_decode_grammar(

python/sglang/srt/managers/scheduler_components/output_streamer.py

@@ -474,9 +474,14 @@ class _GenerationStreamAccumulator:
                 self.output_token_ids_logprobs_idx.append([])
 
         if self.return_hidden_states:
-            self.output_hidden_states.append(
+            if req.return_hidden_states:
-                req.hidden_states if req.return_hidden_states else None
+                # Mirror output_ids_through_stop: spec verify steps can overshoot finished_len.
-            )
+                hs = req.hidden_states
+                if req.finished_len is not None:
+                    hs = hs[: req.finished_len]
+                self.output_hidden_states.append(hs)
+            else:
+                self.output_hidden_states.append(None)
         if self.return_routed_experts:
             self.routed_experts.append(
                 req.routed_experts if req.return_routed_experts else None

test/registered/spec/eagle/test_eagle_hidden_states.py

@@ -0,0 +1,69 @@
+"""Regression test for issue #26163: return_hidden_states under EAGLE spec V2."""
+
+import unittest
+
+import sglang as sgl
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.test_utils import (
+    DEFAULT_DRAFT_MODEL_EAGLE,
+    DEFAULT_TARGET_MODEL_EAGLE,
+    CustomTestCase,
+)
+
+register_cuda_ci(est_time=120, stage="base-b", runner_config="1-gpu-large")
+
+
+class TestEagleReturnHiddenStates(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.engine = sgl.Engine(
+            model_path=DEFAULT_TARGET_MODEL_EAGLE,
+            speculative_algorithm="EAGLE",
+            speculative_draft_model_path=DEFAULT_DRAFT_MODEL_EAGLE,
+            speculative_num_steps=3,
+            speculative_eagle_topk=4,
+            speculative_num_draft_tokens=8,
+            enable_return_hidden_states=True,
+            mem_fraction_static=0.7,
+            attention_backend="triton",
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        if hasattr(cls, "engine") and cls.engine is not None:
+            cls.engine.shutdown()
+
+    def test_hidden_states_length_matches_completion(self):
+        # Two prompts to exercise cross-request stride aliasing.
+        prompts = [
+            "Repeat: the quick brown fox the quick brown fox the quick brown fox",
+            "Count down from ten: ten nine eight",
+        ]
+        max_new_tokens = 32
+        outputs = self.engine.generate(
+            prompts,
+            sampling_params={"temperature": 0, "max_new_tokens": max_new_tokens},
+            return_hidden_states=True,
+        )
+
+        for out in outputs:
+            meta = out["meta_info"]
+            hs = meta["hidden_states"]
+            ct = meta["completion_tokens"]
+            self.assertEqual(
+                len(hs),
+                ct,
+                f"len(hidden_states)={len(hs)} but completion_tokens={ct}",
+            )
+            # hs[0] is the prefill block (List[List[float]]); hs[1:] are decode rows.
+            decode_rows = hs[1:]
+            self.assertGreater(len(decode_rows), 0)
+            hidden_dim = len(decode_rows[0])
+            self.assertGreater(hidden_dim, 0)
+            for row in decode_rows:
+                self.assertIsInstance(row, list)
+                self.assertEqual(len(row), hidden_dim)
+
+
+if __name__ == "__main__":
+    unittest.main()