chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Re-org eagle unit tests (#14909)

Browse Source
This commit is contained in:
Liangsheng Yin
2025-12-12 12:25:39 +09:00
committed by GitHub
parent 10146af099
commit c660d8dfd0
5 changed files with 155 additions and 229 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
docs/references/environment_variables.md
View File

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

1
python/sglang/srt/environ.py
View File

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

9
python/sglang/srt/model_executor/model_runner.py
View File

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

115
python/sglang/test/server_fixtures/eagle_fixture.py Normal file
View File

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

257
test/registered/spec/eagle/test_eagle_infer_b.py
View File

@@ -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] <<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
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__":