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

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)

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]

test/registered/spec/eagle/test_eagle_hidden_states.py

@@ -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

@@ -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):
+class TestEagleLlama2Triton(
-    """EAGLE/Llama-2 topk=8 on triton (spec v1)."""
+    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),)