Re-org eagle unit tests (#14909)

docs/references/environment_variables.md

@@ -105,7 +105,7 @@ SGLang supports various environment variables that can be used to configure its
 | `SGLANG_TEST_RETRACT_NO_PREFILL_BS` | When SGLANG_TEST_RETRACT is enabled, no prefill is performed if the batch size exceeds SGLANG_TEST_RETRACT_NO_PREFILL_BS. | `2 ** 31`     |
 | `SGLANG_RECORD_STEP_TIME` | Record step time for profiling | `false` |
 | `SGLANG_TEST_REQUEST_TIME_STATS` | Test request time statistics | `false` |
-| `SGLANG_CI_SMALL_KV_SIZE` | Use small KV cache size in CI | Not set |
+| `SGLANG_CI_SMALL_KV_SIZE` | Use small KV cache size in CI | `-1` |
 
 ## Profiling & Benchmarking
 

python/sglang/srt/environ.py

@@ -173,6 +173,7 @@ class Envs:
     SGLANG_TEST_RETRACT_NO_PREFILL_BS = EnvInt(2 ** 31)
     SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY = EnvInt(0)
     SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_IDLE = EnvBool(True)
+    SGLANG_CI_SMALL_KV_SIZE = EnvInt(-1)
 
     # Scheduler: new token ratio hyperparameters
     SGLANG_INIT_NEW_TOKEN_RATIO = EnvFloat(0.7)

python/sglang/srt/model_executor/model_runner.py

@@ -237,9 +237,6 @@ def add_chunked_prefix_cache_attention_backend(backend_name):
         )
 
 
-# Use a small KV cache pool size for tests in CI
-SGLANG_CI_SMALL_KV_SIZE = os.getenv("SGLANG_CI_SMALL_KV_SIZE", None)
-
 # Detect stragger ranks in model loading
 UNBALANCED_MODEL_LOADING_TIMEOUT_S = 480  # leave more time for post data processing
 
@@ -1716,8 +1713,10 @@ class ModelRunner:
         log_info_on_rank0(logger, f"Using KV cache dtype: {self.kv_cache_dtype}")
 
         self.max_total_num_tokens = self.profile_max_num_token(total_gpu_memory)
