diff --git a/docs/references/environment_variables.md b/docs/references/environment_variables.md
index 8d450a48f..e45a36d6e 100644
--- a/docs/references/environment_variables.md
+++ b/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
diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py
index 2193ddf67..89b4ab36e 100644
--- a/python/sglang/srt/environ.py
+++ b/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)
diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py
index d65e5ef42..a3fc18e57 100644
--- a/python/sglang/srt/model_executor/model_runner.py
+++ b/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(
diff --git a/python/sglang/test/server_fixtures/eagle_fixture.py b/python/sglang/test/server_fixtures/eagle_fixture.py
new file mode 100644
index 000000000..85e77fc69
--- /dev/null
+++ b/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] <>\\nYou are a helpful assistant.\\n<>\\nToday is a sunny day and I like[/INST]"
+ '[INST] <>\\nYou are a helpful assistant.\\n<>\\nWhat are the mental triggers in Jeff Walker\'s Product Launch Formula and "Launch" book?[/INST]',
+ "[INST] <>\\nYou are a helpful assistant.\\n<>\\nSummarize Russell Brunson's Perfect Webinar Script...[/INST]",
+ "[INST] <>\\nYou are a helpful assistant.\\n<>\\nwho are you?[/INST]",
+ "[INST] <>\\nYou are a helpful assistant.\\n<>\\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)
diff --git a/test/registered/spec/eagle/test_eagle_infer_b.py b/test/registered/spec/eagle/test_eagle_infer_b.py
index a4909e3ef..4bff4953c 100644
--- a/test/registered/spec/eagle/test_eagle_infer_b.py
+++ b/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):
- PROMPTS = [
- "[INST] <>\\nYou are a helpful assistant.\\n<>\\nToday is a sunny day and I like[/INST]"
- '[INST] <>\\nYou are a helpful assistant.\\n<>\\nWhat are the mental triggers in Jeff Walker\'s Product Launch Formula and "Launch" book?[/INST]',
- "[INST] <>\\nYou are a helpful assistant.\\n<>\\nSummarize Russell Brunson's Perfect Webinar Script...[/INST]",
- "[INST] <>\\nYou are a helpful assistant.\\n<>\\nwho are you?[/INST]",
- "[INST] <>\\nYou are a helpful assistant.\\n<>\\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
+class TestEAGLEServerBasic(EagleServerBase):
+ extra_args = ["--chunked-prefill-size", 128, "--max-running-requests", 8]
+ # 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"
- 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",
- 64,
- ],
- )
+ # FIXME(lsyin): use override context manager
+ envs.SGLANG_CI_SMALL_KV_SIZE.set(4500)
+ super().setUpClass()
-class TestEAGLEServerTriton(TestEAGLEServer):
- @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,
- "--attention-backend",
- "triton",
- "--max-running-requests",
- 8,
- ],
- )
+class TestEAGLEServerTriton(TestEAGLEServerBasic):
+ extra_args = ["--attention-backend=triton", "--max-running-requests=8"]
-class TestEAGLEServerPageSize(TestEAGLEServer):
- @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",
- 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 TestEAGLEServerPageSize(TestEAGLEServerBasic):
+ spec_steps = 5
+ spec_topk = 1
+ spec_tokens = 6
+ extra_args = [
+ "--chunked-prefill-size=128",
+ "--max-running-requests=8",
+ "--page-size=4",
+ "--attention-backend=flashinfer",
+ ]
-class TestEAGLEServerPageSizeTopk(TestEAGLEServer):
- @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,
- "--page-size",
- 4,
- "--attention-backend",
- "flashinfer",
- ],
- )
+class TestEAGLEServerPageSizeTopk(TestEAGLEServerBasic):
+ # default topk=8 and tokens=64
+ extra_args = [
+ "--chunked-prefill-size=128",
+ "--max-running-requests=8",
+ "--page-size=4",
+ "--attention-backend=flashinfer",
+ ]
if __name__ == "__main__":