388 lines
11 KiB
Python
388 lines
11 KiB
Python
import os
|
|
import re
|
|
import unittest
|
|
from collections import defaultdict
|
|
from dataclasses import dataclass
|
|
from types import SimpleNamespace
|
|
from typing import List
|
|
|
|
import requests
|
|
import torch
|
|
import torch.multiprocessing as mp
|
|
|
|
from sglang.bench_serving import run_benchmark
|
|
from sglang.srt.environ import envs
|
|
from sglang.srt.managers.prefill_delayer import PrefillDelayer
|
|
from sglang.srt.utils import kill_process_tree
|
|
from sglang.test.run_eval import run_eval
|
|
from sglang.test.test_utils import (
|
|
DEFAULT_MLA_MODEL_NAME_FOR_TEST,
|
|
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
|
DEFAULT_URL_FOR_TEST,
|
|
CustomTestCase,
|
|
get_benchmark_args,
|
|
popen_launch_server,
|
|
)
|
|
|
|
WORLD_SIZE = os.environ.get("SGLANG_TEST_WORLD_SIZE", "8")
|
|
|
|
# ============================ Unit Tests ============================
|
|
|
|
|
|
@dataclass
|
|
class NegotiateCall:
|
|
prefillable: List[bool]
|
|
|
|
|
|
@dataclass
|
|
class NegotiateTestCase:
|
|
name: str
|
|
max_delay_passes: int
|
|
calls: List[NegotiateCall]
|
|
expected_allow: bool
|
|
expected_reason: str
|
|
|
|
|
|
def _run_negotiate_test(rank, world_size, test_cases, results_queue, port):
|
|
torch.distributed.init_process_group(
|
|
backend="gloo",
|
|
init_method=f"tcp://127.0.0.1:{port}",
|
|
world_size=world_size,
|
|
rank=rank,
|
|
)
|
|
cpu_group = torch.distributed.new_group(backend="gloo")
|
|
|
|
for case in test_cases:
|
|
delayer = PrefillDelayer(
|
|
dp_size=world_size,
|
|
attn_tp_size=1,
|
|
cpu_group=cpu_group,
|
|
server_args=SimpleNamespace(
|
|
enable_dp_attention=True,
|
|
disaggregation_mode="null",
|
|
disable_overlap_schedule=False,
|
|
),
|
|
max_delay_passes=case.max_delay_passes,
|
|
)
|
|
|
|
for call in case.calls:
|
|
result = delayer._negotiate_should_allow_prefill(
|
|
local_prefillable=call.prefillable[rank],
|
|
)
|
|
|
|
results_queue.put((rank, case.name, result.output_allow, result.output_reason))
|
|
|
|
torch.distributed.destroy_process_group()
|
|
|
|
|
|
_NEGOTIATE_TEST_CASES = [
|
|
NegotiateTestCase(
|
|
name="all_prefillable",
|
|
max_delay_passes=100,
|
|
calls=[
|
|
NegotiateCall(prefillable=[True, True, True, True]),
|
|
],
|
|
expected_allow=True,
|
|
expected_reason="no_wait",
|
|
),
|
|
NegotiateTestCase(
|
|
name="all_prefillable_with_previous_wait",
|
|
max_delay_passes=100,
|
|
calls=[
|
|
NegotiateCall(prefillable=[True, False, True, False]),
|
|
NegotiateCall(prefillable=[True, True, True, True]),
|
|
],
|
|
expected_allow=True,
|
|
expected_reason="wait_success",
|
|
),
|
|
NegotiateTestCase(
|
|
name="none_prefillable",
|
|
max_delay_passes=100,
|
|
calls=[
|
|
NegotiateCall(prefillable=[False, False, False, False]),
|
|
],
|
|
expected_allow=True,
|
|
expected_reason="",
|
|
),
|
|
NegotiateTestCase(
|
|
name="mixed_delay",
|
|
max_delay_passes=100,
|
|
calls=[
|
|
NegotiateCall(prefillable=[True, False, True, False]),
|
|
],
|
|
expected_allow=False,
|
|
expected_reason="delay",
|
|
),
|
|
NegotiateTestCase(
|
|
name="mixed_timeout",
|
|
max_delay_passes=3,
|
|
calls=[
|
|
NegotiateCall(prefillable=[True, False, True, False]),
|
|
NegotiateCall(prefillable=[True, False, True, False]),
|
|
NegotiateCall(prefillable=[True, False, True, False]),
|
|
],
|
|
expected_allow=True,
|
|
expected_reason="wait_timeout",
|
|
),
|
|
]
|
|
|
|
|
|
class TestPrefillDelayerNegotiate(unittest.TestCase):
|
|
def test_negotiate(self):
|
|
world_size = 4
|
|
test_cases = _NEGOTIATE_TEST_CASES
|
|
|
|
ctx = mp.get_context("spawn")
|
|
results_queue = ctx.Queue()
|
|
port = 29500 + os.getpid() % 1000
|
|
|
|
processes = []
|
|
for rank in range(world_size):
|
|
p = ctx.Process(
|
|
target=_run_negotiate_test,
|
|
args=(rank, world_size, test_cases, results_queue, port),
|
|
)
|
|
p.start()
|
|
processes.append(p)
|
|
|
|
for p in processes:
|
|
p.join()
|
|
|
|
results = defaultdict(dict)
|
|
for _ in range(world_size * len(test_cases)):
|
|
rank, case_name, output_allow, output_reason = results_queue.get()
|
|
results[case_name][rank] = (output_allow, output_reason)
|
|
|
|
for case in test_cases:
|
|
for rank in range(world_size):
|
|
output_allow, output_reason = results[case.name][rank]
|
|
self.assertEqual(
|
|
(output_allow, output_reason),
|
|
(case.expected_allow, case.expected_reason),
|
|
f"Case {case.name} rank {rank}",
|
|
)
|
|
|
|
|
|
# ============================ E2E Tests ============================
|
|
|
|
|
|
class TestPrefillDelayerThroughputOnlineServing(CustomTestCase):
|
|
def test_throughput_comparison(self):
|
|
_run_throughput_comparison(
|
|
self,
|
|
test_name="online_serving",
|
|
other_launch_args=[
|
|
"--schedule-policy",
|
|
"lpm",
|
|
],
|
|
other_benchmark_args=dict(
|
|
num_prompts=500,
|
|
random_input_len=30000,
|
|
random_output_len=256,
|
|
request_rate=32,
|
|
),
|
|
min_improvement_pct=5,
|
|
)
|
|
|
|
|
|
class TestPrefillDelayerThroughputOfflineGen(CustomTestCase):
|
|
def test_throughput_comparison(self):
|
|
_run_throughput_comparison(
|
|
self,
|
|
test_name="offline_gen",
|
|
other_launch_args=["--max-total-tokens", "200000"],
|
|
other_benchmark_args=dict(
|
|
num_prompts=800,
|
|
random_input_len=30000,
|
|
random_output_len=500,
|
|
),
|
|
min_improvement_pct=20,
|
|
)
|
|
|
|
|
|
def _run_throughput_comparison(
|
|
test_case,
|
|
test_name: str,
|
|
other_launch_args,
|
|
other_benchmark_args,
|
|
min_improvement_pct: float,
|
|
):
|
|
common_kwargs = dict(
|
|
debug_name=test_name,
|
|
other_launch_args=other_launch_args,
|
|
other_benchmark_args=other_benchmark_args,
|
|
)
|
|
res_enabled = _run_throughput_test(prefill_delayer=True, **common_kwargs)
|
|
res_disabled = _run_throughput_test(prefill_delayer=False, **common_kwargs)
|
|
|
|
_assert_throughput_improvement(
|
|
test_case,
|
|
test_name=test_name,
|
|
res_enabled=res_enabled,
|
|
res_disabled=res_disabled,
|
|
min_improvement_pct=min_improvement_pct,
|
|
)
|
|
|
|
|
|
def _run_throughput_test(
|
|
debug_name: str,
|
|
prefill_delayer: bool,
|
|
other_launch_args,
|
|
other_benchmark_args,
|
|
):
|
|
model = "Qwen/Qwen3-0.6B"
|
|
base_url = DEFAULT_URL_FOR_TEST
|
|
|
|
process = _launch_server(
|
|
model=model,
|
|
base_url=base_url,
|
|
other_args=other_launch_args,
|
|
prefill_delayer=prefill_delayer,
|
|
)
|
|
|
|
try:
|
|
args = get_benchmark_args(
|
|
base_url=base_url,
|
|
dataset_name="random",
|
|
tokenizer=model,
|
|
**other_benchmark_args,
|
|
)
|
|
res = run_benchmark(args)
|
|
_print_prefill_delayer_metrics(base_url, expect_metrics=prefill_delayer)
|
|
finally:
|
|
kill_process_tree(process.pid)
|
|
|
|
print(f"=== {debug_name} ({prefill_delayer=}) ===")
|
|
res["total_throughput"] = res["input_throughput"] + res["output_throughput"]
|
|
print(f"Input throughput: {res['input_throughput']:.2f} token/s")
|
|
print(f"Output throughput: {res['output_throughput']:.2f} token/s")
|
|
print(f"Total throughput: {res['total_throughput']:.2f} token/s")
|
|
|
|
return res
|
|
|
|
|
|
def _assert_throughput_improvement(
|
|
test_case,
|
|
test_name: str,
|
|
res_enabled: dict,
|
|
res_disabled: dict,
|
|
min_improvement_pct: float,
|
|
):
|
|
test_case.assertEqual(
|
|
WORLD_SIZE,
|
|
"8",
|
|
f"This test requires 8 GPUs to properly measure throughput improvement, got {WORLD_SIZE}",
|
|
)
|
|
|
|
enabled = res_enabled["total_throughput"]
|
|
disabled = res_disabled["total_throughput"]
|
|
improvement_pct = (enabled - disabled) / disabled * 100
|
|
|
|
print(f"\n=== {test_name} Throughput Comparison ===")
|
|
print(
|
|
f"Total: enabled={enabled:.2f}, disabled={disabled:.2f}, improvement={improvement_pct:.2f}%"
|
|
)
|
|
|
|
test_case.assertGreaterEqual(
|
|
improvement_pct,
|
|
min_improvement_pct,
|
|
f"{test_name}: Throughput improvement ({improvement_pct:.2f}%) < {min_improvement_pct}%",
|
|
)
|
|
|
|
|
|
class TestPrefillDelayerAccuracy(CustomTestCase):
|
|
def test_1_mgsm_en_has_prefill_delayer(self):
|
|
self._run_accuracy_test(prefill_delayer=True)
|
|
|
|
def test_2_mgsm_en_no_prefill_delayer(self):
|
|
self._run_accuracy_test(prefill_delayer=False)
|
|
|
|
def _run_accuracy_test(self, prefill_delayer: bool):
|
|
model = DEFAULT_MLA_MODEL_NAME_FOR_TEST
|
|
base_url = DEFAULT_URL_FOR_TEST
|
|
process = _launch_server(
|
|
prefill_delayer=prefill_delayer,
|
|
model=model,
|
|
base_url=base_url,
|
|
other_args=[
|
|
"--schedule-policy",
|
|
"lpm",
|
|
"--max-total-tokens",
|
|
"4096",
|
|
],
|
|
)
|
|
try:
|
|
args = SimpleNamespace(
|
|
base_url=base_url,
|
|
model=model,
|
|
eval_name="mgsm_en",
|
|
num_examples=None,
|
|
num_threads=1024,
|
|
)
|
|
metrics = run_eval(args)
|
|
print(f"=== mgsm_en ({prefill_delayer=}) ===")
|
|
print(f"{metrics=}")
|
|
self.assertGreater(metrics["score"], 0.87)
|
|
finally:
|
|
kill_process_tree(process.pid)
|
|
|
|
|
|
def _launch_server(
|
|
*,
|
|
model,
|
|
base_url,
|
|
prefill_delayer: bool,
|
|
other_args,
|
|
max_delay_passes: int = 100,
|
|
):
|
|
os.environ["SGLANG_PREFILL_DELAYER_DEBUG_LOG"] = "1"
|
|
|
|
with envs.SGLANG_SCHEDULER_DECREASE_PREFILL_IDLE.override(
|
|
prefill_delayer
|
|
), envs.SGLANG_PREFILL_DELAYER_MAX_DELAY_PASSES.override(max_delay_passes):
|
|
return popen_launch_server(
|
|
model,
|
|
base_url,
|
|
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
|
other_args=[
|
|
"--trust-remote-code",
|
|
"--tp",
|
|
WORLD_SIZE,
|
|
"--enable-dp-attention",
|
|
"--dp",
|
|
WORLD_SIZE,
|
|
"--chunked-prefill-size",
|
|
"131072",
|
|
"--mem-fraction-static",
|
|
"0.6",
|
|
"--enable-metrics",
|
|
*(other_args or []),
|
|
],
|
|
)
|
|
|
|
|
|
def _print_prefill_delayer_metrics(base_url: str, expect_metrics: bool) -> str:
|
|
metrics_response = requests.get(f"{base_url}/metrics")
|
|
assert metrics_response.status_code == 200
|
|
metrics_text = metrics_response.text
|
|
prefill_delayer_metrics = [
|
|
line for line in metrics_text.split("\n") if "prefill_delayer" in line
|
|
]
|
|
print("=== PrefillDelayer Metrics ===")
|
|
for line in prefill_delayer_metrics:
|
|
print(line)
|
|
if expect_metrics:
|
|
assert "sglang:prefill_delayer_wait_forward_passes" in metrics_text
|
|
assert "sglang:prefill_delayer_wait_seconds" in metrics_text
|
|
assert "sglang:prefill_delayer_outcomes_total" in metrics_text
|
|
return metrics_text
|
|
|
|
|
|
def _sum_prometheus_metric_values(metrics_text: str, label_value: str) -> int:
|
|
matches = re.findall(rf'{label_value}".*?\}} (\d+)', metrics_text)
|
|
return sum(int(m) for m in matches)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|