chore: vendor sglang v0.5.10 snapshot

This commit is contained in:
2026-04-24 12:29:36 +00:00
parent 78f0d15221
commit bded08301f
4308 changed files with 1200894 additions and 2 deletions

View File

@@ -0,0 +1,257 @@
import os
import random
import tempfile
import unittest
from typing import Dict
import requests
from sglang.benchmark.utils import get_tokenizer
from sglang.test.server_fixtures.disaggregation_fixture import (
PDDisaggregationServerBase,
)
from sglang.test.test_utils import (
DEFAULT_MODEL_NAME_FOR_TEST,
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
popen_launch_pd_server,
)
class DisaggregationHiCacheBase(PDDisaggregationServerBase):
"""Base class for disaggregation with HiCache tests"""
@classmethod
def setUpClass(cls):
super(DisaggregationHiCacheBase, cls).setUpClass()
cls.model = DEFAULT_MODEL_NAME_FOR_TEST
cls.tokenizer = get_tokenizer(cls.model)
cls.temp_dir = tempfile.mkdtemp()
cls.start_prefill()
cls.start_decode()
# Block until both
cls.wait_server_ready(cls.prefill_url + "/health", process=cls.process_prefill)
cls.wait_server_ready(cls.decode_url + "/health", process=cls.process_decode)
cls.launch_lb()
@classmethod
def start_prefill(cls):
# Prefill with HiCache enabled
prefill_args = [
"--trust-remote-code",
"--disaggregation-mode",
"prefill",
"--tp-size",
"1",
"--page-size",
"64",
"--enable-hierarchical-cache",
"--hicache-ratio",
"1.2",
"--hicache-size",
"0",
"--hicache-write-policy",
"write_through",
"--hicache-storage-backend",
"file",
"--hicache-storage-prefetch-policy",
"wait_complete",
"--mem-fraction-static",
"0.8",
]
prefill_args += cls.transfer_backend + cls.rdma_devices
env = {
**os.environ,
"SGLANG_HICACHE_FILE_BACKEND_STORAGE_DIR": cls.temp_dir,
}
cls.process_prefill = popen_launch_pd_server(
cls.model,
cls.prefill_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=prefill_args,
env=env,
)
@classmethod
def start_decode(cls):
pass
def gen_prompt(self, token_num: int) -> str:
all_available_tokens = list(self.tokenizer.get_vocab().values())
selected_tokens = random.choices(all_available_tokens, k=token_num)
return self.tokenizer.decode(selected_tokens)
def send_request(
self, prompt: str, max_tokens: int = 100, temperature: float = 0.0
) -> Dict:
"""Send a generate request and return response"""
response = requests.post(
f"{self.lb_url}/generate",
json={
"text": prompt,
"sampling_params": {
"temperature": temperature,
"max_new_tokens": max_tokens,
"ignore_eos": True,
},
},
timeout=60,
)
self.assertEqual(
response.status_code,
200,
f"Request failed: {response.status_code} - {response.text}",
)
return response.json()
def trigger_offloading_and_flush(self):
"""Helper method to trigger offloading and flush cache"""
# Trigger offloading
self.send_request(self.gen_prompt(1), max_tokens=150)
# Flush device cache to force remote storage access.
res = requests.post(
f"{self.prefill_url}/flush_cache",
params={"timeout": 30},
timeout=40,
)
res.raise_for_status()
class TestDisaggregationPrefillWithHiCache(DisaggregationHiCacheBase):
"""Test disaggregation with HiCache enabled only on Prefill side"""
@classmethod
def start_decode(cls):
# Decode without HiCache offload
decode_args = [
"--trust-remote-code",
"--disaggregation-mode",
"decode",
"--tp-size",
"1",
"--page-size",
"64",
"--mem-fraction-static",
"0.8",
"--base-gpu-id",
"1",
]
decode_args += cls.transfer_backend + cls.rdma_devices
env = {
**os.environ,
"SGLANG_HICACHE_FILE_BACKEND_STORAGE_DIR": cls.temp_dir,
}
cls.process_decode = popen_launch_pd_server(
cls.model,
cls.decode_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=decode_args,
env=env,
)
def test_prefill_cache_hit(self):
"""Test that prefill cache works with repeated queries"""
repeated_prompt = self.gen_prompt(800)
# First request - should miss cache
self.send_request(repeated_prompt, max_tokens=100)
# Flush cache
self.trigger_offloading_and_flush()
# Second request - should hit cache (faster)
response2 = self.send_request(repeated_prompt, max_tokens=100)
# Assert cached tokens cnt
self.assertGreater(response2["meta_info"]["cached_tokens"], 700)
class TestDisaggregationDecodeWithHiCache(DisaggregationHiCacheBase):
"""Test disaggregation with HiCache enabled on both Prefill and Decode sides"""
@classmethod
def start_decode(cls):
# Decode with HiCache offload enabled
decode_args = [
"--trust-remote-code",
"--disaggregation-mode",
"decode",
"--tp-size",
"1",
"--page-size",
"64",
"--mem-fraction-static",
"0.8",
"--base-gpu-id",
"1",
"--disaggregation-decode-enable-offload-kvcache",
"--hicache-ratio",
"1.2",
"--hicache-size",
"0",
"--hicache-storage-backend",
"file",
"--hicache-storage-prefetch-policy",
"wait_complete",
]
decode_args += cls.transfer_backend + cls.rdma_devices
env = {
**os.environ,
"SGLANG_HICACHE_FILE_BACKEND_STORAGE_DIR": cls.temp_dir,
}
cls.process_decode = popen_launch_pd_server(
cls.model,
cls.decode_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=decode_args,
env=env,
)
def test_multi_turn_conversation_cache(self):
"""Test multi-turn conversation scenario with cache hit improvement"""
print("=== Multi-turn Conversation Cache Test ===")
# Turn 1
initial_prompt = self.gen_prompt(300)
response1 = self.send_request(initial_prompt, max_tokens=200, temperature=0.1)
current_context = initial_prompt + response1["text"]
# Turns 2-4: Continue generation based on previous context
previous_cached_tokens = 0
for turn in range(2, 5):
print(f"\nTurn {turn}: Continuing from previous context")
response = self.send_request(
current_context, max_tokens=200, temperature=0.1
)
cached_tokens = response["meta_info"]["cached_tokens"]
print(f"Turn {turn} cached tokens: {cached_tokens}")
print(f"Improvement: {cached_tokens - previous_cached_tokens} tokens")
# Assert cache improvement
self.assertGreater(
cached_tokens,
previous_cached_tokens,
f"Turn {turn} should have more cached tokens than turn {turn-1}",
)
# Update context and cached tokens for next iteration
current_context += response["text"]
previous_cached_tokens = cached_tokens
# Flush prefill cache
self.trigger_offloading_and_flush()
if __name__ == "__main__":
unittest.main()