-        if SGLANG_CI_SMALL_KV_SIZE:
+
-            self.max_total_num_tokens = int(SGLANG_CI_SMALL_KV_SIZE)
+        if (small_kv_size := envs.SGLANG_CI_SMALL_KV_SIZE.get()) > 0:
+            # Use a small KV cache pool size for local tests
+            self.max_total_num_tokens = small_kv_size
 
         if max_num_reqs is None:
             max_num_reqs = min(

python/sglang/test/server_fixtures/eagle_fixture.py

@@ -0,0 +1,115 @@
+import json
+import random
+import time
+
+import requests
+
+from sglang.srt.utils.common import kill_process_tree
+from sglang.test.test_utils import (
+    DEFAULT_EAGLE_DRAFT_MODEL_FOR_TEST,
+    DEFAULT_EAGLE_TARGET_MODEL_FOR_TEST,
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    popen_launch_server,
+)
+
+PROMPTS = [
+    "[INST] <<SYS>>\\nYou are a helpful assistant.\\n<</SYS>>\\nToday is a sunny day and I like[/INST]"
+    '[INST] <<SYS>>\\nYou are a helpful assistant.\\n<</SYS>>\\nWhat are the mental triggers in Jeff Walker\'s Product Launch Formula and "Launch" book?[/INST]',
+    "[INST] <<SYS>>\\nYou are a helpful assistant.\\n<</SYS>>\\nSummarize Russell Brunson's Perfect Webinar Script...[/INST]",
+    "[INST] <<SYS>>\\nYou are a helpful assistant.\\n<</SYS>>\\nwho are you?[/INST]",
+    "[INST] <<SYS>>\\nYou are a helpful assistant.\\n<</SYS>>\\nwhere are you from?[/INST]",
+]
+
+
+class EagleServerBase(CustomTestCase):
+    target_model = DEFAULT_EAGLE_TARGET_MODEL_FOR_TEST
+    draft_model = DEFAULT_EAGLE_DRAFT_MODEL_FOR_TEST
+    spec_algo = "EAGLE"
+    spec_steps = 5
+    spec_topk = 8
+    spec_tokens = 64
+    mem_fraction_static = 0.7
+    extra_args = []
+
+    @classmethod
+    def setUpClass(cls):
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        cls.process = popen_launch_server(
+            cls.target_model,
+            cls.base_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=[
+                f"--speculative-algorithm={cls.spec_algo}",
+                f"--speculative-draft-model-path={cls.draft_model}",
+                f"--speculative-num-steps={cls.spec_steps}",
+                f"--speculative-eagle-topk={cls.spec_topk}",
+                f"--speculative-num-draft-tokens={cls.spec_tokens}",
+                f"--mem-fraction-static={cls.mem_fraction_static}",
+            ]
+            + cls.extra_args,
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid)
+
+    def send_request(self):
+        time.sleep(random.uniform(0, 2))
+        for prompt in PROMPTS:
+            url = self.base_url + "/generate"
+            data = {
+                "text": prompt,
+                "sampling_params": {
+                    "temperature": 0,
+                    "max_new_tokens": 1024,
+                },
+            }
+            response = requests.post(url, json=data)
+            assert response.status_code == 200
+
+    def send_requests_abort(self):
+        for prompt in PROMPTS:
+            try:
+                time.sleep(random.uniform(0, 2))
+                url = self.base_url + "/generate"
+                data = {
+                    "model": "base",
+                    "text": prompt,
+                    "sampling_params": {
+                        "temperature": 0,
+                        "max_new_tokens": 1024,
+                    },
+                }
+                # set timeout = 1s, mock disconnected
+                requests.post(url, json=data, timeout=1)
+            except Exception as e:
+                print(e)
+                pass
+
+    def run_decode(self, sampling_params):
+        return_logprob = True
+        top_logprobs_num = 5
+        return_text = True
+        n = 1
+
+        response = requests.post(
+            self.base_url + "/generate",
+            json={
+                "text": "Human: Write a travel blog post to Hawaii.\n\nAssistant:",
+                "sampling_params": {
+                    "max_new_tokens": 48,
+                    "n": n,
+                    "temperature": 0.7,
+                    **sampling_params,
+                },
+                "return_logprob": return_logprob,
+                "top_logprobs_num": top_logprobs_num,
+                "return_text_in_logprobs": return_text,
+                "logprob_start_len": 0,
+            },
+        )
+        self.assertEqual(response.status_code, 200)
+        print(json.dumps(response.json()))
+        print("=" * 100)

test/registered/spec/eagle/test_eagle_infer_b.py

@@ -1,5 +1,4 @@
 import json
-import os
 import random
 import threading
 import time
@@ -11,95 +10,22 @@ from types import SimpleNamespace
 import numpy as np
 import requests
 
-from sglang.srt.utils import kill_process_tree
+from sglang.srt.environ import envs
 from sglang.test.ci.ci_register import register_cuda_ci
-from sglang.test.few_shot_gsm8k import run_eval
+from sglang.test.few_shot_gsm8k import run_eval as run_gsm8k_eval
+from sglang.test.server_fixtures.eagle_fixture import EagleServerBase
 from sglang.test.test_utils import (
-    DEFAULT_EAGLE_DRAFT_MODEL_FOR_TEST,
     DEFAULT_EAGLE_TARGET_MODEL_FOR_TEST,
-    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
-    DEFAULT_URL_FOR_TEST,
-    CustomTestCase,
-    popen_launch_server,
     run_logprob_check,
 )
 
 register_cuda_ci(est_time=473, suite="stage-b-test-small-1-gpu")
 
 
-class TestEAGLEServer(CustomTestCase):
+class TestEAGLEServerBasic(EagleServerBase):
-    PROMPTS = [
+    extra_args = ["--chunked-prefill-size", 128, "--max-running-requests", 8]
-        "[INST] <<SYS>>\\nYou are a helpful assistant.\\n<</SYS>>\\nToday is a sunny day and I like[/INST]"
-        '[INST] <<SYS>>\\nYou are a helpful assistant.\\n<</SYS>>\\nWhat are the mental triggers in Jeff Walker\'s Product Launch Formula and "Launch" book?[/INST]',
-        "[INST] <<SYS>>\\nYou are a helpful assistant.\\n<</SYS>>\\nSummarize Russell Brunson's Perfect Webinar Script...[/INST]",
-        "[INST] <<SYS>>\\nYou are a helpful assistant.\\n<</SYS>>\\nwho are you?[/INST]",
-        "[INST] <<SYS>>\\nYou are a helpful assistant.\\n<</SYS>>\\nwhere are you from?[/INST]",
-    ]
-
-    @classmethod
-    def setUpClass(cls):
-        cls.base_url = DEFAULT_URL_FOR_TEST
-        cls.process = popen_launch_server(
-            DEFAULT_EAGLE_TARGET_MODEL_FOR_TEST,
-            cls.base_url,
-            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
-            other_args=[
-                "--speculative-algorithm",
-                "EAGLE",
-                "--speculative-draft-model-path",
-                DEFAULT_EAGLE_DRAFT_MODEL_FOR_TEST,
-                "--speculative-num-steps",
-                5,
-                "--speculative-eagle-topk",
-                8,
-                "--speculative-num-draft-tokens",
-                64,
-                "--mem-fraction-static",
-                0.7,
-                "--chunked-prefill-size",
-                128,
-                "--max-running-requests",
-                8,
-            ],
-        )
-
-    @classmethod
-    def tearDownClass(cls):
-        kill_process_tree(cls.process.pid)
-
-    def send_request(self):
-        time.sleep(random.uniform(0, 2))
-        for prompt in self.PROMPTS:
-            url = self.base_url + "/generate"
-            data = {
-                "text": prompt,
-                "sampling_params": {
-                    "temperature": 0,
-                    "max_new_tokens": 1024,
-                },
-            }
-            response = requests.post(url, json=data)
-            assert response.status_code == 200
-
-    def send_requests_abort(self):
-        for prompt in self.PROMPTS:
-            try:
-                time.sleep(random.uniform(0, 2))
-                url = self.base_url + "/generate"
-                data = {
-                    "model": "base",
-                    "text": prompt,
-                    "sampling_params": {
-                        "temperature": 0,
-                        "max_new_tokens": 1024,
-                    },
-                }
-                # set timeout = 1s, mock disconnected
-                requests.post(url, json=data, timeout=1)
-            except Exception as e:
-                print(e)
-                pass
 
+    # FIXME(lsyin): move the test methods to kits
     def test_request_abort(self):
         concurrency = 4
         threads = [
@@ -127,7 +53,7 @@ class TestEAGLEServer(CustomTestCase):
         )
 
         # Just run and check it does not hang
-        metrics = run_eval(args)
+        metrics = run_gsm8k_eval(args)
         self.assertGreater(metrics["output_throughput"], 50)
 
     def test_gsm8k(self):
@@ -143,7 +69,7 @@ class TestEAGLEServer(CustomTestCase):
             port=int(self.base_url.split(":")[-1]),
         )
 
