minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[Test] Fold EAGLE return_hidden_states regression into spec triton suite (#28577)

Browse Source
This commit is contained in:
Liangsheng Yin
2026-06-17 16:42:02 -07:00
committed by GitHub
parent b88bada64e
commit a663500ea9
4 changed files with 62 additions and 72 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/test/kits/spec_server_kits.py
+46
View File
@@ -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)
python/sglang/test/server_fixtures/spec_eagle_fixture.py
+5
View File
@@ -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]
test/registered/spec/eagle/test_eagle_hidden_states.py
-69
View File
@@ -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()
test/registered/spec/eagle/test_spec_eagle_triton.py
+11 -3
View File
@@ -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),)