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

[CI] Migrate Eagle 1-GPU tests to test/registered/ (#14529)

Browse Source
This commit is contained in:
Alison Shao
2025-12-09 12:56:36 +09:00
committed by GitHub
parent af20657cd4
commit e6f0ddda44
9 changed files with 58 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/srt/run_suite.py
-5
View File
@@ -7,7 +7,6 @@ from sglang.test.ci.ci_utils import TestFile, run_unittest_files
# NOTE: please sort the test cases alphabetically by the test file name
suites = {
"per-commit-1-gpu": [
TestFile("test_eagle_constrained_decoding.py", 100),
TestFile("debug_utils/test_tensor_dump_forward_hook.py", 9),
TestFile("hicache/test_hicache_storage.py", 96),
TestFile("hicache/test_hicache_variants.py", 368),
@@ -58,13 +57,9 @@ suites = {
# TestFile("rl/test_update_weights_from_disk.py", 210), # Temporarily disabled, see https://github.com/sgl-project/sglang/pull/13998
TestFile("rl/test_update_weights_from_tensor.py", 195),
TestFile("test_abort.py", 131),
TestFile("test_build_eagle_tree.py", 3),
TestFile("test_chunked_prefill.py", 312),
TestFile("test_create_kvindices.py", 7),
TestFile("test_deterministic.py", 228),
TestFile("test_eagle_infer_a.py", 470),
TestFile("test_eagle_infer_b.py", 473),
TestFile("test_eagle_infer_beta.py", 194),
TestFile("test_constrained_decoding.py", 111),
TestFile("test_eval_fp8_accuracy.py", 250),
TestFile("test_external_models.py", 30),
test/srt/test_build_eagle_tree.py
-308
View File
@@ -1,308 +0,0 @@
import unittest
import torch
from sglang.srt.speculative.eagle_utils import (
build_tree_kernel_efficient,
organize_draft_results,
)
class TestBuildEagleTree(unittest.TestCase):
"""Unit tests for build_eagle_tree functionality."""
def test_build_tree_kernel_efficient(self):
"""Test the build_tree_kernel_efficient function with known inputs and expected outputs."""
verified_id = torch.tensor([29974, 13], device="cuda", dtype=torch.int32)
score_list = [
torch.tensor(
[
[[7.1127e-01, 2.8292e-01, 2.2995e-03, 1.7357e-03]],
[[9.7476e-01, 2.2219e-02, 6.5031e-04, 1.3212e-04]],
],
dtype=torch.float32,
device="cuda",
),
torch.tensor(
[
[
[6.9142e-01, 1.2863e-02, 1.6873e-03, 1.1871e-03],
[2.4787e-01, 1.8818e-02, 1.4204e-02, 9.2235e-04],
[2.2971e-03, 1.6700e-06, 1.8737e-07, 8.3146e-08],
[1.2771e-03, 2.4374e-04, 1.7832e-04, 1.1947e-05],
],
[
[8.4832e-02, 6.6068e-02, 5.8304e-02, 5.7851e-02],
[2.3616e-03, 1.1243e-03, 5.4368e-04, 2.7768e-04],
[2.5286e-04, 1.5578e-04, 2.8817e-05, 1.2888e-05],
[1.2834e-04, 2.5417e-06, 1.1279e-06, 1.6088e-08],
],
],
dtype=torch.float32,
device="cuda",
),
torch.tensor(
[
[
[6.6438e-01, 2.6997e-02, 2.4236e-05, 4.0821e-06],
[2.4402e-01, 2.8409e-03, 5.0935e-04, 2.9022e-04],
[1.6178e-02, 2.0567e-03, 4.5892e-04, 3.0034e-05],
[1.3023e-02, 5.0497e-04, 3.6371e-04, 8.7750e-05],
],
[
[2.3263e-02, 2.0054e-02, 9.3990e-03, 2.7783e-03],
[6.4156e-02, 5.5506e-04, 1.0429e-04, 9.7211e-05],
[4.9950e-02, 5.0630e-03, 9.0068e-04, 3.3656e-04],
[7.5817e-03, 8.5731e-04, 6.9972e-04, 6.0793e-04],
],
],
dtype=torch.float32,
device="cuda",
),
torch.tensor(
[
[
[6.6420e-01, 1.0525e-04, 6.5864e-05, 1.2253e-06],
[1.3019e-01, 1.0461e-01, 5.2083e-03, 1.6777e-03],
[2.0103e-02, 6.7335e-03, 1.2625e-04, 1.0364e-05],
[1.5142e-02, 7.0819e-04, 9.6595e-05, 8.7951e-05],
],
[
[5.8608e-02, 1.8840e-03, 7.8535e-04, 4.4400e-04],
[1.2185e-02, 2.0684e-03, 1.7418e-03, 1.4327e-03],
[6.2455e-03, 6.1487e-03, 2.6862e-03, 1.8034e-03],
[1.8590e-03, 1.6151e-03, 1.2481e-03, 3.6038e-04],
],
],
dtype=torch.float32,
device="cuda",
),
]
token_list = [
torch.tensor(
[[29896, 29906, 29900, 29945], [13, 2, 29871, 28956]],
dtype=torch.int64,
device="cuda",
),
torch.tensor(
[
[
29889,
29974,
29945,
29900,
29974,
29922,
29930,
29958,
29889,
29974,
29930,
29945,
29974,
29922,
29930,
29958,
],
[
22550,
4136,
16492,
8439,
29871,
2,
3001,
13,
2,
13,
29906,
29946,
2,
13,
29871,
259,
],
],
device="cuda",
),
torch.tensor(
[
[
29946,
29945,
29953,
29906,
29896,
29945,
29900,
29906,
29896,
29945,
29906,
29953,
29896,
29945,
29906,
29946,
],
[
29871,
2,
29901,
29889,
29871,
2,
395,
259,
29901,
29871,
2,
29889,
3001,
1234,
7146,
2186,
],
],
device="cuda",
),
torch.tensor(
[
[
29946,
29974,
29945,
29930,
29889,
29922,
29974,
29930,
29974,
29946,
29930,
29922,
29889,
29974,
29945,
29922,
],
[
29941,
29906,
2,
29946,
29871,
450,
319,
14990,
29946,
29941,
2,
29906,
29871,
2,
3001,
13,
],
],
device="cuda",
),
]
parents_list = [
torch.tensor(
[[-1, 0, 1, 2, 3], [-1, 0, 1, 2, 3]], dtype=torch.int64, device="cuda"
),
torch.tensor(
[[4, 8, 9, 10], [4, 5, 6, 7]], dtype=torch.int64, device="cuda"
),
torch.tensor(
[[20, 24, 21, 28], [24, 28, 20, 21]], dtype=torch.int64, device="cuda"
),
torch.tensor(
[[36, 40, 41, 44], [36, 40, 44, 45]], dtype=torch.int64, device="cuda"
),
]
seq_lens = torch.tensor([5, 10], dtype=torch.int64, device="cuda")
topk = 4
depth = 4
num_draft_token = 8
parent_list, top_scores_index, draft_tokens = organize_draft_results(
score_list, token_list, parents_list, num_draft_token
)
(
tree_mask,
position,
retrieve_index,
retrieve_next_token,
retrieve_next_sibling,
draft_tokens,
) = build_tree_kernel_efficient(
verified_id=verified_id,
parent_list=parent_list,
top_scores_index=top_scores_index,
draft_tokens=draft_tokens,
seq_lens=seq_lens,
seq_lens_sum=torch.sum(seq_lens).item(),
topk=topk,
spec_steps=depth,
num_verify_tokens=num_draft_token,
)
# Verify expected outputs
self.assertEqual(
position.tolist(),
[5, 6, 6, 7, 7, 8, 8, 9, 10, 11, 12, 12, 12, 12, 13, 14],
"Position tensor does not match expected values",
)
self.assertEqual(
retrieve_index.tolist(),
[
[0, 1, 2, 3, 4, 5, 6, 7],
[8, 9, 10, 11, 12, 13, 14, 15],
],
"Retrieve index tensor does not match expected values",
)
self.assertEqual(
retrieve_next_token.tolist(),
[
[1, 3, 4, 5, 6, 7, -1, -1],
[1, 2, -1, 6, -1, -1, 7, -1],
],
"Retrieve next token tensor does not match expected values",
)
self.assertEqual(
retrieve_next_sibling.tolist(),
[
[-1, 2, -1, -1, -1, -1, -1, -1],
[-1, -1, 3, 4, 5, -1, -1, -1],
],
"Retrieve next sibling tensor does not match expected values",
)
self.assertEqual(
draft_tokens.tolist(),
[
29974,
29896,
29906,
29889,
29974,
29946,
29896,
29946,
13,
13,
22550,
4136,
16492,
8439,
29871,
29941,
],
"Draft tokens tensor does not match expected values",
)
if __name__ == "__main__":
unittest.main()
test/srt/test_eagle_constrained_decoding.py
-77
View File
@@ -1,77 +0,0 @@
import unittest
from sglang.srt.environ import envs
from sglang.srt.utils import kill_process_tree
from sglang.test.kits.json_constrained_kit import TestJSONConstrainedMixin
from sglang.test.kits.regex_constrained_kit import TestRegexConstrainedMixin
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,
)
class TestEagleConstrainedDecoding(
CustomTestCase, TestRegexConstrainedMixin, TestJSONConstrainedMixin
):
max_running_requests = 64
attention_backend = "triton"
spec_steps = 5
spec_topk = 1
spec_draft_tokens = 6
page_size = 1
other_launch_args = []
model = DEFAULT_EAGLE_TARGET_MODEL_FOR_TEST
draft_model = DEFAULT_EAGLE_DRAFT_MODEL_FOR_TEST
grammar_backend = "xgrammar"
eagle_v2 = False
@classmethod
def setUpClass(cls):
cls.base_url = DEFAULT_URL_FOR_TEST
launch_args = [
"--trust-remote-code",
"--attention-backend",
cls.attention_backend,
"--speculative-algorithm",
"EAGLE",
"--speculative-draft-model",
cls.draft_model,
"--speculative-num-steps",
cls.spec_steps,
"--speculative-eagle-topk",
cls.spec_topk,
"--speculative-num-draft-tokens",
cls.spec_draft_tokens,
"--page-size",
str(cls.page_size),
"--mem-fraction-static",
"0.75",
"--max-running-requests",
str(cls.max_running_requests),
"--grammar-backend",
cls.grammar_backend,
]
launch_args.extend(cls.other_launch_args)
with envs.SGLANG_ENABLE_SPEC_V2.override(cls.eagle_v2):
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=launch_args,
)
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
class TestEagleConstrainedDecodingV2(TestEagleConstrainedDecoding):
eagle_v2 = True
if __name__ == "__main__":
unittest.main()
test/srt/test_eagle_infer_a.py
-438
View File
@@ -1,438 +0,0 @@
import os
import unittest
import requests
import torch
import sglang as sgl
from sglang.srt.utils import kill_process_tree
from sglang.srt.utils.hf_transformers_utils import get_tokenizer
from sglang.test.test_utils import (
DEFAULT_EAGLE_DRAFT_MODEL_FOR_TEST,
DEFAULT_EAGLE_TARGET_MODEL_FOR_TEST,
DEFAULT_EAGLE_TARGET_MODEL_FOR_TEST_EAGLE3,
DEFAULT_MODEL_NAME_FOR_TEST_EAGLE3,
DEFAULT_MODEL_NAME_FOR_TEST_MLA,
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
CustomTestCase,
is_in_ci,
popen_launch_server,
)
torch_dtype = torch.float16
prefill_tolerance = 5e-2
decode_tolerance: float = 5e-2
class TestEAGLEEngine(CustomTestCase):
BASE_CONFIG = {
"model_path": DEFAULT_EAGLE_TARGET_MODEL_FOR_TEST,
"speculative_draft_model_path": DEFAULT_EAGLE_DRAFT_MODEL_FOR_TEST,
"speculative_algorithm": "EAGLE",
"speculative_num_steps": 5,
"speculative_eagle_topk": 4,
"speculative_num_draft_tokens": 8,
"mem_fraction_static": 0.7,
"cuda_graph_max_bs": 5,
"trust_remote_code": True,
}
NUM_CONFIGS = 2
THRESHOLDS = {
"batch_avg_accept_len": 1.9,
"accept_len": 3.6,
}
def setUp(self):
self.prompt = "Today is a sunny day and I like"
self.sampling_params = {"temperature": 0, "max_new_tokens": 8}
ref_engine = sgl.Engine(
model_path=self.BASE_CONFIG["model_path"], cuda_graph_max_bs=1
)
self.ref_output = ref_engine.generate(self.prompt, self.sampling_params)["text"]
ref_engine.shutdown()
def test_correctness(self):
configs = [
# Basic config
self.BASE_CONFIG,
# Chunked prefill
{**self.BASE_CONFIG, "chunked_prefill_size": 4},
]
for i, config in enumerate(configs[: self.NUM_CONFIGS]):
with self.subTest(i=i):
print(f"{config=}")
engine = sgl.Engine(**config, log_level="info", decode_log_interval=10)
try:
self._test_single_generation(engine)
self._test_batch_generation(engine)
self._test_eos_token(engine)
self._test_acc_length(engine)
finally:
engine.flush_cache() # check engine alive
engine.shutdown()
print("=" * 100)
def _test_single_generation(self, engine):
output = engine.generate(self.prompt, self.sampling_params)["text"]
print(f"{output=}, {self.ref_output=}")
self.assertEqual(output, self.ref_output)
def _test_batch_generation(self, engine):
prompts = [
"Hello, my name is",
"The president of the United States is",
"The capital of France is",
"The future of AI is",
]
params = {"temperature": 0, "max_new_tokens": 50}
outputs = engine.generate(prompts, params)
for prompt, output in zip(prompts, outputs):
print(f"Prompt: {prompt}")
print(f"Generated: {output['text']}")
print("-" * 40)
print(f"{engine.get_server_info()=}")
avg_spec_accept_length = engine.get_server_info()["internal_states"][0][
"avg_spec_accept_length"
]
print(f"{avg_spec_accept_length=}")
self.assertGreater(
avg_spec_accept_length, self.THRESHOLDS["batch_avg_accept_len"]
)
def _test_eos_token(self, engine):
prompt = "[INST] <<SYS>>\nYou are a helpful assistant.\n<</SYS>>\nToday is a sunny day and I like [/INST]"
params = {
"temperature": 0.1,
"max_new_tokens": 1024,
"skip_special_tokens": False,
}
tokenizer = get_tokenizer(DEFAULT_EAGLE_TARGET_MODEL_FOR_TEST)
output = engine.generate(prompt, params)["text"]
print(f"{output=}")
tokens = tokenizer.encode(output, truncation=False)
self.assertNotIn(tokenizer.eos_token_id, tokens)
def _test_acc_length(self, engine):
prompt = [
"Human: Give me a fully functional FastAPI server. Show the python code.\n\nAssistant:",
] * 5 # test batched generation
sampling_params = {"temperature": 0, "max_new_tokens": 512}
output = engine.generate(prompt, sampling_params)
output = output[0]
if "spec_verify_ct" in output["meta_info"]:
acc_length = (
output["meta_info"]["completion_tokens"]
/ output["meta_info"]["spec_verify_ct"]
)
else:
acc_length = 1.0
speed = (
output["meta_info"]["completion_tokens"]
/ output["meta_info"]["e2e_latency"]
)
print(f"{acc_length=:.4f}, {speed=}")
self.assertGreater(acc_length, self.THRESHOLDS["accept_len"])
class TestEAGLEEngineTokenMap(TestEAGLEEngine):
BASE_CONFIG = {
"model_path": "meta-llama/Meta-Llama-3-8B-Instruct",
"speculative_draft_model_path": "lmsys/sglang-EAGLE-LLaMA3-Instruct-8B",
"speculative_algorithm": "EAGLE",
"speculative_num_steps": 5,
"speculative_eagle_topk": 4,
"speculative_num_draft_tokens": 8,
"speculative_token_map": "thunlp/LLaMA3-Instruct-8B-FR-Spec/freq_32768.pt",
"mem_fraction_static": 0.7,
"cuda_graph_max_bs": 5,
"dtype": "float16",
}
NUM_CONFIGS = 1
THRESHOLDS = {
"batch_avg_accept_len": 1.9,
"accept_len": 2.5,
}
class TestEAGLE3Engine(TestEAGLEEngine):
BASE_CONFIG = {
"model_path": DEFAULT_EAGLE_TARGET_MODEL_FOR_TEST_EAGLE3,
"speculative_draft_model_path": DEFAULT_MODEL_NAME_FOR_TEST_EAGLE3,
"speculative_algorithm": "EAGLE3",
"speculative_num_steps": 5,
"speculative_eagle_topk": 16,
"speculative_num_draft_tokens": 64,
"mem_fraction_static": 0.7,
"cuda_graph_max_bs": 5,
"dtype": "float16",
}
NUM_CONFIGS = 1
THRESHOLDS = {
"batch_avg_accept_len": 1.75,
"accept_len": 3.1,
}
class TestEAGLERadixCache(CustomTestCase):
BASE_CONFIG = {
"model_path": DEFAULT_EAGLE_TARGET_MODEL_FOR_TEST_EAGLE3,
"speculative_draft_model_path": DEFAULT_MODEL_NAME_FOR_TEST_EAGLE3,
"speculative_algorithm": "EAGLE3",
"speculative_num_steps": 2,
"speculative_eagle_topk": 2,
"speculative_num_draft_tokens": 5,
"mem_fraction_static": 0.7,
"dtype": "float16",
"trust_remote_code": True,
"attention_backend": "fa3",
"skip_server_warmup": True,
"cuda_graph_max_bs": 5,
}
def test_correctness(self):
os.environ["SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN"] = "1"
configs = [
# Basic config
self.BASE_CONFIG,
# Chunked prefill & Page Size > 1
{**self.BASE_CONFIG, "chunked_prefill_size": 64, "page_size": 4},
{**self.BASE_CONFIG, "page_size": 4},
# Preferred by some kernels
{**self.BASE_CONFIG, "page_size": 64},
# Disable CUDA Graph
{
**self.BASE_CONFIG,
"disable_cuda_graph": True,
"page_size": 4,
},
]
for i, config in enumerate(configs):
with self.subTest(i=i):
print(f"{config=}")
engine = sgl.Engine(**config, log_level="info", decode_log_interval=10)
try:
self._test_acc_length(engine)
self._test_batch_generation(engine)
finally:
engine.shutdown()
print("=" * 100)
del os.environ["SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN"]
def _test_acc_length(self, engine):
warmup_prompt = [
"Human: Give me a fully functional FastAPI server. Show the python code.\n\nAssistant:",
]
sampling_params = {"temperature": 0, "max_new_tokens": 512}
output = engine.generate(warmup_prompt, sampling_params)
test_prompt = [
"<|start_header_id|>system<|end_header_id|>\n\nYou are a helpful, respectful and honest assistant. Always answer as helpfully as possible, while being safe. Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure that your responses are socially unbiased and positive in nature.\n\nIf a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information.<|eot_id|><|start_header_id|>user<|end_header_id|>\n\nGive me a fully functional FastAPI server. Show the python code.<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
]
output = engine.generate(test_prompt, sampling_params)
output = output[0]
if "spec_verify_ct" in output["meta_info"]:
acc_length = (
output["meta_info"]["completion_tokens"]
/ output["meta_info"]["spec_verify_ct"]
)
else:
acc_length = 1.0
speed = (
output["meta_info"]["completion_tokens"]
/ output["meta_info"]["e2e_latency"]
)
print(f"{acc_length=:.4f}, {speed=}")
self.assertGreater(acc_length, 2.5)
def _test_batch_generation(self, engine):
prompts = [
"Hello, my name is",
"The president of the United States is",
"The capital of France is",
"The future of AI is",
]
params = {"temperature": 0, "max_new_tokens": 50}
outputs = engine.generate(prompts, params)
for prompt, output in zip(prompts, outputs):
print(f"Prompt: {prompt}")
print(f"Generated: {output['text']}")
print("-" * 40)
print(f"{engine.get_server_info()=}")
avg_spec_accept_length = engine.get_server_info()["internal_states"][0][
"avg_spec_accept_length"
]
print(f"{avg_spec_accept_length=}")
self.assertGreater(avg_spec_accept_length, 2.0)
@unittest.skipIf(is_in_ci(), "To reduce the CI execution time.")
class TestEAGLEDraftExtend(CustomTestCase):
@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",
1,
"--speculative-eagle-topk",
1,
"--speculative-num-draft-tokens",
2,
"--max-running-requests",
4,
"--attention-backend",
"fa3",
],
)
cls.accept_len_threshold = 1.50
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
def test_one_batch_accept_length(self):
resp = requests.get(self.base_url + "/flush_cache")
self.assertEqual(resp.status_code, 200)
prompts = [
"Hello, my name is",
"The president of the United States is",
"The capital of France is",
"The future of AI is",
]
url = self.base_url + "/generate"
data = {
"text": prompts,
"sampling_params": {
"temperature": 0,
"max_new_tokens": 512,
},
}
response = requests.post(url, json=data)
self.assertEqual(response.status_code, 200)
outputs = response.json()
for i in range(len(prompts)):
output = outputs[i]
if "spec_verify_ct" in output["meta_info"]:
acc_length = (
output["meta_info"]["completion_tokens"]
/ output["meta_info"]["spec_verify_ct"]
)
else:
acc_length = 1.0
print(f"{acc_length=}")
self.assertGreater(acc_length, self.accept_len_threshold)
class TestEAGLEDraftExtendFlashinfer(TestEAGLEDraftExtend):
@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",
1,
"--speculative-eagle-topk",
1,
"--speculative-num-draft-tokens",
2,
"--max-running-requests",
4,
"--attention-backend",
"flashinfer",
],
)
cls.accept_len_threshold = 1.50
@unittest.skipIf(is_in_ci(), "To reduce the CI execution time.")
class TestEAGLEDraftExtendTriton(TestEAGLEDraftExtend):
@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",
1,
"--speculative-eagle-topk",
1,
"--speculative-num-draft-tokens",
2,
"--max-running-requests",
4,
"--attention-backend",
"triton",
],
)
cls.accept_len_threshold = 1.50
@unittest.skipIf(is_in_ci(), "To reduce the CI execution time.")
class TestEAGLEDraftExtendFlashinferMLA(TestEAGLEDraftExtend):
@classmethod
def setUpClass(cls):
cls.base_url = DEFAULT_URL_FOR_TEST
cls.process = popen_launch_server(
DEFAULT_MODEL_NAME_FOR_TEST_MLA,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=[
"--speculative-algorithm",
"EAGLE",
"--speculative-num-steps",
1,
"--speculative-eagle-topk",
1,
"--speculative-num-draft-tokens",
2,
"--max-running-requests",
4,
"--attention-backend",
"flashinfer",
],
)
cls.accept_len_threshold = 1.85
if __name__ == "__main__":
unittest.main()
test/srt/test_eagle_infer_b.py
-511
View File
@@ -1,511 +0,0 @@
import json
import os
import random
import threading
import time
import unittest
from concurrent.futures import ThreadPoolExecutor
from functools import partial
from types import SimpleNamespace
import numpy as np
import requests
from sglang.srt.utils import kill_process_tree
from sglang.test.few_shot_gsm8k import run_eval
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,
)
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
def test_request_abort(self):
concurrency = 4
threads = [
threading.Thread(target=self.send_request) for _ in range(concurrency)
] + [
threading.Thread(target=self.send_requests_abort)
for _ in range(concurrency)
]
for worker in threads:
worker.start()
for p in threads:
p.join()
def test_max_token_one(self):
requests.get(self.base_url + "/flush_cache")
args = SimpleNamespace(
num_shots=5,
data_path=None,
num_questions=200,
max_new_tokens=1,
parallel=128,
host="http://127.0.0.1",
port=int(self.base_url.split(":")[-1]),
)
# Just run and check it does not hang
metrics = run_eval(args)
self.assertGreater(metrics["output_throughput"], 50)
def test_gsm8k(self):
requests.get(self.base_url + "/flush_cache")
args = SimpleNamespace(
num_shots=5,
data_path=None,
num_questions=200,
max_new_tokens=512,
parallel=128,
host="http://127.0.0.1",
port=int(self.base_url.split(":")[-1]),
)
metrics = run_eval(args)
print(f"{metrics=}")
self.assertGreater(metrics["accuracy"], 0.20)
server_info = requests.get(self.base_url + "/get_server_info").json()
avg_spec_accept_length = server_info["internal_states"][0][
"avg_spec_accept_length"
]
print(f"{avg_spec_accept_length=}")
speculative_eagle_topk = server_info["speculative_eagle_topk"]
if speculative_eagle_topk == 1:
self.assertGreater(avg_spec_accept_length, 2.5)
else:
self.assertGreater(avg_spec_accept_length, 3.5)
# Wait a little bit so that the memory check happens.
time.sleep(4)
def test_logprob_start_len(self):
logprob_start_len = 4
new_tokens = 4
prompts = [
"I have a very good idea on",
"Today is a sunndy day and",
]
response = requests.post(
self.base_url + "/generate",
json={
"text": prompts,
"sampling_params": {
"temperature": 0,
"max_new_tokens": new_tokens,
},
"return_logprob": True,
"top_logprobs_num": 5,
"logprob_start_len": logprob_start_len,
},
)
response_json = response.json()
print(json.dumps(response_json, indent=2))
for res in response_json:
self.assertEqual(
res["meta_info"]["prompt_tokens"],
logprob_start_len + len(res["meta_info"]["input_token_logprobs"]),
)
self.assertEqual(res["meta_info"]["completion_tokens"], new_tokens)
self.assertEqual(len(res["meta_info"]["output_token_logprobs"]), new_tokens)
def test_logprob_match(self):
"""Test the output logprobs are close to the input logprobs if we run a prefill again."""
def run_generate(
prompt,
return_logprob=False,
max_new_tokens=512,
logprob_start_len=-1,
temperature=1.0,
):
if isinstance(prompt, str):
prompt_kwargs = {"text": prompt}
else:
prompt_kwargs = {"input_ids": prompt}
response = requests.post(
self.base_url + "/generate",
json={
**prompt_kwargs,
"sampling_params": {
"temperature": temperature,
"max_new_tokens": max_new_tokens,
"ignore_eos": True,
},
"return_logprob": return_logprob,
"return_text_in_logprobs": True,
"logprob_start_len": logprob_start_len,
"temp_scaled_logprobs": True,
},
)
return response.json()
prompt = "I have a very good idea on how to"
for temperature in [1.0]:
gen = run_generate(
prompt,
return_logprob=True,
logprob_start_len=0,
temperature=temperature,
)
output_logprobs = np.array(
[x[0] for x in gen["meta_info"]["output_token_logprobs"]]
)
num_prompts_tokens = gen["meta_info"]["prompt_tokens"]
input_tokens = [x[1] for x in gen["meta_info"]["input_token_logprobs"]]
output_tokens = [x[1] for x in gen["meta_info"]["output_token_logprobs"]]
new_prompt = input_tokens + output_tokens
score = run_generate(
new_prompt,
return_logprob=True,
logprob_start_len=0,
max_new_tokens=0,
temperature=temperature,
)
output_logprobs_score = np.array(
[
x[0]
for x in score["meta_info"]["input_token_logprobs"][
num_prompts_tokens:
]
]
)
print(f"{output_logprobs[-10:]=}")
print(f"{output_logprobs_score[-10:]=}")
diff = np.abs(output_logprobs - output_logprobs_score)
max_diff = np.max(diff)
self.assertLess(max_diff, 0.255)
def test_logprob_mixed(self):
args = []
temperature = 0
# input_len, output_len, temperature, logprob_start_len, return_logprob, top_logprobs_num
# Llama 2 context length seems to be only 2k, so we can only test small length.
for input_len in [200, 500, 1000, 2000]:
for output_len in [4, 8]:
for logprob_start_len in [0, 100, 300, 800, 1998]:
for return_logprob in [True, False]:
for top_logprobs_num in [0, 5]:
if logprob_start_len >= input_len:
continue
args.append(
(
input_len,
output_len,
temperature,
logprob_start_len,
return_logprob,
top_logprobs_num,
)
)
random.shuffle(args)
func = partial(run_logprob_check, self)
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 = [
{},
{},
{},
{"frequency_penalty": 2},
{"presence_penalty": 1},
{"min_new_tokens": 16},
{"frequency_penalty": 0.2},
{"presence_penalty": 0.4},
{"min_new_tokens": 8},
{"frequency_penalty": 0.4, "presence_penalty": 0.8},
{"frequency_penalty": 0.4, "min_new_tokens": 12},
{"presence_penalty": 0.8, "min_new_tokens": 12},
{"presence_penalty": -0.3, "frequency_penalty": 1.3, "min_new_tokens": 32},
{"presence_penalty": 0.3, "frequency_penalty": -1.3, "min_new_tokens": 32},
]
random.shuffle(args * 5)
with ThreadPoolExecutor(8) as executor:
list(executor.map(self.run_decode, args))
def test_constrained_decoding(self):
messages = [
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "Give me a json"},
]
response = requests.post(
self.base_url + "/v1/chat/completions",
json={
"model": DEFAULT_EAGLE_TARGET_MODEL_FOR_TEST,
"messages": messages,
"temperature": 0,
"response_format": {"type": "json_object"},
},
)
self.assertEqual(response.status_code, 200)
res = response.json()
# Validate response structure
self.assertIn("choices", res)
self.assertEqual(len(res["choices"]), 1)
self.assertIn("message", res["choices"][0])
self.assertIn("content", res["choices"][0]["message"])
# Validate JSON content
content_json = res["choices"][0]["message"]["content"]
is_valid_json = True
try:
content = json.loads(content_json)
self.assertIsInstance(content, dict)
except Exception:
print(f"parse JSON failed: {content_json}")
is_valid_json = False
self.assertTrue(is_valid_json)
class TestEAGLERetract(TestEAGLEServer):
@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,
],
)
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 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 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",
],
)
if __name__ == "__main__":
unittest.main()
test/srt/test_eagle_infer_beta.py
-96
View File
@@ -1,96 +0,0 @@
import unittest
from types import SimpleNamespace
from sglang.srt.environ import envs
from sglang.srt.utils import kill_process_tree
from sglang.test.few_shot_gsm8k import run_eval
from sglang.test.kits.matched_stop_kit import MatchedStopMixin
from sglang.test.kits.radix_cache_server_kit import run_radix_attention_test
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,
)
class TestEagleServerBase(CustomTestCase, MatchedStopMixin):
max_running_requests = 64
attention_backend = "triton"
spec_steps = 5
spec_topk = 1
spec_draft_tokens = 6
page_size = 1
other_launch_args = []
model = DEFAULT_EAGLE_TARGET_MODEL_FOR_TEST
draft_model = DEFAULT_EAGLE_DRAFT_MODEL_FOR_TEST
@classmethod
def setUpClass(cls):
cls.base_url = DEFAULT_URL_FOR_TEST
launch_args = [
"--trust-remote-code",
"--attention-backend",
cls.attention_backend,
"--speculative-algorithm",
"EAGLE",
"--speculative-draft-model",
cls.draft_model,
"--speculative-num-steps",
cls.spec_steps,
"--speculative-eagle-topk",
cls.spec_topk,
"--speculative-num-draft-tokens",
cls.spec_draft_tokens,
"--page-size",
str(cls.page_size),
"--mem-fraction-static",
"0.75",
"--max-running-requests",
str(cls.max_running_requests),
"--cuda-graph-bs",
*[str(i) for i in range(1, cls.max_running_requests + 1)],
]
launch_args.extend(cls.other_launch_args)
with envs.SGLANG_ENABLE_SPEC_V2.override(True):
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=launch_args,
)
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
def test_radix_attention(self):
run_radix_attention_test(self.base_url)
assert self.process.poll() is None
def test_gsm8k(self):
args = SimpleNamespace(
num_shots=5,
data_path=None,
num_questions=1000,
max_new_tokens=512,
parallel=128,
host="http://127.0.0.1",
port=int(self.base_url.split(":")[-1]),
)
metrics = run_eval(args)
print(f"TestEagleLargeBS -- {metrics=}")
self.assertGreater(
metrics["accuracy"], 0.23
) # 0.3333 for 60 questions; 0.234 for 1319 questions
assert self.process.poll() is None
class TestEagleServerPage(TestEagleServerBase):
other_launch_args = ["--page-size", "64"]
if __name__ == "__main__":
unittest.main()