diff --git a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py b/python/sglang/srt/managers/scheduler_components/batch_result_processor.py index 503d31a24..7ac74073f 100644 --- a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py +++ b/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( - logits_output.hidden_states[i].cpu().clone().tolist() - ) + if batch.spec_algorithm.is_none(): + 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( diff --git a/python/sglang/srt/managers/scheduler_components/output_streamer.py b/python/sglang/srt/managers/scheduler_components/output_streamer.py index 8d26f7934..f95c59f6f 100644 --- a/python/sglang/srt/managers/scheduler_components/output_streamer.py +++ b/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( - req.hidden_states if req.return_hidden_states else None - ) + if req.return_hidden_states: + # 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 diff --git a/test/registered/spec/eagle/test_eagle_hidden_states.py b/test/registered/spec/eagle/test_eagle_hidden_states.py new file mode 100644 index 000000000..6b61ad813 --- /dev/null +++ b/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()