Refactor EAGLE infer tests: shared fixture + kits + overlap matrix (#26871)
This commit is contained in:
Liangsheng Yin
@@ -0,0 +1,566 @@
|
||||
"""Reusable test-method mixins (kits) for EAGLE/EAGLE3 spec-decoding servers.
|
||||
|
||||
Pair these with ``SpecEagleServerBase`` (sglang.test.server_fixtures.spec_eagle_fixture).
|
||||
Each kit is a cohesive group of ``test_*`` methods with no launch logic; concrete
|
||||
test classes mix in the fixture (which owns launch knobs) + whichever kits apply.
|
||||
|
||||
Thresholds are read off ``self`` so a config can tune them as class attributes.
|
||||
"""
|
||||
|
||||
import concurrent.futures
|
||||
import json
|
||||
import random
|
||||
import threading
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from functools import partial
|
||||
from types import SimpleNamespace
|
||||
|
||||
import numpy as np
|
||||
import requests
|
||||
|
||||
from sglang.srt.utils.common import kill_process_tree
|
||||
from sglang.test.kits.radix_cache_server_kit import run_radix_attention_test
|
||||
from sglang.test.run_eval import run_eval
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
popen_launch_server,
|
||||
run_logprob_check,
|
||||
)
|
||||
|
||||
|
||||
class SpecCorrectnessKit:
|
||||
"""Acceptance-quality + EOS checks (single server, cheap)."""
|
||||
|
||||
# Tunable thresholds (override per config class).
|
||||
acc_length_thres = 3.1
|
||||
batch_accept_len_thres = 1.75
|
||||
|
||||
def test_acc_length(self):
|
||||
prompt = [
|
||||
"Human: Give me a fully functional FastAPI server. Show the python code.\n\nAssistant:",
|
||||
] * 5
|
||||
sampling_params = {"temperature": 0, "max_new_tokens": 512}
|
||||
output = requests.post(
|
||||
self.base_url + "/generate",
|
||||
json={"text": prompt, "sampling_params": sampling_params},
|
||||
).json()[0]
|
||||
|
||||
meta = output["meta_info"]
|
||||
if "spec_verify_ct" in meta and meta["spec_verify_ct"] > 0:
|
||||
acc_length = meta["completion_tokens"] / meta["spec_verify_ct"]
|
||||
else:
|
||||
acc_length = 1.0
|
||||
print(f"{acc_length=:.4f}")
|
||||
self.assertGreater(acc_length, self.acc_length_thres)
|
||||
|
||||
def test_batch_generation(self):
|
||||
prompts = [
|
||||
"Hello, my name is",
|
||||
"The president of the United States is",
|
||||
"The capital of France is",
|
||||
"The future of AI is",
|
||||
]
|
||||
results = requests.post(
|
||||
self.base_url + "/generate",
|
||||
json={
|
||||
"text": prompts,
|
||||
"sampling_params": {"temperature": 0, "max_new_tokens": 50},
|
||||
},
|
||||
).json()
|
||||
# Accept length from per-request meta_info (self-contained). The
|
||||
# internal_states `avg_spec_accept_length` isn't populated on the v1 /
|
||||
# disable-overlap path after a small batch, so don't read server_info.
|
||||
total_completion, total_verify = 0, 0
|
||||
for r in results:
|
||||
self.assertIn("text", r, f"Server error: {r}")
|
||||
meta = r["meta_info"]
|
||||
total_completion += meta["completion_tokens"]
|
||||
total_verify += meta.get("spec_verify_ct", 0)
|
||||
if total_verify > 0:
|
||||
acc_length = total_completion / total_verify
|
||||
print(f"batch {acc_length=:.4f}")
|
||||
self.assertGreater(acc_length, self.batch_accept_len_thres)
|
||||
|
||||
def test_eos_token(self):
|
||||
prompt = "[INST] <<SYS>>\nYou are a helpful assistant.\n<</SYS>>\nToday is a sunny day and I like [/INST]"
|
||||
res = requests.post(
|
||||
self.base_url + "/generate",
|
||||
json={
|
||||
"text": prompt,
|
||||
"sampling_params": {
|
||||
"temperature": 0.1,
|
||||
"max_new_tokens": 1024,
|
||||
"skip_special_tokens": False,
|
||||
},
|
||||
},
|
||||
).json()
|
||||
output = res["text"]
|
||||
tokens = self.tokenizer.encode(output, truncation=False)
|
||||
self.assertNotIn(self.tokenizer.eos_token_id, tokens)
|
||||
|
||||
def test_first_token_finish(self):
|
||||
# Very short max_new_tokens (1-3): exercise the immediate-finish path,
|
||||
# where a request stops within the first draft window. Just must not crash.
|
||||
prompts = [
|
||||
f"There are {i} apples on the table. How to divide them equally?"
|
||||
for i in range(8)
|
||||
]
|
||||
sampling_params = [
|
||||
{"temperature": 0, "max_new_tokens": random.randint(1, 3)} for _ in range(8)
|
||||
]
|
||||
results = requests.post(
|
||||
self.base_url + "/generate",
|
||||
json={"text": prompts, "sampling_params": sampling_params},
|
||||
).json()
|
||||
for r in results:
|
||||
self.assertIn("text", r, f"Server error: {r}")
|
||||
|
||||
|
||||
def _greedy(url, text, max_new_tokens=48):
|
||||
return requests.post(
|
||||
url + "/generate",
|
||||
json={
|
||||
"text": text,
|
||||
"sampling_params": {"temperature": 0, "max_new_tokens": max_new_tokens},
|
||||
},
|
||||
).json()["text"]
|
||||
|
||||
|
||||
class SpecParityKit:
|
||||
"""Lossless output parity vs a non-spec reference.
|
||||
|
||||
Sequential (NOT concurrent): launch a non-spec reference server on the
|
||||
standard port, capture greedy outputs, tear it down, THEN let the fixture
|
||||
launch the spec server. Only one model is resident at a time -- two 8B
|
||||
servers don't fit on one GPU. Mix this kit FIRST in the bases so its
|
||||
setUpClass runs before the fixture's: ``class T(SpecParityKit, Eagle3Base)``.
|
||||
"""
|
||||
|
||||
parity_prompts = [
|
||||
"The capital of France is",
|
||||
"Once upon a time, there was a",
|
||||
"The three primary colors are",
|
||||
"def fibonacci(n):",
|
||||
]
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
ref_url = DEFAULT_URL_FOR_TEST
|
||||
ref_proc = popen_launch_server(
|
||||
cls.model,
|
||||
ref_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=[
|
||||
"--mem-fraction-static",
|
||||
"0.8", # ref alone -> full GPU available
|
||||
"--attention-backend",
|
||||
cls.attention_backend,
|
||||
"--page-size",
|
||||
"1",
|
||||
"--dtype",
|
||||
cls.dtype,
|
||||
*(["--trust-remote-code"] if cls.trust_remote_code else []),
|
||||
],
|
||||
)
|
||||
try:
|
||||
cls.parity_ref_outputs = {
|
||||
p: _greedy(ref_url, p) for p in cls.parity_prompts
|
||||
}
|
||||
finally:
|
||||
kill_process_tree(ref_proc.pid, wait_timeout=60)
|
||||
# Now the spec server (same port; ref is gone).
|
||||
super().setUpClass()
|
||||
|
||||
def test_parity_vs_reference(self):
|
||||
"""Spec decode greedy output must equal the non-spec reference."""
|
||||
for prompt in self.parity_prompts:
|
||||
spec_out = _greedy(self.base_url, prompt)
|
||||
self.assertEqual(
|
||||
spec_out,
|
||||
self.parity_ref_outputs[prompt],
|
||||
f"spec != ref for prompt {prompt!r}",
|
||||
)
|
||||
|
||||
|
||||
class SpecAccuracyKit:
|
||||
"""gsm8k accuracy + acceptance length, and throughput at max_tokens=1."""
|
||||
|
||||
gsm8k_num_examples = 200
|
||||
gsm8k_score_thres = 0.20
|
||||
gsm8k_check_accept_len = True
|
||||
# If set, use this; else fall back to topk-based default (2.5 / 3.47).
|
||||
gsm8k_accept_len_thres = None
|
||||
|
||||
def test_gsm8k(self):
|
||||
requests.get(self.base_url + "/flush_cache")
|
||||
args = SimpleNamespace(
|
||||
base_url=self.base_url,
|
||||
model=self.model,
|
||||
eval_name="gsm8k",
|
||||
api="completion",
|
||||
max_tokens=512,
|
||||
num_examples=self.gsm8k_num_examples,
|
||||
num_threads=128,
|
||||
)
|
||||
metrics = run_eval(args)
|
||||
print(f"{metrics=}")
|
||||
self.assertGreater(metrics["score"], self.gsm8k_score_thres)
|
||||
|
||||
if self.gsm8k_check_accept_len:
|
||||
server_info = requests.get(self.base_url + "/server_info").json()
|
||||
avg_spec_accept_length = server_info["internal_states"][0].get(
|
||||
"avg_spec_accept_length"
|
||||
)
|
||||
print(f"{avg_spec_accept_length=}")
|
||||
# The metric isn't always populated (e.g. v1 / disable-overlap).
|
||||
# Only enforce the threshold when it's reported.
|
||||
if avg_spec_accept_length is not None:
|
||||
topk = server_info["speculative_eagle_topk"]
|
||||
thres = self.gsm8k_accept_len_thres
|
||||
if thres is None:
|
||||
thres = 2.5 if topk == 1 else 3.47
|
||||
self.assertGreater(avg_spec_accept_length, thres)
|
||||
|
||||
|
||||
class SpecPerfKit:
|
||||
"""Throughput perf check (GPU-specific -> run on the reference/Hopper runner)."""
|
||||
|
||||
perf_output_throughput_thres = 50
|
||||
|
||||
def test_max_token_one(self):
|
||||
requests.get(self.base_url + "/flush_cache")
|
||||
args = SimpleNamespace(
|
||||
base_url=self.base_url,
|
||||
model=self.model,
|
||||
eval_name="gsm8k",
|
||||
api="completion",
|
||||
max_tokens=1,
|
||||
num_examples=200,
|
||||
num_threads=128,
|
||||
)
|
||||
metrics = run_eval(args)
|
||||
self.assertGreater(
|
||||
metrics["output_throughput"], self.perf_output_throughput_thres
|
||||
)
|
||||
|
||||
|
||||
class SpecLogprobKit:
|
||||
"""Logprob correctness: start_len, prefill-rescore match, mixed sweep,
|
||||
spec-v2 decode-vs-prefill match, and ragged token_ids_logprob."""
|
||||
|
||||
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()
|
||||
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):
|
||||
"""Output logprobs should match a fresh prefill of the same sequence."""
|
||||
|
||||
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,
|
||||
},
|
||||
)
|
||||
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:
|
||||
]
|
||||
]
|
||||
)
|
||||
|
||||
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
|
||||
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 test_logprob_spec_v2_match(self):
|
||||
"""Verify spec v2 decode logprobs match prefill scoring logprobs."""
|
||||
top_k = 5
|
||||
probe_token_ids = [1, 2, 10, 100, 1000]
|
||||
prompts = [
|
||||
"The capital of France is",
|
||||
"Explain quantum computing in simple terms:",
|
||||
]
|
||||
|
||||
for round_idx, prompt in enumerate(prompts):
|
||||
with self.subTest(round=round_idx, prompt=prompt):
|
||||
gen_res = requests.post(
|
||||
self.base_url + "/generate",
|
||||
json={
|
||||
"text": prompt,
|
||||
"sampling_params": {
|
||||
"temperature": 0,
|
||||
"max_new_tokens": 32,
|
||||
"ignore_eos": True,
|
||||
},
|
||||
"return_logprob": True,
|
||||
"top_logprobs_num": top_k,
|
||||
"token_ids_logprob": probe_token_ids,
|
||||
"logprob_start_len": 0,
|
||||
},
|
||||
).json()
|
||||
|
||||
decode_logprobs = gen_res["meta_info"]["output_token_logprobs"]
|
||||
decode_top_logprobs = gen_res["meta_info"]["output_top_logprobs"]
|
||||
decode_tid_logprobs = gen_res["meta_info"]["output_token_ids_logprobs"]
|
||||
input_token_ids = [
|
||||
t[1] for t in gen_res["meta_info"]["input_token_logprobs"]
|
||||
]
|
||||
output_token_ids = [t[1] for t in decode_logprobs]
|
||||
num_prompt_tokens = gen_res["meta_info"]["prompt_tokens"]
|
||||
|
||||
score_res = requests.post(
|
||||
self.base_url + "/generate",
|
||||
json={
|
||||
"input_ids": input_token_ids + output_token_ids,
|
||||
"sampling_params": {
|
||||
"temperature": 0,
|
||||
"max_new_tokens": 0,
|
||||
},
|
||||
"return_logprob": True,
|
||||
"top_logprobs_num": top_k,
|
||||
"token_ids_logprob": probe_token_ids,
|
||||
"logprob_start_len": 0,
|
||||
},
|
||||
).json()
|
||||
|
||||
score_logprobs = score_res["meta_info"]["input_token_logprobs"][
|
||||
num_prompt_tokens:
|
||||
]
|
||||
score_top_logprobs = score_res["meta_info"]["input_top_logprobs"][
|
||||
num_prompt_tokens:
|
||||
]
|
||||
score_tid_logprobs = score_res["meta_info"]["input_token_ids_logprobs"][
|
||||
num_prompt_tokens:
|
||||
]
|
||||
|
||||
self.assertEqual(len(decode_logprobs), len(score_logprobs))
|
||||
|
||||
decode_vals = np.array([t[0] for t in decode_logprobs])
|
||||
score_vals = np.array([t[0] for t in score_logprobs])
|
||||
max_diff = np.max(np.abs(decode_vals - score_vals))
|
||||
print(f"[round {round_idx}] logprob max_diff={max_diff:.6f}")
|
||||
self.assertLess(max_diff, 0.255)
|
||||
|
||||
for pos in range(len(decode_logprobs)):
|
||||
dec_top = {t[1]: t[0] for t in decode_top_logprobs[pos]}
|
||||
scr_top = {t[1]: t[0] for t in score_top_logprobs[pos]}
|
||||
common_ids = set(dec_top.keys()) & set(scr_top.keys())
|
||||
self.assertGreater(len(common_ids), 0)
|
||||
for tid in common_ids:
|
||||
self.assertAlmostEqual(dec_top[tid], scr_top[tid], delta=0.255)
|
||||
|
||||
self.assertEqual(len(decode_tid_logprobs), len(score_tid_logprobs))
|
||||
for pos in range(len(decode_tid_logprobs)):
|
||||
dec_tid = {t[1]: t[0] for t in decode_tid_logprobs[pos]}
|
||||
scr_tid = {t[1]: t[0] for t in score_tid_logprobs[pos]}
|
||||
self.assertEqual(set(dec_tid.keys()), set(scr_tid.keys()))
|
||||
for tid in dec_tid:
|
||||
self.assertAlmostEqual(dec_tid[tid], scr_tid[tid], delta=0.255)
|
||||
|
||||
def test_token_ids_logprob_ragged(self):
|
||||
"""Regression: ragged token_ids_logprob lists in one batch must not crash."""
|
||||
|
||||
def send(probe_ids):
|
||||
return requests.post(
|
||||
self.base_url + "/generate",
|
||||
json={
|
||||
"text": "Hello world",
|
||||
"sampling_params": {"temperature": 0, "max_new_tokens": 8},
|
||||
"return_logprob": True,
|
||||
"top_logprobs_num": 3,
|
||||
"token_ids_logprob": probe_ids,
|
||||
},
|
||||
).json()
|
||||
|
||||
ragged_probes = [
|
||||
[1, 2],
|
||||
[3, 4, 5],
|
||||
[6],
|
||||
[10, 20, 30, 40],
|
||||
[1, 2],
|
||||
[3, 4, 5],
|
||||
[6],
|
||||
[10, 20, 30, 40],
|
||||
]
|
||||
with concurrent.futures.ThreadPoolExecutor(max_workers=8) as pool:
|
||||
futs = [pool.submit(send, ids) for ids in ragged_probes]
|
||||
for f in concurrent.futures.as_completed(futs):
|
||||
res = f.result()
|
||||
self.assertIn("text", res, f"Server error: {res}")
|
||||
|
||||
|
||||
class SpecPenaltyKit:
|
||||
"""Penalty parameters under concurrency must not crash / corrupt output."""
|
||||
|
||||
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))
|
||||
|
||||
|
||||
class SpecFeatureKit:
|
||||
"""Radix attention, constrained decoding, concurrent abort."""
|
||||
|
||||
def test_radix_attention(self):
|
||||
run_radix_attention_test(self.base_url)
|
||||
self.assertIsNone(self.process.poll())
|
||||
|
||||
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_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": self.model,
|
||||
"messages": messages,
|
||||
"temperature": 0,
|
||||
"response_format": {"type": "json_object"},
|
||||
},
|
||||
)
|
||||
self.assertEqual(response.status_code, 200)
|
||||
res = response.json()
|
||||
self.assertIn("choices", res)
|
||||
self.assertEqual(len(res["choices"]), 1)
|
||||
self.assertIn("message", res["choices"][0])
|
||||
self.assertIn("content", res["choices"][0]["message"])
|
||||
|
||||
content_json = res["choices"][0]["message"]["content"]
|
||||
try:
|
||||
content = json.loads(content_json)
|
||||
self.assertIsInstance(content, dict)
|
||||
except Exception:
|
||||
self.fail(f"parse JSON failed: {content_json}")
|
||||
@@ -0,0 +1,227 @@
|
||||
"""Unified EAGLE/EAGLE3 speculative-decoding server fixture.
|
||||
|
||||
A single popen-server base whose launch is fully described by class attributes,
|
||||
so concrete test classes only flip knobs (overlap on/off, model, topk, page size,
|
||||
backend, env overrides). Pair it with the kits in
|
||||
``sglang.test.kits.spec_server_kits`` to assemble test classes.
|
||||
|
||||
The primary axis is ``disable_overlap``:
|
||||
- ``False`` -> spec v2 (overlap scheduler)
|
||||
- ``True`` -> spec v1 (overlap disabled)
|
||||
"""
|
||||
|
||||
import contextlib
|
||||
import random
|
||||
import time
|
||||
|
||||
import requests
|
||||
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.utils.common import kill_process_tree
|
||||
from sglang.srt.utils.hf_transformers_utils import get_tokenizer
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_DRAFT_MODEL_EAGLE,
|
||||
DEFAULT_DRAFT_MODEL_EAGLE3,
|
||||
DEFAULT_TARGET_MODEL_EAGLE,
|
||||
DEFAULT_TARGET_MODEL_EAGLE3,
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
popen_launch_server,
|
||||
)
|
||||
|
||||
# Chat-style prompts shared by send_request / send_requests_abort.
|
||||
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 SpecEagleServerBase(CustomTestCase):
|
||||
"""Launch a single EAGLE/EAGLE3 spec server from class-attribute knobs."""
|
||||
|
||||
# -- model --
|
||||
model = DEFAULT_TARGET_MODEL_EAGLE3
|
||||
draft_model = DEFAULT_DRAFT_MODEL_EAGLE3
|
||||
spec_algo = "EAGLE3"
|
||||
|
||||
# -- speculative config --
|
||||
spec_steps = 5
|
||||
spec_topk = 1
|
||||
spec_tokens = 6
|
||||
|
||||
# -- runtime config --
|
||||
page_size = 1
|
||||
attention_backend = "flashinfer"
|
||||
# Primary axis: False -> spec v2 (overlap); True -> spec v1 (overlap off).
|
||||
disable_overlap = False
|
||||
mem_fraction_static = 0.75
|
||||
max_running_requests = 8
|
||||
chunked_prefill_size = 128
|
||||
dtype = "float16"
|
||||
cuda_graph_max_bs = None
|
||||
trust_remote_code = True
|
||||
|
||||
# -- extras --
|
||||
# env_overrides: iterable of (env_var_obj, value) applied only around launch.
|
||||
env_overrides = ()
|
||||
extra_args = ()
|
||||
|
||||
@classmethod
|
||||
def _launch_args(cls):
|
||||
args = [
|
||||
"--speculative-algorithm",
|
||||
cls.spec_algo,
|
||||
"--speculative-draft-model-path",
|
||||
cls.draft_model,
|
||||
"--speculative-num-steps",
|
||||
str(cls.spec_steps),
|
||||
"--speculative-eagle-topk",
|
||||
str(cls.spec_topk),
|
||||
"--speculative-num-draft-tokens",
|
||||
str(cls.spec_tokens),
|
||||
"--page-size",
|
||||
str(cls.page_size),
|
||||
"--attention-backend",
|
||||
cls.attention_backend,
|
||||
"--mem-fraction-static",
|
||||
str(cls.mem_fraction_static),
|
||||
"--max-running-requests",
|
||||
str(cls.max_running_requests),
|
||||
"--chunked-prefill-size",
|
||||
str(cls.chunked_prefill_size),
|
||||
"--dtype",
|
||||
cls.dtype,
|
||||
]
|
||||
if cls.disable_overlap:
|
||||
args.append("--disable-overlap-schedule")
|
||||
if cls.trust_remote_code:
|
||||
args.append("--trust-remote-code")
|
||||
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]
|
||||
return args
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
# Alias so kit methods can use either name.
|
||||
cls.target_model = cls.model
|
||||
cls._tokenizer = None
|
||||
with contextlib.ExitStack() as stack:
|
||||
stack.enter_context(envs.SGLANG_ENABLE_ASYNC_ASSERT.override(True))
|
||||
stack.enter_context(
|
||||
envs.SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN.override(True)
|
||||
)
|
||||
for env_var, value in cls.env_overrides:
|
||||
stack.enter_context(env_var.override(value))
|
||||
cls.process = popen_launch_server(
|
||||
cls.model,
|
||||
cls.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=cls._launch_args(),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
kill_process_tree(cls.process.pid, wait_timeout=60)
|
||||
|
||||
@property
|
||||
def tokenizer(self):
|
||||
if type(self)._tokenizer is None:
|
||||
type(self)._tokenizer = get_tokenizer(self.model)
|
||||
return type(self)._tokenizer
|
||||
|
||||
# -- helpers used by kits --
|
||||
|
||||
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):
|
||||
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": 1,
|
||||
"temperature": 0.7,
|
||||
**sampling_params,
|
||||
},
|
||||
"return_logprob": True,
|
||||
"top_logprobs_num": 5,
|
||||
"return_text_in_logprobs": True,
|
||||
"logprob_start_len": 0,
|
||||
},
|
||||
)
|
||||
self.assertEqual(response.status_code, 200)
|
||||
|
||||
|
||||
class Eagle3Base(SpecEagleServerBase):
|
||||
"""EAGLE3 (Llama-3.1) config preset, topk=1 / page_size=1 by default."""
|
||||
|
||||
model = DEFAULT_TARGET_MODEL_EAGLE3
|
||||
draft_model = DEFAULT_DRAFT_MODEL_EAGLE3
|
||||
spec_algo = "EAGLE3"
|
||||
spec_steps = 5
|
||||
spec_topk = 1
|
||||
spec_tokens = 6
|
||||
attention_backend = "flashinfer"
|
||||
chunked_prefill_size = 1024
|
||||
# EAGLE3 topk=1 accepts modestly; tune against CI if needed.
|
||||
acc_length_thres = 1.6
|
||||
batch_accept_len_thres = 1.3
|
||||
gsm8k_score_thres = 0.7
|
||||
gsm8k_accept_len_thres = 1.3
|
||||
|
||||
|
||||
class EagleLlama2Base(SpecEagleServerBase):
|
||||
"""EAGLE (Llama-2) config preset. topk=8 tree -> spec v1; gsm8k is low."""
|
||||
|
||||
model = DEFAULT_TARGET_MODEL_EAGLE
|
||||
draft_model = DEFAULT_DRAFT_MODEL_EAGLE
|
||||
spec_algo = "EAGLE"
|
||||
spec_steps = 5
|
||||
spec_topk = 8
|
||||
spec_tokens = 64
|
||||
attention_backend = "flashinfer"
|
||||
chunked_prefill_size = 128
|
||||
mem_fraction_static = 0.7
|
||||
gsm8k_score_thres = 0.20
|
||||
acc_length_thres = 3.0
|
||||
batch_accept_len_thres = 1.8
|
||||
# EAGLE topk>1 already routes to v1; force it explicitly to preserve intent.
|
||||
env_overrides = ((envs.SGLANG_ENABLE_SPEC_V2, False),)
|
||||
Reference in New Issue
Block a user