From a663500ea933357bb754ec06fa581f5626b6d594 Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Wed, 17 Jun 2026 16:42:02 -0700 Subject: [PATCH] [Test] Fold EAGLE `return_hidden_states` regression into spec triton suite (#28577) --- python/sglang/test/kits/spec_server_kits.py | 46 +++++++++++++ .../server_fixtures/spec_eagle_fixture.py | 5 ++ .../spec/eagle/test_eagle_hidden_states.py | 69 ------------------- .../spec/eagle/test_spec_eagle_triton.py | 14 +++- 4 files changed, 62 insertions(+), 72 deletions(-) delete mode 100644 test/registered/spec/eagle/test_eagle_hidden_states.py diff --git a/python/sglang/test/kits/spec_server_kits.py b/python/sglang/test/kits/spec_server_kits.py index c29639446..89997d178 100644 --- a/python/sglang/test/kits/spec_server_kits.py +++ b/python/sglang/test/kits/spec_server_kits.py @@ -574,3 +574,49 @@ class SpecFeatureKit: self.assertIsInstance(content, dict) except Exception: self.fail(f"parse JSON failed: {content_json}") + + +class SpecHiddenStatesKit: + """return_hidden_states under spec V2 (regression for issue #26163). + + Requires the server launched with --enable-return-hidden-states + (set ``enable_return_hidden_states = True`` on the fixture class). + """ + + def test_return_hidden_states(self): + # Two prompts of different lengths to exercise the per-req stride + # window: under spec V2 hidden_states is [bs * num_draft_tokens, dim], + # so a wrong index aliases a neighbor request's accepted rows. + prompts = [ + "Repeat: the quick brown fox the quick brown fox the quick brown fox", + "Count down from ten: ten nine eight", + ] + res = requests.post( + self.base_url + "/generate", + json={ + "text": prompts, + "sampling_params": {"temperature": 0, "max_new_tokens": 32}, + "return_hidden_states": True, + }, + ) + self.assertEqual(res.status_code, 200) + outputs = res.json() + + for out in outputs: + meta = out["meta_info"] + hs = meta["hidden_states"] + ct = meta["completion_tokens"] + # One hidden-state entry per completion token: hs[0] is the prefill + # block (List[List[float]]), hs[1:] are per-decode-token rows. + self.assertEqual( + len(hs), + ct, + f"len(hidden_states)={len(hs)} but completion_tokens={ct}", + ) + 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) diff --git a/python/sglang/test/server_fixtures/spec_eagle_fixture.py b/python/sglang/test/server_fixtures/spec_eagle_fixture.py index c18538f4e..ddb36c704 100644 --- a/python/sglang/test/server_fixtures/spec_eagle_fixture.py +++ b/python/sglang/test/server_fixtures/spec_eagle_fixture.py @@ -66,6 +66,9 @@ class SpecEagleServerBase(CustomTestCase): dtype = "bfloat16" cuda_graph_max_bs = None trust_remote_code = True + # Launch with --enable-return-hidden-states so SpecHiddenStatesKit can probe + # per-request hidden states; per-request gated, so other requests don't pay. + enable_return_hidden_states = False # -- extras -- # env_overrides: (env_var_obj, value) pairs applied only around launch. @@ -104,6 +107,8 @@ class SpecEagleServerBase(CustomTestCase): args.append("--disable-overlap-schedule") if cls.trust_remote_code: args.append("--trust-remote-code") + if cls.enable_return_hidden_states: + args.append("--enable-return-hidden-states") if cls.cuda_graph_max_bs is not None: args += ["--cuda-graph-max-bs", str(cls.cuda_graph_max_bs)] args += [str(a) for a in cls.extra_args] diff --git a/test/registered/spec/eagle/test_eagle_hidden_states.py b/test/registered/spec/eagle/test_eagle_hidden_states.py deleted file mode 100644 index 6b61ad813..000000000 --- a/test/registered/spec/eagle/test_eagle_hidden_states.py +++ /dev/null @@ -1,69 +0,0 @@ -"""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() diff --git a/test/registered/spec/eagle/test_spec_eagle_triton.py b/test/registered/spec/eagle/test_spec_eagle_triton.py index 63c5ce04b..dde5f04f0 100644 --- a/test/registered/spec/eagle/test_spec_eagle_triton.py +++ b/test/registered/spec/eagle/test_spec_eagle_triton.py @@ -1,4 +1,4 @@ -"""triton attention backend (EAGLE3 spec v2 + EAGLE/Llama-2 spec v1). +"""triton attention backend (EAGLE3 topk=1 chain + EAGLE/Llama-2 topk=8 tree). triton runs everywhere, so this stays on the cheap (5090) runner. """ @@ -11,6 +11,7 @@ from sglang.test.kits.matched_stop_kit import MatchedStopMixin from sglang.test.kits.spec_server_kits import ( SpecAccuracyKit, SpecFeatureKit, + SpecHiddenStatesKit, SpecLogprobKit, SpecPenaltyKit, ) @@ -37,10 +38,17 @@ class TestEagle3Triton( env_overrides = ((envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY, 1),) -class TestEagleLlama2Triton(EagleLlama2Base, SpecAccuracyKit, SpecFeatureKit): - """EAGLE/Llama-2 topk=8 on triton (spec v1).""" +class TestEagleLlama2Triton( + EagleLlama2Base, SpecAccuracyKit, SpecFeatureKit, SpecHiddenStatesKit +): + """EAGLE/Llama-2 topk=8 tree on triton. + + Hosts SpecHiddenStatesKit: topk>1 exercises the tree accept-path + compaction that the per-req hidden-state stride slicing depends on. + """ attention_backend = "triton" + enable_return_hidden_states = True env_overrides = ((envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY, 1),)