-        metrics = run_eval(args)
+        metrics = run_gsm8k_eval(args)
         print(f"{metrics=}")
         self.assertGreater(metrics["accuracy"], 0.20)
 
@@ -301,32 +227,6 @@ class TestEAGLEServer(CustomTestCase):
         with ThreadPoolExecutor(8) as executor:
             list(executor.map(func, args))
 
-    def run_decode(self, sampling_params):
-        return_logprob = True
-        top_logprobs_num = 5
-        return_text = True
-        n = 1
-
-        response = requests.post(
-            self.base_url + "/generate",
-            json={
-                "text": "Human: Write a travel blog post to Hawaii.\n\nAssistant:",
-                "sampling_params": {
-                    "max_new_tokens": 48,
-                    "n": n,
-                    "temperature": 0.7,
-                    **sampling_params,
-                },
-                "return_logprob": return_logprob,
-                "top_logprobs_num": top_logprobs_num,
-                "return_text_in_logprobs": return_text,
-                "logprob_start_len": 0,
-            },
-        )
-        self.assertEqual(response.status_code, 200)
-        print(json.dumps(response.json()))
-        print("=" * 100)
-
     def test_penalty_mixed(self):
         args = [
             {},
@@ -384,130 +284,41 @@ class TestEAGLEServer(CustomTestCase):
         self.assertTrue(is_valid_json)
 
 
-class TestEAGLERetract(TestEAGLEServer):
+class TestEAGLERetract(TestEAGLEServerBasic):
+    extra_args = ["--chunked-prefill-size", 128, "--max-running-requests", 64]
+
     @classmethod
     def setUpClass(cls):
         # These config helps find a leak.
-        os.environ["SGLANG_CI_SMALL_KV_SIZE"] = "4500"
+        # FIXME(lsyin): use override context manager
-        cls.base_url = DEFAULT_URL_FOR_TEST
+        envs.SGLANG_CI_SMALL_KV_SIZE.set(4500)
-        cls.process = popen_launch_server(
+        super().setUpClass()
-            DEFAULT_EAGLE_TARGET_MODEL_FOR_TEST,
-            cls.base_url,
-            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
-            other_args=[
-                "--speculative-algorithm",
-                "EAGLE",
-                "--speculative-draft-model-path",
-                DEFAULT_EAGLE_DRAFT_MODEL_FOR_TEST,
-                "--speculative-num-steps",
-                5,
-                "--speculative-eagle-topk",
-                8,
-                "--speculative-num-draft-tokens",
-                64,
-                "--mem-fraction-static",
-                0.7,
-                "--chunked-prefill-size",
-                128,
-                "--max-running-requests",
-                64,
-            ],
-        )
 
 
-class TestEAGLEServerTriton(TestEAGLEServer):
+class TestEAGLEServerTriton(TestEAGLEServerBasic):
-    @classmethod
+    extra_args = ["--attention-backend=triton", "--max-running-requests=8"]
-    def setUpClass(cls):
-        cls.base_url = DEFAULT_URL_FOR_TEST
-        cls.process = popen_launch_server(
-            DEFAULT_EAGLE_TARGET_MODEL_FOR_TEST,
-            cls.base_url,
-            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
-            other_args=[
-                "--speculative-algorithm",
-                "EAGLE",
-                "--speculative-draft-model-path",
-                DEFAULT_EAGLE_DRAFT_MODEL_FOR_TEST,
-                "--speculative-num-steps",
-                5,
-                "--speculative-eagle-topk",
-                8,
-                "--speculative-num-draft-tokens",
-                64,
-                "--mem-fraction-static",
-                0.7,
-                "--attention-backend",
-                "triton",
-                "--max-running-requests",
-                8,
-            ],
-        )
 
 
-class TestEAGLEServerPageSize(TestEAGLEServer):
+class TestEAGLEServerPageSize(TestEAGLEServerBasic):
-    @classmethod
+    spec_steps = 5
-    def setUpClass(cls):
+    spec_topk = 1
-        cls.base_url = DEFAULT_URL_FOR_TEST
+    spec_tokens = 6
-        cls.process = popen_launch_server(
+    extra_args = [
-            DEFAULT_EAGLE_TARGET_MODEL_FOR_TEST,
+        "--chunked-prefill-size=128",
-            cls.base_url,
+        "--max-running-requests=8",
-            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+        "--page-size=4",
-            other_args=[
+        "--attention-backend=flashinfer",
-                "--speculative-algorithm",
+    ]
-                "EAGLE",
-                "--speculative-draft-model-path",
-                DEFAULT_EAGLE_DRAFT_MODEL_FOR_TEST,
-                "--speculative-num-steps",
-                5,
-                "--speculative-eagle-topk",
-                1,
-                "--speculative-num-draft-tokens",
-                6,
-                "--mem-fraction-static",
-                0.7,
-                "--chunked-prefill-size",
-                128,
-                "--max-running-requests",
-                8,
-                "--page-size",
-                4,
-                "--attention-backend",
-                "flashinfer",
-            ],
-        )
 
 
-class TestEAGLEServerPageSizeTopk(TestEAGLEServer):
+class TestEAGLEServerPageSizeTopk(TestEAGLEServerBasic):
-    @classmethod
+    # default topk=8 and tokens=64
-    def setUpClass(cls):
+    extra_args = [
-        cls.base_url = DEFAULT_URL_FOR_TEST
+        "--chunked-prefill-size=128",
-        cls.process = popen_launch_server(
+        "--max-running-requests=8",
-            DEFAULT_EAGLE_TARGET_MODEL_FOR_TEST,
+        "--page-size=4",
-            cls.base_url,
+        "--attention-backend=flashinfer",
-            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    ]
-            other_args=[
-                "--speculative-algorithm",
-                "EAGLE",
-                "--speculative-draft-model-path",
-                DEFAULT_EAGLE_DRAFT_MODEL_FOR_TEST,
-                "--speculative-num-steps",
-                5,
-                "--speculative-eagle-topk",
-                8,
-                "--speculative-num-draft-tokens",
-                64,
-                "--mem-fraction-static",
-                0.7,
-                "--chunked-prefill-size",
-                128,
-                "--max-running-requests",
-                8,
-                "--page-size",
-                4,
-                "--attention-backend",
-                "flashinfer",
-            ],
-        )
 
 
 if __name__ == "__main__":