View File

@@ -0,0 +1,214 @@
"""
Usage:
python3 -m unittest test_pp_with_hicache.TestPPWithHiCache.test_eval_accuracy
"""
import os
import subprocess
import time
import unittest
from types import SimpleNamespace
from urllib.parse import urlparse
import requests
from sglang.srt.utils import kill_process_tree
from sglang.test.run_eval import run_eval
from sglang.test.test_utils import (
DEFAULT_MODEL_NAME_FOR_TEST,
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
find_available_port,
popen_launch_server,
)
class TestPPWithHiCache(unittest.TestCase):
@classmethod
def setUpClass(cls):
cls.base_url = f"http://127.0.0.1:{find_available_port(23337)}"
parsed_url = urlparse(cls.base_url)
cls.base_host = parsed_url.hostname
cls.base_port = str(parsed_url.port)
cls.model = DEFAULT_MODEL_NAME_FOR_TEST
cls._start_mooncake_services()
server_args_dict = {
"--enable-hierarchical-cache": True,
"--mem-fraction-static": 0.6,
"--hicache-ratio": 1.2,
"--page-size": 64,
"--enable-cache-report": True,
"--hicache-storage-prefetch-policy": "wait_complete",
"--hicache-storage-backend": "mooncake",
"--tp-size": 2,
"--pp-size": 2,
"--chunked-prefill-size": 256,
"--hicache-mem-layout": "page_first",
}
final_server_args = []
for key, value in server_args_dict.items():
final_server_args.append(str(key))
if value is not True:
final_server_args.append(str(value))
env_vars = {**os.environ, **cls._mooncake_env()}
try:
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=final_server_args,
env=env_vars,
)
except Exception:
cls._stop_mooncake_services()
raise
@classmethod
def tearDownClass(cls):
if hasattr(cls, "process"):
kill_process_tree(cls.process.pid)
cls._stop_mooncake_services()
@classmethod
def _start_mooncake_services(cls):
try:
import mooncake.http_metadata_server # type: ignore # noqa: F401
except Exception as exc: # pragma: no cover - environment dependent
raise unittest.SkipTest(
f"Mooncake metadata server module unavailable: {exc}"
) from exc
cls._mooncake_master_port = find_available_port(50051)
cls._mooncake_metadata_port = find_available_port(8080)
try:
cls._mooncake_metadata_process = subprocess.Popen(
[
"python3",
"-m",
"mooncake.http_metadata_server",
"--port",
str(cls._mooncake_metadata_port),
],
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
preexec_fn=os.setsid,
)
except (FileNotFoundError, subprocess.SubprocessError) as exc:
cls._stop_mooncake_services()
raise unittest.SkipTest(
f"Could not start Mooncake metadata service: {exc}"
) from exc
try:
cls._mooncake_master_process = subprocess.Popen(
["mooncake_master", "--port", str(cls._mooncake_master_port)],
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
preexec_fn=os.setsid,
)
except (FileNotFoundError, subprocess.SubprocessError) as exc:
cls._stop_mooncake_services()
raise unittest.SkipTest(f"Could not start mooncake_master: {exc}") from exc
if not cls._wait_for_mooncake_ready():
cls._stop_mooncake_services()
raise unittest.SkipTest("Mooncake services did not become ready in time")
@classmethod
def _stop_mooncake_services(cls):
for attr in ("_mooncake_metadata_process", "_mooncake_master_process"):
proc = getattr(cls, attr, None)
if proc:
try:
os.killpg(os.getpgid(proc.pid), 9)
proc.wait(timeout=5)
except Exception:
pass
cls._mooncake_metadata_process = None
cls._mooncake_master_process = None
@classmethod
def _mooncake_env(cls):
return {
"MOONCAKE_MASTER": f"127.0.0.1:{cls._mooncake_master_port}",
"MOONCAKE_PROTOCOL": "tcp",
"MC_MS_AUTO_DISC": "0",
"MOONCAKE_DEVICE": "",
"MOONCAKE_TE_META_DATA_SERVER": f"http://127.0.0.1:{cls._mooncake_metadata_port}/metadata",
"MOONCAKE_GLOBAL_SEGMENT_SIZE": "4294967296",
"SGLANG_ENABLE_DETERMINISTIC_INFERENCE": "1",
}
@classmethod
def _wait_for_mooncake_ready(cls, timeout: int = 30) -> bool:
start_time = time.time()
while time.time() - start_time < timeout:
metadata_ready = False
master_ready = False
if (
getattr(cls, "_mooncake_metadata_process", None)
and cls._mooncake_metadata_process.poll() is None
):
try:
resp = requests.get(
f"http://127.0.0.1:{cls._mooncake_metadata_port}/metadata",
timeout=2,
)
print(resp)
metadata_ready = True
except requests.RequestException:
metadata_ready = False
if (
getattr(cls, "_mooncake_master_process", None)
and cls._mooncake_master_process.poll() is None
):
if time.time() - start_time > 3:
master_ready = True
if metadata_ready and master_ready:
return True
time.sleep(1.5)
return False
def flush_cache(self):
res = requests.post(
f"{self.base_url}/flush_cache",
params={"timeout": 30},
timeout=40,
)
res.raise_for_status()
def test_eval_accuracy(self):
args = SimpleNamespace(
base_url=self.base_url,
model=self.model,
eval_name="gsm8k",
api="completion",
max_tokens=512,
num_examples=40,
num_threads=24,
)
metrics_initial = run_eval(args)
self.assertGreater(metrics_initial["score"], 0.6)
self.flush_cache()
metrics_cached = run_eval(args)
self.assertGreater(metrics_cached["score"], 0.6)
accuracy_diff = abs(metrics_initial["score"] - metrics_cached["score"])
self.assertLess(accuracy_diff, 0.05)
if __name__ == "__main__":
unittest.main()