From 1bff7a290f98c9fb92f3844d0fa58727b88e218e Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Mon, 1 Jun 2026 03:55:01 -0700 Subject: [PATCH] Refactor EAGLE infer tests: shared fixture + kits + overlap matrix (#26871) --- python/sglang/test/kits/spec_server_kits.py | 566 ++++++++++++++++++ .../server_fixtures/spec_eagle_fixture.py | 227 +++++++ .../spec/eagle/test_eagle_infer_a.py | 203 ------- .../spec/eagle/test_eagle_infer_b.py | 388 ------------ .../spec/eagle/test_eagle_infer_beta.py | 281 --------- test/registered/spec/eagle/test_spec_eagle.py | 51 ++ .../spec/eagle/test_spec_eagle_fa3.py | 49 ++ .../spec/eagle/test_spec_eagle_page.py | 47 ++ .../spec/eagle/test_spec_eagle_parity.py | 28 + .../spec/eagle/test_spec_eagle_stress.py | 65 ++ .../spec/eagle/test_spec_eagle_topk.py | 67 +++ .../spec/eagle/test_spec_eagle_triton.py | 47 ++ 12 files changed, 1147 insertions(+), 872 deletions(-) create mode 100644 python/sglang/test/kits/spec_server_kits.py create mode 100644 python/sglang/test/server_fixtures/spec_eagle_fixture.py delete mode 100644 test/registered/spec/eagle/test_eagle_infer_a.py delete mode 100644 test/registered/spec/eagle/test_eagle_infer_b.py delete mode 100644 test/registered/spec/eagle/test_eagle_infer_beta.py create mode 100644 test/registered/spec/eagle/test_spec_eagle.py create mode 100644 test/registered/spec/eagle/test_spec_eagle_fa3.py create mode 100644 test/registered/spec/eagle/test_spec_eagle_page.py create mode 100644 test/registered/spec/eagle/test_spec_eagle_parity.py create mode 100644 test/registered/spec/eagle/test_spec_eagle_stress.py create mode 100644 test/registered/spec/eagle/test_spec_eagle_topk.py create mode 100644 test/registered/spec/eagle/test_spec_eagle_triton.py diff --git a/python/sglang/test/kits/spec_server_kits.py b/python/sglang/test/kits/spec_server_kits.py new file mode 100644 index 000000000..f7c7c3ff0 --- /dev/null +++ b/python/sglang/test/kits/spec_server_kits.py @@ -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] <>\nYou are a helpful assistant.\n<>\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}") diff --git a/python/sglang/test/server_fixtures/spec_eagle_fixture.py b/python/sglang/test/server_fixtures/spec_eagle_fixture.py new file mode 100644 index 000000000..32f5b86ba --- /dev/null +++ b/python/sglang/test/server_fixtures/spec_eagle_fixture.py @@ -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] <>\\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 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),) diff --git a/test/registered/spec/eagle/test_eagle_infer_a.py b/test/registered/spec/eagle/test_eagle_infer_a.py deleted file mode 100644 index 0354b958c..000000000 --- a/test/registered/spec/eagle/test_eagle_infer_a.py +++ /dev/null @@ -1,203 +0,0 @@ -import random -import unittest - -import sglang as sgl -from sglang.srt.environ import envs -from sglang.srt.utils.hf_transformers_utils import get_tokenizer -from sglang.test.ci.ci_register import register_cuda_ci -from sglang.test.test_utils import ( - DEFAULT_DRAFT_MODEL_EAGLE, - DEFAULT_DRAFT_MODEL_EAGLE3, - DEFAULT_TARGET_MODEL_EAGLE, - DEFAULT_TARGET_MODEL_EAGLE3, - CustomTestCase, -) - -register_cuda_ci(est_time=357, stage="extra-a", runner_config="1-gpu-large") - - -class TestEAGLEEngine(CustomTestCase): - BASE_CONFIG = { - "model_path": DEFAULT_TARGET_MODEL_EAGLE, - "speculative_draft_model_path": DEFAULT_DRAFT_MODEL_EAGLE, - "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, - } - - @classmethod - def setUpClass(cls): - envs.SGLANG_ENABLE_SPEC_V2.set(False) - - @classmethod - def tearDownClass(cls): - envs.SGLANG_ENABLE_SPEC_V2.clear() - - 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_first_token_finish(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_first_token_finish(self, engine): - prompt = [ - f"There are {i} apples on the table. How to divide them equally?" - for i in range(8) - ] - params = [ - {"temperature": 0, "max_new_tokens": random.randint(1, 3)} for _ in range(8) - ] - outputs = engine.generate(prompt, params) - for i, output in enumerate(outputs): - print(f"Prompt: {prompt[i]}") - print(f"Generated: {output['text']}") - print("-" * 40) - - def _test_eos_token(self, engine): - prompt = "[INST] <>\nYou are a helpful assistant.\n<>\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_TARGET_MODEL_EAGLE) - 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_TARGET_MODEL_EAGLE3, - "speculative_draft_model_path": DEFAULT_DRAFT_MODEL_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, - } - - -if __name__ == "__main__": - unittest.main() diff --git a/test/registered/spec/eagle/test_eagle_infer_b.py b/test/registered/spec/eagle/test_eagle_infer_b.py deleted file mode 100644 index 973565016..000000000 --- a/test/registered/spec/eagle/test_eagle_infer_b.py +++ /dev/null @@ -1,388 +0,0 @@ -import json -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.environ import envs -from sglang.test.ci.ci_register import register_cuda_ci -from sglang.test.kits.abort_timeout_kit import ( - AbortAllMixin, - RunningTimeoutTwoWaveMixin, - WaitingTimeoutMixin, -) -from sglang.test.kits.radix_cache_server_kit import run_radix_attention_test -from sglang.test.run_eval import run_eval -from sglang.test.server_fixtures.eagle_fixture import EagleServerBase -from sglang.test.test_utils import DEFAULT_TARGET_MODEL_EAGLE, run_logprob_check - -register_cuda_ci(est_time=847, stage="base-b", runner_config="1-gpu-large") - - -class TestEAGLEServerBasic(EagleServerBase): - """Core tests that run on every server config variant.""" - - extra_args = ["--chunked-prefill-size", 128, "--max-running-requests", 8] - - @classmethod - def setUpClass(cls): - with envs.SGLANG_ENABLE_SPEC_V2.override(False): - super().setUpClass() - - # FIXME(lsyin): move the test methods to kits - 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_gsm8k(self): - requests.get(self.base_url + "/flush_cache") - - args = SimpleNamespace( - base_url=self.base_url, - model=self.target_model, - eval_name="gsm8k", - api="completion", - max_tokens=512, - num_examples=200, - num_threads=128, - ) - - metrics = run_eval(args) - print(f"{metrics=}") - self.assertGreater(metrics["score"], 0.20) - - server_info = requests.get(self.base_url + "/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.47) - - # Wait a little bit so that the memory check happens. - time.sleep(4) - - -class TestEAGLEServerAdditional(TestEAGLEServerBasic): - spec_topk = 5 - spec_steps = 8 - spec_tokens = 64 - extra_args = [ - "--max-running-requests", - 8, - "--cuda-graph-max-bs", - 5, - "--attention-backend", - "fa3", - "--page-size", - 256, - "--dtype", - "float16", - ] - - def test_radix_attention(self): - run_radix_attention_test(self.base_url) - self.assertIsNone(self.process.poll()) - - def test_max_token_one(self): - requests.get(self.base_url + "/flush_cache") - - args = SimpleNamespace( - base_url=self.base_url, - model=self.target_model, - eval_name="gsm8k", - api="completion", - max_tokens=1, - num_examples=200, - num_threads=128, - ) - - metrics = run_eval(args) - self.assertGreater(metrics["output_throughput"], 50) - - 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, - }, - ) - 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 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_TARGET_MODEL_EAGLE, - "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(TestEAGLEServerBasic): - extra_args = [ - "--chunked-prefill-size=128", - "--max-running-requests=64", - "--max-total-tokens=4500", # Set a smaller KV cache to trigger retract more easily - ] - - @classmethod - def setUpClass(cls): - # These config helps find a leak. - with envs.SGLANG_TEST_RETRACT.override(True): - super().setUpClass() - - -class TestEAGLEServerTriton(TestEAGLEServerBasic): - extra_args = ["--attention-backend=triton", "--max-running-requests=8"] - - -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", - ] - - @classmethod - def setUpClass(cls): - # Runtime check only supported for topk=1, and can help to find a leak. - with envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY.override(1): - super().setUpClass() - - -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", - ] - - -class TestEAGLEAbortAll(AbortAllMixin, EagleServerBase): - abort_all_max_new_tokens = 4000 - extra_args = ["--max-running-requests=8"] - - -class TestEAGLEWaitingTimeout(WaitingTimeoutMixin, EagleServerBase): - extra_args = ["--max-running-requests=1"] - - @classmethod - def setUpClass(cls): - with envs.SGLANG_REQ_WAITING_TIMEOUT.override(0.001): - super().setUpClass() - - -class TestEAGLERunningTimeout(RunningTimeoutTwoWaveMixin, EagleServerBase): - # Regression test for https://github.com/sgl-project/sglang/pull/18760 - extra_args = ["--max-running-requests=16"] - - @classmethod - def setUpClass(cls): - with envs.SGLANG_REQ_RUNNING_TIMEOUT.override(3): - super().setUpClass() - - -if __name__ == "__main__": - unittest.main() diff --git a/test/registered/spec/eagle/test_eagle_infer_beta.py b/test/registered/spec/eagle/test_eagle_infer_beta.py deleted file mode 100644 index e0b50b0a8..000000000 --- a/test/registered/spec/eagle/test_eagle_infer_beta.py +++ /dev/null @@ -1,281 +0,0 @@ -import unittest -from types import SimpleNamespace - -import numpy as np -import requests - -from sglang.srt.environ import envs -from sglang.srt.utils import kill_process_tree -from sglang.test.ci.ci_register import register_cuda_ci -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.run_eval import run_eval -from sglang.test.test_utils import ( - DEFAULT_DRAFT_MODEL_EAGLE3, - DEFAULT_TARGET_MODEL_EAGLE3, - DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, - DEFAULT_URL_FOR_TEST, - CustomTestCase, - popen_launch_server, -) - -register_cuda_ci(est_time=369, stage="base-b", runner_config="1-gpu-small") - - -class TestEagle3ServerBase(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_TARGET_MODEL_EAGLE3 - draft_model = DEFAULT_DRAFT_MODEL_EAGLE3 - - @classmethod - def setUpClass(cls): - cls.base_url = DEFAULT_URL_FOR_TEST - launch_args = [ - "--trust-remote-code", - "--dtype=float16", - "--chunked-prefill-size", - "1024", - "--attention-backend", - cls.attention_backend, - "--speculative-algorithm", - "EAGLE3", - "--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_STRICT_MEM_CHECK_DURING_BUSY.override(1), - envs.SGLANG_ENABLE_ASYNC_ASSERT.override(True), - envs.SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN.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( - base_url=self.base_url, - model=self.model, - eval_name="gsm8k", - api="completion", - max_tokens=512, - num_examples=1000, - num_threads=128, - ) - metrics = run_eval(args) - print(f"TestEagle3LargeBS -- {metrics=}") - self.assertGreater(metrics["score"], 0.7) - assert self.process.poll() is None - - def test_logprob_spec_v2_match(self): - """Verify spec v2 decode logprobs match prefill scoring logprobs. - - Generate tokens with spec v2, then score the same sequence via - prefill-only (no speculation). The two sets of logprobs should be - close, validating that spec v2 computes logprobs correctly. - - Runs two rounds with different prompts to catch state-dependent bugs. - """ - 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)) - - # Check per-token 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}] prompt={prompt!r} " - f"logprob max_diff={max_diff:.6f}" - ) - print(f"[round {round_idx}] decode_vals[-5:]={decode_vals[-5:]}") - print(f"[round {round_idx}] score_vals[-5:]={score_vals[-5:]}") - self.assertLess(max_diff, 0.255) - - # Check top-k logprobs - 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) - - # Check token_ids_logprob - 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: get_token_ids_logprobs_raw crashes on ragged token_ids_logprob lists. - - Sends concurrent requests with different-length token_ids_logprob lists - so they land in the same batch. torch.tensor() on ragged input will crash. - """ - import concurrent.futures - - 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}") - - def test_penalty(self): - """Verify spec v2 handles penalty parameters without crashing.""" - import concurrent.futures - - args = [ - {"max_new_tokens": 32}, - {"max_new_tokens": 16, "frequency_penalty": 2}, - {"max_new_tokens": 48, "presence_penalty": 1}, - {"max_new_tokens": 8, "frequency_penalty": 0.4, "presence_penalty": 0.8}, - {"max_new_tokens": 64, "frequency_penalty": -0.5, "presence_penalty": 0.3}, - {"max_new_tokens": 24, "min_new_tokens": 8, "frequency_penalty": 0.4}, - {"max_new_tokens": 32, "repetition_penalty": 1.5}, - ] - - def run_decode(sampling_params): - response = requests.post( - self.base_url + "/generate", - json={ - "text": "The capital of France is", - "sampling_params": sampling_params, - }, - ) - self.assertEqual(response.status_code, 200) - res = response.json() - self.assertIn("text", res, f"Server error: {res}") - self.assertIsInstance( - res["text"], - str, - f"Expected 'text' to be str, got {type(res['text']).__name__}: {res}", - ) - - with concurrent.futures.ThreadPoolExecutor(max_workers=8) as pool: - list(pool.map(run_decode, args * 3)) - assert self.process.poll() is None - - -class TestEagle3ServerPage(TestEagle3ServerBase): - other_launch_args = ["--page-size", "64"] - - -if __name__ == "__main__": - unittest.main() diff --git a/test/registered/spec/eagle/test_spec_eagle.py b/test/registered/spec/eagle/test_spec_eagle.py new file mode 100644 index 000000000..9d603b5ec --- /dev/null +++ b/test/registered/spec/eagle/test_spec_eagle.py @@ -0,0 +1,51 @@ +"""EAGLE3 spec-decoding core: overlap (spec v2) x no-overlap (spec v1) matrix, +same standard config (topk=1, page_size=1), only ``disable_overlap`` differs. +flashinfer is pinned (the 5090 default) so a default-selection change can't +silently alter what this exercises. +""" + +import unittest + +from sglang.srt.environ import envs +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.kits.matched_stop_kit import MatchedStopMixin +from sglang.test.kits.spec_server_kits import ( + SpecAccuracyKit, + SpecCorrectnessKit, + SpecFeatureKit, + SpecLogprobKit, + SpecPenaltyKit, +) +from sglang.test.server_fixtures.spec_eagle_fixture import Eagle3Base + +register_cuda_ci(est_time=480, stage="base-b", runner_config="1-gpu-small") + +_KITS = ( + SpecCorrectnessKit, + SpecAccuracyKit, + SpecLogprobKit, + SpecPenaltyKit, + SpecFeatureKit, + MatchedStopMixin, +) + + +class _Core(Eagle3Base): + # Busy-time pool accounting check (topk=1 only). + env_overrides = ((envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY, 1),) + + +class TestEagle3Overlap(_Core, *_KITS): + """Spec v2 (overlap scheduler on).""" + + disable_overlap = False + + +class TestEagle3NoOverlap(_Core, *_KITS): + """Spec v1 (overlap scheduler off).""" + + disable_overlap = True + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/spec/eagle/test_spec_eagle_fa3.py b/test/registered/spec/eagle/test_spec_eagle_fa3.py new file mode 100644 index 000000000..7d00c2bfc --- /dev/null +++ b/test/registered/spec/eagle/test_spec_eagle_fa3.py @@ -0,0 +1,49 @@ +"""fa3 attention backend -- Hopper-only (FlashAttention-3 is sm_90). + +fa3 is the real H200 default for MHA spec at topk=1, so this also covers the +"what an H200 user actually runs" path. Requires the large (Hopper) runner. +""" + +import unittest + +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.kits.spec_server_kits import ( + SpecAccuracyKit, + SpecCorrectnessKit, + SpecFeatureKit, + SpecLogprobKit, + SpecPenaltyKit, + SpecPerfKit, +) +from sglang.test.server_fixtures.spec_eagle_fixture import Eagle3Base, EagleLlama2Base + +register_cuda_ci(est_time=600, stage="base-b", runner_config="1-gpu-large") + + +class TestEagle3Fa3(Eagle3Base, SpecCorrectnessKit, SpecAccuracyKit, SpecLogprobKit): + """EAGLE3 spec v2 topk=1 on fa3 (the H200 default backend).""" + + attention_backend = "fa3" + disable_overlap = False + + +class TestEagleLlama2Fa3Page256( + EagleLlama2Base, + SpecAccuracyKit, + SpecLogprobKit, + SpecPenaltyKit, + SpecPerfKit, + SpecFeatureKit, +): + """EAGLE/Llama-2 topk=5 tree on fa3 + page_size=256 (spec v1).""" + + spec_topk = 5 + spec_steps = 8 + attention_backend = "fa3" + page_size = 256 + chunked_prefill_size = 4096 # must be divisible by page_size (256) + cuda_graph_max_bs = 5 + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/spec/eagle/test_spec_eagle_page.py b/test/registered/spec/eagle/test_spec_eagle_page.py new file mode 100644 index 000000000..dc415db9c --- /dev/null +++ b/test/registered/spec/eagle/test_spec_eagle_page.py @@ -0,0 +1,47 @@ +"""page_size > 1 variants (flashinfer). + +EAGLE3 page64 (spec v2) + EAGLE/Llama-2 page4 (topk1 and topk8, spec v1). +Runs on the cheap (5090) runner. +""" + +import unittest + +from sglang.srt.environ import envs +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.kits.spec_server_kits import ( + SpecAccuracyKit, + SpecFeatureKit, + SpecLogprobKit, +) +from sglang.test.server_fixtures.spec_eagle_fixture import Eagle3Base, EagleLlama2Base + +register_cuda_ci(est_time=540, stage="base-b", runner_config="1-gpu-small") + + +class TestEagle3Page64(Eagle3Base, SpecAccuracyKit, SpecLogprobKit, SpecFeatureKit): + """EAGLE3 spec v2, page_size=64 (flashinfer): + logprob losslessness.""" + + page_size = 64 + disable_overlap = False + + +class TestEagleLlama2Page4Topk1(EagleLlama2Base, SpecAccuracyKit, SpecFeatureKit): + """Llama-2 topk=1 + page_size=4; busy-time pool check (topk=1 only).""" + + spec_topk = 1 + spec_tokens = 6 + page_size = 4 + env_overrides = ( + (envs.SGLANG_ENABLE_SPEC_V2, False), + (envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY, 1), + ) + + +class TestEagleLlama2Page4Topk8(EagleLlama2Base, SpecAccuracyKit, SpecFeatureKit): + """Llama-2 topk>1 tree + page_size=4 (spec v1).""" + + page_size = 4 + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/spec/eagle/test_spec_eagle_parity.py b/test/registered/spec/eagle/test_spec_eagle_parity.py new file mode 100644 index 000000000..f20cb9c29 --- /dev/null +++ b/test/registered/spec/eagle/test_spec_eagle_parity.py @@ -0,0 +1,28 @@ +"""Lossless output parity: spec-decode greedy output == a non-spec reference. + +The reference is a separate non-spec server, launched and torn down BEFORE the +spec server (sequential -- one model resident at a time; see SpecParityKit). +""" + +import unittest + +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.kits.spec_server_kits import SpecParityKit +from sglang.test.server_fixtures.spec_eagle_fixture import Eagle3Base + +register_cuda_ci(est_time=360, stage="base-b", runner_config="1-gpu-large") + + +class TestEagle3Parity(SpecParityKit, Eagle3Base): + """EAGLE3 spec v2 (flashinfer) greedy output == non-spec reference. + + SpecParityKit is first so its setUpClass runs the reference server (and tears + it down) before the fixture launches the spec server -- sequential, one model + at a time. + """ + + disable_overlap = False + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/spec/eagle/test_spec_eagle_stress.py b/test/registered/spec/eagle/test_spec_eagle_stress.py new file mode 100644 index 000000000..eff30e1db --- /dev/null +++ b/test/registered/spec/eagle/test_spec_eagle_stress.py @@ -0,0 +1,65 @@ +"""Perf + stress: throughput, retract-under-pressure, abort storms, timeouts. + +These need memory headroom / measure load behavior, so they run on the large +(Hopper) runner. +""" + +import unittest + +from sglang.srt.environ import envs +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.kits.abort_timeout_kit import ( + AbortAllMixin, + RunningTimeoutTwoWaveMixin, + WaitingTimeoutMixin, +) +from sglang.test.kits.spec_server_kits import ( + SpecAccuracyKit, + SpecFeatureKit, + SpecPerfKit, +) +from sglang.test.server_fixtures.spec_eagle_fixture import Eagle3Base, EagleLlama2Base + +register_cuda_ci(est_time=600, stage="base-b", runner_config="1-gpu-large") + + +class TestEagle3Perf(Eagle3Base, SpecPerfKit): + """Decode throughput (max_new_tokens=1) on EAGLE3 spec v2.""" + + disable_overlap = False + + +class TestEagleLlama2Retract(EagleLlama2Base, SpecAccuracyKit, SpecFeatureKit): + """Retract under a small KV budget; must not leak.""" + + max_running_requests = 64 + extra_args = ("--max-total-tokens", 4500) # small KV to trigger retract + env_overrides = ( + (envs.SGLANG_ENABLE_SPEC_V2, False), + (envs.SGLANG_TEST_RETRACT, True), + ) + + +class TestEagleLlama2AbortAll(EagleLlama2Base, AbortAllMixin): + abort_all_max_new_tokens = 4000 + + +class TestEagleLlama2WaitingTimeout(EagleLlama2Base, WaitingTimeoutMixin): + max_running_requests = 1 + env_overrides = ( + (envs.SGLANG_ENABLE_SPEC_V2, False), + (envs.SGLANG_REQ_WAITING_TIMEOUT, 0.001), + ) + + +class TestEagleLlama2RunningTimeout(EagleLlama2Base, RunningTimeoutTwoWaveMixin): + # Regression: https://github.com/sgl-project/sglang/pull/18760 + max_running_requests = 16 + env_overrides = ( + (envs.SGLANG_ENABLE_SPEC_V2, False), + (envs.SGLANG_REQ_RUNNING_TIMEOUT, 3), + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/spec/eagle/test_spec_eagle_topk.py b/test/registered/spec/eagle/test_spec_eagle_topk.py new file mode 100644 index 000000000..2b3914bba --- /dev/null +++ b/test/registered/spec/eagle/test_spec_eagle_topk.py @@ -0,0 +1,67 @@ +"""topk > 1 tree drafting (EAGLE3 topk16 + EAGLE/Llama-2 topk8). + +topk > 1 always routes to spec v1; flashinfer is pinned (topk > 1 can't use fa3). +Runs on the cheap (5090) runner -- functional sanity only, no perf/stress. +""" + +import unittest + +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.kits.spec_server_kits import ( + SpecAccuracyKit, + SpecCorrectnessKit, + SpecFeatureKit, + SpecLogprobKit, + SpecPenaltyKit, +) +from sglang.test.server_fixtures.spec_eagle_fixture import Eagle3Base, EagleLlama2Base + +register_cuda_ci(est_time=840, stage="base-b", runner_config="1-gpu-small") + + +class TestEagle3Topk16(Eagle3Base, SpecCorrectnessKit, SpecAccuracyKit, SpecLogprobKit): + """EAGLE3 topk=16 tree (spec v1): correctness + gsm8k + logprob losslessness.""" + + spec_topk = 16 + spec_tokens = 64 + disable_overlap = True # topk>1 -> spec v1 + cuda_graph_max_bs = 5 + acc_length_thres = 3.1 + batch_accept_len_thres = 1.75 + gsm8k_accept_len_thres = 2.4 # EAGLE3 topk16 gsm8k accept ~2.48 + + +class TestEagleLlama2Suite( + EagleLlama2Base, + SpecCorrectnessKit, + SpecAccuracyKit, + SpecLogprobKit, + SpecPenaltyKit, + SpecFeatureKit, +): + """EAGLE/Llama-2 topk=8 full coverage (kits listed in bases).""" + + +class TestEagleLlama2Chunked4(EagleLlama2Base, SpecCorrectnessKit): + """Correctness under tiny chunked prefill.""" + + chunked_prefill_size = 4 + + +class TestEagleLlama3TokenMap(EagleLlama2Base, SpecAccuracyKit): + """EAGLE on Llama-3-8B with a FR-Spec token map (topk=4).""" + + model = "meta-llama/Meta-Llama-3-8B-Instruct" + draft_model = "lmsys/sglang-EAGLE-LLaMA3-Instruct-8B" + spec_topk = 4 + spec_tokens = 8 + cuda_graph_max_bs = 5 + gsm8k_accept_len_thres = 2.5 # FR-Spec token map lowers accept (~2.57) + extra_args = ( + "--speculative-token-map", + "thunlp/LLaMA3-Instruct-8B-FR-Spec/freq_32768.pt", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/spec/eagle/test_spec_eagle_triton.py b/test/registered/spec/eagle/test_spec_eagle_triton.py new file mode 100644 index 000000000..d5d422c1f --- /dev/null +++ b/test/registered/spec/eagle/test_spec_eagle_triton.py @@ -0,0 +1,47 @@ +"""triton attention backend (EAGLE3 spec v2 + EAGLE/Llama-2 spec v1). + +triton runs everywhere, so this stays on the cheap (5090) runner. +""" + +import unittest + +from sglang.srt.environ import envs +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.kits.matched_stop_kit import MatchedStopMixin +from sglang.test.kits.spec_server_kits import ( + SpecAccuracyKit, + SpecFeatureKit, + SpecLogprobKit, + SpecPenaltyKit, +) +from sglang.test.server_fixtures.spec_eagle_fixture import Eagle3Base, EagleLlama2Base + +register_cuda_ci(est_time=480, stage="base-b", runner_config="1-gpu-small") + + +class TestEagle3Triton( + Eagle3Base, + MatchedStopMixin, + SpecAccuracyKit, + SpecLogprobKit, + SpecPenaltyKit, + SpecFeatureKit, +): + """EAGLE3 spec v2 on triton (kits listed in bases).""" + + attention_backend = "triton" + max_running_requests = 64 + cuda_graph_max_bs = 64 + gsm8k_num_examples = 1000 + gsm8k_check_accept_len = False + env_overrides = ((envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY, 1),) + + +class TestEagleLlama2Triton(EagleLlama2Base, SpecAccuracyKit, SpecFeatureKit): + """EAGLE/Llama-2 topk=8 on triton (spec v1).""" + + attention_backend = "triton" + + +if __name__ == "__main__": + unittest.main()