From f4e7a98fe58347c3b82c8b0f985e30b9f458b321 Mon Sep 17 00:00:00 2001 From: shuwenn <47200617+alphabetc1@users.noreply.github.com> Date: Wed, 3 Jun 2026 12:42:04 +0800 Subject: [PATCH] [HiCache] feat: support draft offload for mooncake (#24984) Co-authored-by: huangtingwei9988 <141888744+huangtingwei9988@users.noreply.github.com> Co-authored-by: stmatengss <11641725+stmatengss@users.noreply.github.com> --- .../sglang/srt/managers/cache_controller.py | 132 +++++++-- .../sglang/srt/mem_cache/hicache_storage.py | 3 + .../storage/mooncake_store/mooncake_store.py | 25 +- .../test/hicache_spec_storage_common.py | 273 ++++++++++++++++++ .../hicache/test_hicache_spec_file_storage.py | 244 +--------------- .../test_hicache_spec_mooncake_storage.py | 187 ++++++++++++ 6 files changed, 607 insertions(+), 257 deletions(-) create mode 100644 python/sglang/test/hicache_spec_storage_common.py create mode 100644 test/registered/hicache/test_hicache_spec_mooncake_storage.py diff --git a/python/sglang/srt/managers/cache_controller.py b/python/sglang/srt/managers/cache_controller.py index 480717aed..56057ade6 100644 --- a/python/sglang/srt/managers/cache_controller.py +++ b/python/sglang/srt/managers/cache_controller.py @@ -25,6 +25,8 @@ from sglang.srt.mem_cache.hicache_storage import ( STORAGE_BATCH_SIZE, HiCacheStorageConfig, HiCacheStorageExtraInfo, + PoolName, + PoolTransfer, ) if TYPE_CHECKING: @@ -288,6 +290,8 @@ class HiCacheController: self.has_draft = False self.mem_pool_device_draft = None self.mem_pool_host_draft = None + self.draft_page_get_func = None + self.draft_page_set_func = None # Default storage page IO functions (may be overridden by attach). self.page_get_func = self._generic_page_get @@ -529,6 +533,8 @@ class HiCacheController: self.page_get_func = self._page_get_zero_copy self.page_set_func = self._page_set_zero_copy + self._maybe_register_draft_with_storage() + # Ensure stop_event is clear before starting threads. self.storage_stop_event.clear() self._start_storage_threads() @@ -553,6 +559,8 @@ class HiCacheController: self.enable_storage = False self.page_get_func = self._generic_page_get self.page_set_func = self._generic_page_set + self.draft_page_get_func = None + self.draft_page_set_func = None raise def detach_storage_backend(self): @@ -594,6 +602,8 @@ class HiCacheController: self.enable_storage = False self.page_get_func = self._generic_page_get self.page_set_func = self._generic_page_set + self.draft_page_get_func = None + self.draft_page_set_func = None # Now it's safe to clear the stop event for future re-attach. self.storage_stop_event.clear() @@ -844,6 +854,47 @@ class HiCacheController: draft_host_pool.size, ) + # If storage is already attached, wire up the draft I/O path now. + # Otherwise this will be deferred until attach_storage_backend(). + self._maybe_register_draft_with_storage() + + def _maybe_register_draft_with_storage(self) -> None: + """Pick the draft L3 IO implementation.""" + self.draft_page_get_func = None + self.draft_page_set_func = None + if not self.has_draft or not self.enable_storage: + return + + backend = self.storage_backend_type + + # Multi-pool zero-copy backends. + if backend == "mooncake": + if self.storage_config.should_split_heads: + logger.warning( + "HiCache draft L3 disabled: should_split_heads not yet " + "supported on the mooncake v2 path." + ) + return + self.storage_backend.register_mem_host_pool_v2( + self.mem_pool_host_draft, PoolName.DRAFT + ) + self.draft_page_get_func = self._draft_page_get_v2 + self.draft_page_set_func = self._draft_page_set_v2 + return + + # TODO: support "hf3fs", "eic", "nixl", "simm" + if backend in {"hf3fs", "eic", "nixl", "simm"}: + logger.warning( + "HiCache draft L3 disabled: backend %s does not yet support " + "draft pool registration.", + backend, + ) + return + + # Generic backends. + self.draft_page_get_func = self._draft_page_get_generic + self.draft_page_set_func = self._draft_page_set_generic + def prefetch( self, request_id: str, @@ -1075,44 +1126,71 @@ class HiCacheController: ) def _draft_page_set(self, hash_values, host_indices) -> None: - """Best-effort write draft KV pages to L3 with 'd:' prefixed keys. - - TODO: support batch_set_v1 (zero-copy) for high-performance backends. - """ + """Best-effort write draft KV pages to L3 alongside the target backup.""" + if self.draft_page_set_func is None: + return try: - draft_keys = [f"d:{h}" for h in hash_values] - draft_data = [ - self.mem_pool_host_draft.get_data_page(host_indices[i * self.page_size]) - for i in range(len(draft_keys)) - ] - self.storage_backend.batch_set(draft_keys, draft_data) + self.draft_page_set_func(hash_values, host_indices) except Exception: logger.debug( "Draft L3 write failed (best-effort), skipping.", exc_info=True ) def _draft_page_get(self, hash_values, host_indices) -> None: - """Best-effort read draft KV pages from L3 with 'd:' prefixed keys. - - TODO: support batch_get_v1 (zero-copy) for high-performance backends. - """ + """Best-effort read draft KV pages from L3 (mirrors `_draft_page_set`).""" + if self.draft_page_get_func is None: + return try: - draft_keys = [f"d:{h}" for h in hash_values] - draft_dummy = [ - self.mem_pool_host_draft.get_dummy_flat_data_page() for _ in draft_keys - ] - draft_pages = self.storage_backend.batch_get(draft_keys, draft_dummy) - if draft_pages is None: - return - - for i, p in enumerate(draft_pages): - if p is not None: - self.mem_pool_host_draft.set_from_flat_data_page( - host_indices[i * self.page_size], p - ) + self.draft_page_get_func(hash_values, host_indices) except Exception: logger.debug("Draft L3 read failed (best-effort), skipping.", exc_info=True) + def _draft_page_set_v2(self, hash_values, host_indices) -> None: + self.storage_backend.batch_set_v2( + [ + PoolTransfer( + name=PoolName.DRAFT, + host_indices=host_indices, + keys=list(hash_values), + ) + ] + ) + + def _draft_page_get_v2(self, hash_values, host_indices) -> None: + self.storage_backend.batch_get_v2( + [ + PoolTransfer( + name=PoolName.DRAFT, + host_indices=host_indices, + keys=list(hash_values), + ) + ] + ) + + def _draft_page_set_generic(self, hash_values, host_indices) -> None: + # `{hash}.draft` mirrors HiCacheStorage._get_component_key's + # `{key}.{pool_name}` convention so target/draft pages never collide. + draft_keys = [f"{h}.{PoolName.DRAFT}" for h in hash_values] + draft_data = [ + self.mem_pool_host_draft.get_data_page(host_indices[i * self.page_size]) + for i in range(len(draft_keys)) + ] + self.storage_backend.batch_set(draft_keys, draft_data) + + def _draft_page_get_generic(self, hash_values, host_indices) -> None: + draft_keys = [f"{h}.{PoolName.DRAFT}" for h in hash_values] + draft_dummy = [ + self.mem_pool_host_draft.get_dummy_flat_data_page() for _ in draft_keys + ] + draft_pages = self.storage_backend.batch_get(draft_keys, draft_dummy) + if draft_pages is None: + return + for i, p in enumerate(draft_pages): + if p is not None: + self.mem_pool_host_draft.set_from_flat_data_page( + host_indices[i * self.page_size], p + ) + # Backup batch by batch def _page_backup(self, operation): # Backup batch by batch diff --git a/python/sglang/srt/mem_cache/hicache_storage.py b/python/sglang/srt/mem_cache/hicache_storage.py index effbdde8e..a982cb7c8 100644 --- a/python/sglang/srt/mem_cache/hicache_storage.py +++ b/python/sglang/srt/mem_cache/hicache_storage.py @@ -68,6 +68,9 @@ class PoolName(str, Enum): DEEPSEEK_V4_C4_INDEXER_STATE = "deepseek_v4_c4_indexer_state" DEEPSEEK_V4_C128_STATE = "deepseek_v4_c128_state" + # Draft KV pool + DRAFT = "draft" + def __str__(self) -> str: return self.value diff --git a/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py b/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py index cd163ba63..6124a60fe 100644 --- a/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py +++ b/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py @@ -20,7 +20,11 @@ from sglang.srt.mem_cache.hicache_storage import ( PoolTransfer, PoolTransferResult, ) -from sglang.srt.mem_cache.memory_pool_host import HostKVCache, HostTensorAllocator +from sglang.srt.mem_cache.memory_pool_host import ( + HostKVCache, + HostTensorAllocator, + MLATokenToKVPoolHost, +) from sglang.srt.observability.metrics_collector import StorageMetrics DEFAULT_LOCAL_BUFFER_SIZE = 16 * 1024 * 1024 # 16 MB @@ -618,6 +622,11 @@ class MooncakeStore(HiCacheStorage, MooncakeBaseStore): # v2 here only registers additional hybrid pools. if host_pool_name == PoolName.KV: return + if host_pool_name == PoolName.DRAFT: + self.registered_pools[host_pool_name] = host_pool + super().register_buffer(host_pool.kv_buffer) + return + # Keep a name->pool mapping so batch v2 can resolve PoolTransfer.name to # the corresponding host pool implementation at runtime. self.registered_pools[host_pool_name] = host_pool @@ -652,6 +661,20 @@ class MooncakeStore(HiCacheStorage, MooncakeBaseStore): suffixes = [f"{base_suffix}_temporal"] + [ f"{base_suffix}_conv_{i}" for i in range(conv_num) ] + elif name == PoolName.DRAFT: + # Draft pool's MLA/MHA layout is independent from the target + # (e.g. EAGLE-MHA draft on top of an MLA target), so pick the + # suffix scheme from the draft pool's own class. The `_draft` + # tag is what keeps these keys from colliding with target's + # `{rank}_k` / `{rank}_k` + `{rank}_v` keys. + draft_pool = self.registered_pools.get(PoolName.DRAFT) + if isinstance(draft_pool, MLATokenToKVPoolHost): + suffixes = [f"_{self.mla_suffix}_{PoolName.DRAFT}_k"] + else: + suffixes = [ + f"_{self.mha_suffix}_{PoolName.DRAFT}_k", + f"_{self.mha_suffix}_{PoolName.DRAFT}_v", + ] key_multiplier = len(suffixes) component_keys = [ f"{page_key}{suffix}" for page_key in page_keys for suffix in suffixes diff --git a/python/sglang/test/hicache_spec_storage_common.py b/python/sglang/test/hicache_spec_storage_common.py new file mode 100644 index 000000000..9a3a5efa3 --- /dev/null +++ b/python/sglang/test/hicache_spec_storage_common.py @@ -0,0 +1,273 @@ +"""Shared helpers for HiCache storage + EAGLE3 speculative decoding tests.""" + +import json +import os +import subprocess +from typing import Dict, List + +import psutil +import requests + +from sglang.benchmark.utils import get_tokenizer +from sglang.srt.utils import kill_process_tree +from sglang.test.test_utils import ( + DEFAULT_DRAFT_MODEL_EAGLE3, + DEFAULT_TARGET_MODEL_EAGLE3, + DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + DEFAULT_URL_FOR_TEST, + find_available_port, + popen_launch_server, +) +from sglang.utils import wait_for_http_ready + + +class HiCacheSpecStorageMixin: + """Common EAGLE3 + HiCache storage loadback flow. + + Subclasses provide the storage backend, environment, and the backend-specific + wait condition before server restart. + """ + + model = DEFAULT_TARGET_MODEL_EAGLE3 + draft_model = DEFAULT_DRAFT_MODEL_EAGLE3 + + input_token_len = 1024 + max_new_tokens = 200 + first_measure_new_tokens = 128 + page_size = 64 + min_expected_accept_length = 7.0 + min_second_to_first_accept_ratio = 0.9 + + storage_backend = None + expected_storage_backend = None + + @classmethod + def _get_storage_backend_extra_config(cls): + return { + "hicache_storage_pass_prefix_keys": True, + } + + @classmethod + def _get_spec_server_env(cls) -> Dict[str, str]: + return {} + + @classmethod + def _get_spec_server_args(cls) -> List[str]: + if cls.storage_backend is None: + raise ValueError("storage_backend must be set by subclasses.") + + return [ + "--enable-hierarchical-cache", + "--enable-cache-report", + "--mem-fraction-static", + "0.3", + "--hicache-ratio", + "1.5", + "--disable-cuda-graph", + "--page-size", + str(cls.page_size), + "--hicache-storage-backend", + cls.storage_backend, + "--hicache-storage-prefetch-policy", + "wait_complete", + "--hicache-storage-backend-extra-config", + json.dumps(cls._get_storage_backend_extra_config()), + "--speculative-algorithm", + "EAGLE3", + "--speculative-draft-model-path", + cls.draft_model, + "--speculative-num-steps", + "7", + "--speculative-eagle-topk", + "1", + "--speculative-num-draft-tokens", + "8", + "--dtype", + "float16", + ] + + @classmethod + def _build_spec_prompt(cls): + cls.tokenizer = get_tokenizer(cls.model) + cls.prompt_input_ids = cls._build_long_repetitive_prompt_ids( + cls.tokenizer, cls.input_token_len + ) + + @classmethod + def _launch_spec_server(cls): + default_port = int(DEFAULT_URL_FOR_TEST.rsplit(":", 1)[1]) + cls.base_url = f"http://127.0.0.1:{find_available_port(default_port)}" + cls._build_spec_prompt() + cls.other_args = cls._get_spec_server_args() + cls.env = { + **os.environ, + "SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN": "1", + **cls._get_spec_server_env(), + } + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=cls.other_args, + env=cls.env, + ) + wait_for_http_ready( + url=f"{cls.base_url}/health", + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + process=cls.process, + ) + + @classmethod + def _restart_spec_server(cls): + cls._stop_spec_server() + cls._launch_spec_server() + + @classmethod + def _stop_spec_server(cls): + if getattr(cls, "process", None) is None: + return + + process = cls.process + try: + root = psutil.Process(process.pid) + watched_procs = [root] + root.children(recursive=True) + except psutil.NoSuchProcess: + watched_procs = [] + + try: + try: + process.terminate() + process.wait(timeout=120) + return + except subprocess.TimeoutExpired: + pass + + try: + kill_process_tree(process.pid, wait_timeout=60) + except RuntimeError: + non_zombie_procs = [] + for proc in watched_procs: + try: + if proc.is_running() and proc.status() != psutil.STATUS_ZOMBIE: + non_zombie_procs.append(proc) + except psutil.NoSuchProcess: + pass + if non_zombie_procs: + raise + finally: + cls.process = None + + @classmethod + def _encode_without_special_tokens(cls, tokenizer, text: str) -> List[int]: + return tokenizer.encode(text, add_special_tokens=False) + + @classmethod + def _build_long_repetitive_prompt_ids(cls, tokenizer, target_len: int) -> List[int]: + bos_ids = ( + [tokenizer.bos_token_id] + if getattr(tokenizer, "bos_token_id", None) is not None + else [] + ) + suffix_ids = cls._encode_without_special_tokens( + tokenizer, + "\n\nContinue the sequence with only the word apple separated by spaces.\n" + "Answer: apple apple apple apple", + ) + repeat_ids = cls._encode_without_special_tokens(tokenizer, " apple") + if not repeat_ids: + raise ValueError( + "Tokenizer produced no ids for the repetitive prompt seed." + ) + if len(bos_ids) + len(suffix_ids) >= target_len: + raise ValueError( + "Prompt suffix is too long: " + f"{len(bos_ids)=}, {len(suffix_ids)=}, {target_len=}." + ) + + prefix_len = target_len - len(bos_ids) - len(suffix_ids) + repeats = (prefix_len + len(repeat_ids) - 1) // len(repeat_ids) + prefix_ids = (repeat_ids * repeats)[:prefix_len] + prompt_ids = bos_ids + prefix_ids + suffix_ids + assert len(prompt_ids) == target_len + return prompt_ids + + def _send_long_prompt(self, max_new_tokens: int = None) -> Dict: + if max_new_tokens is None: + max_new_tokens = self.max_new_tokens + response = requests.post( + f"{self.base_url}/generate", + json={ + "input_ids": self.prompt_input_ids, + "sampling_params": { + "temperature": 0, + "max_new_tokens": max_new_tokens, + "ignore_eos": True, + }, + }, + timeout=900, + ) + self.assertEqual( + response.status_code, + 200, + f"Request failed: {response.status_code} - {response.text}", + ) + return response.json() + + def _get_spec_accept_length(self, response_json: Dict) -> float: + meta_info = response_json.get("meta_info", {}) + self.assertIn( + "spec_accept_length", + meta_info, + f"Missing spec_accept_length in meta_info: {meta_info}", + ) + return float(meta_info["spec_accept_length"]) + + def _wait_for_storage_before_restart(self): + raise NotImplementedError + + def _run_storage_loadback_keeps_spec_accept_length(self): + first = self._send_long_prompt(max_new_tokens=self.first_measure_new_tokens) + first_accept_length = self._get_spec_accept_length(first) + self.assertGreaterEqual( + first_accept_length, + self.min_expected_accept_length, + f"First prompt accept length is too low: {first_accept_length=}", + ) + + self._wait_for_storage_before_restart() + self._restart_spec_server() + + second = self._send_long_prompt() + second_accept_length = self._get_spec_accept_length(second) + second_meta = second.get("meta_info", {}) + cached_details = second_meta.get("cached_tokens_details") or {} + storage_cached_tokens = int(cached_details.get("storage", 0)) + + print( + f"{first_accept_length=:.3f}, {second_accept_length=:.3f}, " + f"{storage_cached_tokens=}, {cached_details=}" + ) + + self.assertGreaterEqual( + storage_cached_tokens, + self.input_token_len - 2 * self.page_size, + "Expected the second request to load the long prompt KV cache from " + f"{self.storage_backend} storage, got {cached_details=}", + ) + self.assertEqual( + cached_details.get("storage_backend"), + self.expected_storage_backend, + f"Expected {self.expected_storage_backend} in cache report, " + f"got {cached_details=}", + ) + self.assertGreaterEqual( + second_accept_length, + self.min_expected_accept_length, + f"Second prompt accept length is too low: {second_accept_length=}", + ) + self.assertGreaterEqual( + second_accept_length, + first_accept_length * self.min_second_to_first_accept_ratio, + f"Spec accept length dropped after {self.storage_backend}-storage " + f"loadback: {first_accept_length=:.3f}, {second_accept_length=:.3f}", + ) diff --git a/test/registered/hicache/test_hicache_spec_file_storage.py b/test/registered/hicache/test_hicache_spec_file_storage.py index 4ea14c26e..b05a63184 100644 --- a/test/registered/hicache/test_hicache_spec_file_storage.py +++ b/test/registered/hicache/test_hicache_spec_file_storage.py @@ -5,144 +5,36 @@ Usage: python3 -m pytest test/registered/hicache/test_hicache_spec_file_storage.py -v """ -import json import os import shutil import tempfile import time import unittest -from typing import Dict, List -import psutil -import requests - -from sglang.benchmark.utils import get_tokenizer -from sglang.srt.utils import is_hip, kill_process_tree +from sglang.srt.mem_cache.hicache_storage import PoolName +from sglang.srt.utils import is_hip from sglang.test.ci.ci_register import register_cuda_ci -from sglang.test.test_utils import ( - DEFAULT_DRAFT_MODEL_EAGLE3, - DEFAULT_TARGET_MODEL_EAGLE3, - DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, - DEFAULT_URL_FOR_TEST, - CustomTestCase, - find_available_port, - popen_launch_server, -) -from sglang.utils import wait_for_http_ready +from sglang.test.hicache_spec_storage_common import HiCacheSpecStorageMixin +from sglang.test.test_utils import CustomTestCase register_cuda_ci(est_time=200, stage="extra-a", runner_config="1-gpu-large") @unittest.skipIf(is_hip(), "HiCache + EAGLE3 file-storage loadback e2e is CUDA-only.") -class TestHiCacheSpecFileStorage(CustomTestCase): - model = DEFAULT_TARGET_MODEL_EAGLE3 - draft_model = DEFAULT_DRAFT_MODEL_EAGLE3 - - input_token_len = 1024 - max_new_tokens = 200 - page_size = 64 - min_expected_accept_length = 7.0 - min_second_to_first_accept_ratio = 0.9 +class TestHiCacheSpecFileStorage(HiCacheSpecStorageMixin, CustomTestCase): + storage_backend = "file" + expected_storage_backend = "HiCacheFile" storage_wait_timeout = 30 - first_measure_new_tokens = 128 @classmethod def setUpClass(cls): cls.temp_dir = tempfile.mkdtemp() - default_port = int(DEFAULT_URL_FOR_TEST.rsplit(":", 1)[1]) - cls.base_url = f"http://127.0.0.1:{find_available_port(default_port)}" - - cls.tokenizer = get_tokenizer(cls.model) - cls.prompt_input_ids = cls._build_long_repetitive_prompt_ids( - cls.tokenizer, cls.input_token_len - ) - - extra_config = { - "hicache_storage_pass_prefix_keys": True, - } - cls.other_args = [ - "--enable-hierarchical-cache", - "--enable-cache-report", - "--mem-fraction-static", - "0.3", - "--hicache-ratio", - "1.5", - "--disable-cuda-graph", - "--page-size", - str(cls.page_size), - "--hicache-storage-backend", - "file", - "--hicache-storage-prefetch-policy", - "wait_complete", - "--hicache-storage-backend-extra-config", - json.dumps(extra_config), - "--speculative-algorithm", - "EAGLE3", - "--speculative-draft-model-path", - cls.draft_model, - "--speculative-num-steps", - "7", - "--speculative-eagle-topk", - "1", - "--speculative-num-draft-tokens", - "8", - "--dtype", - "float16", - ] - cls.env = { - **os.environ, - "SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN": "1", - "SGLANG_HICACHE_FILE_BACKEND_STORAGE_DIR": cls.temp_dir, - } cls.process = None - cls._launch_server() + cls._launch_spec_server() @classmethod - def _launch_server(cls): - cls.process = popen_launch_server( - cls.model, - cls.base_url, - timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, - other_args=cls.other_args, - env=cls.env, - ) - wait_for_http_ready( - url=f"{cls.base_url}/health", - timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, - process=cls.process, - ) - - @classmethod - def _stop_server(cls): - if getattr(cls, "process", None) is None: - return - - process = cls.process - try: - root = psutil.Process(process.pid) - watched_procs = [root] + root.children(recursive=True) - except psutil.NoSuchProcess: - watched_procs = [] - - try: - kill_process_tree(process.pid, wait_timeout=60) - except RuntimeError: - non_zombie_procs = [] - for proc in watched_procs: - try: - if proc.is_running() and proc.status() != psutil.STATUS_ZOMBIE: - non_zombie_procs.append(proc) - except psutil.NoSuchProcess: - pass - if non_zombie_procs: - raise - finally: - cls.process = None - - @classmethod - def _restart_server(cls): - cls._stop_server() - cls._launch_server() + def _get_spec_server_env(cls): + return {"SGLANG_HICACHE_FILE_BACKEND_STORAGE_DIR": cls.temp_dir} @classmethod def _count_file_storage_pages(cls): @@ -156,7 +48,7 @@ class TestHiCacheSpecFileStorage(CustomTestCase): for filename in filenames: if not filename.endswith(".bin"): continue - if filename.startswith("d:"): + if f".{PoolName.DRAFT}" in filename: draft_pages += 1 else: target_pages += 1 @@ -181,122 +73,16 @@ class TestHiCacheSpecFileStorage(CustomTestCase): @classmethod def tearDownClass(cls): - cls._stop_server() + cls._stop_spec_server() if hasattr(cls, "temp_dir"): shutil.rmtree(cls.temp_dir, ignore_errors=True) - @classmethod - def _encode_without_special_tokens(cls, tokenizer, text: str) -> List[int]: - return tokenizer.encode(text, add_special_tokens=False) - - @classmethod - def _build_long_repetitive_prompt_ids(cls, tokenizer, target_len: int) -> List[int]: - bos_ids = ( - [tokenizer.bos_token_id] - if getattr(tokenizer, "bos_token_id", None) is not None - else [] - ) - suffix_ids = cls._encode_without_special_tokens( - tokenizer, - "\n\nContinue the sequence with only the word apple separated by spaces.\n" - "Answer: apple apple apple apple", - ) - repeat_ids = cls._encode_without_special_tokens(tokenizer, " apple") - if not repeat_ids: - raise ValueError( - "Tokenizer produced no ids for the repetitive prompt seed." - ) - if len(bos_ids) + len(suffix_ids) >= target_len: - raise ValueError( - "Prompt suffix is too long: " - f"{len(bos_ids)=}, {len(suffix_ids)=}, {target_len=}." - ) - - prefix_len = target_len - len(bos_ids) - len(suffix_ids) - repeats = (prefix_len + len(repeat_ids) - 1) // len(repeat_ids) - prefix_ids = (repeat_ids * repeats)[:prefix_len] - prompt_ids = bos_ids + prefix_ids + suffix_ids - assert len(prompt_ids) == target_len - return prompt_ids - - def _send_long_prompt(self, max_new_tokens: int = None) -> Dict: - if max_new_tokens is None: - max_new_tokens = self.max_new_tokens - response = requests.post( - f"{self.base_url}/generate", - json={ - "input_ids": self.prompt_input_ids, - "sampling_params": { - "temperature": 0, - "max_new_tokens": max_new_tokens, - "ignore_eos": True, - }, - }, - timeout=900, - ) - self.assertEqual( - response.status_code, - 200, - f"Request failed: {response.status_code} - {response.text}", - ) - return response.json() - - def _get_spec_accept_length(self, response_json: Dict) -> float: - meta_info = response_json.get("meta_info", {}) - self.assertIn( - "spec_accept_length", - meta_info, - f"Missing spec_accept_length in meta_info: {meta_info}", - ) - return float(meta_info["spec_accept_length"]) - - def test_file_storage_loadback_keeps_spec_accept_length(self): - first = self._send_long_prompt(max_new_tokens=self.first_measure_new_tokens) - first_accept_length = self._get_spec_accept_length(first) - self.assertGreaterEqual( - first_accept_length, - self.min_expected_accept_length, - f"First prompt accept length is too low: {first_accept_length=}", - ) - + def _wait_for_storage_before_restart(self): target_pages, draft_pages = self._wait_for_file_storage_pages() print(f"file_storage_before_restart: {target_pages=}, {draft_pages=}") - self._restart_server() - - second = self._send_long_prompt() - second_accept_length = self._get_spec_accept_length(second) - second_meta = second.get("meta_info", {}) - cached_details = second_meta.get("cached_tokens_details") or {} - storage_cached_tokens = int(cached_details.get("storage", 0)) - - print( - f"{first_accept_length=:.3f}, {second_accept_length=:.3f}, " - f"{storage_cached_tokens=}, {cached_details=}" - ) - - self.assertGreaterEqual( - storage_cached_tokens, - self.input_token_len - 2 * self.page_size, - "Expected the second request to load the long prompt KV cache from " - f"file storage, got {cached_details=}", - ) - self.assertEqual( - cached_details.get("storage_backend"), - "HiCacheFile", - f"Expected file storage backend in cache report, got {cached_details=}", - ) - self.assertGreaterEqual( - second_accept_length, - self.min_expected_accept_length, - f"Second prompt accept length is too low: {second_accept_length=}", - ) - self.assertGreaterEqual( - second_accept_length, - first_accept_length * self.min_second_to_first_accept_ratio, - "Spec accept length dropped after file-storage loadback: " - f"{first_accept_length=:.3f}, {second_accept_length=:.3f}", - ) + def test_file_storage_loadback_keeps_spec_accept_length(self): + self._run_storage_loadback_keeps_spec_accept_length() if __name__ == "__main__": diff --git a/test/registered/hicache/test_hicache_spec_mooncake_storage.py b/test/registered/hicache/test_hicache_spec_mooncake_storage.py new file mode 100644 index 000000000..3de31cd9c --- /dev/null +++ b/test/registered/hicache/test_hicache_spec_mooncake_storage.py @@ -0,0 +1,187 @@ +""" +E2E test for HiCache mooncake storage with EAGLE3 speculative decoding. + +Verifies that after a server restart, both target and draft KV cache for a +long deterministic prompt are reloaded from the mooncake_master process and +the spec accept length does not regress. + +Mooncake_master + http metadata server lifecycle is reused from +``test_hicache_storage_mooncake_backend.HiCacheStorageMooncakeBackendBaseMixin``; +this file only overrides the SGLang-side setup for EAGLE3 + spec loadback. + +Usage: + python3 -m pytest test/registered/hicache/test_hicache_spec_mooncake_storage.py -v +""" + +import os +import signal +import socket +import subprocess +import time +import unittest + +from test_hicache_storage_mooncake_backend import ( + HiCacheStorageMooncakeBackendBaseMixin, +) + +from sglang.srt.utils import is_hip +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.hicache_spec_storage_common import HiCacheSpecStorageMixin +from sglang.test.test_utils import CustomTestCase, find_available_port + +register_cuda_ci(est_time=240, stage="extra-a", runner_config="2-gpu-large") + + +@unittest.skipIf( + is_hip(), "HiCache + EAGLE3 mooncake-storage loadback e2e is CUDA-only." +) +class TestHiCacheSpecMooncakeStorage( + HiCacheSpecStorageMixin, HiCacheStorageMooncakeBackendBaseMixin, CustomTestCase +): + """EAGLE3 + mooncake L3 loadback. After a server restart the long + deterministic prompt's target+draft KV must be reloaded from + mooncake_master and the spec accept length must not regress.""" + + storage_backend = "mooncake" + expected_storage_backend = "MooncakeStore" + # Mooncake exposes no external page count; wait this long after the + # first prompt for HiCacheController's backup queue to drain into the + # mooncake_master before restart. + mooncake_backup_drain_seconds = 15 + mooncake_store_port_base = 50052 + mooncake_store_http_port_base = 8081 + + # The inherited test from HiCacheStorageBaseMixin assumes a basic-storage + # setup; it is already covered by test_hicache_storage_mooncake_backend.py. + @unittest.skip("Covered by test_hicache_storage_mooncake_backend.py") + def test_basic_backup_and_prefetch(self): # noqa: D401 + pass + + @classmethod + def setUpClass(cls): + # Bypass the parent chain's basic-storage server bootstrap; only + # reuse the inherited mooncake_master / metadata-server lifecycle. + cls.mooncake_master_port = find_available_port(cls.mooncake_master_port_base) + cls.mooncake_metadata_port = find_available_port( + cls.mooncake_metadata_port_base + ) + cls.mooncake_store_port = find_available_port(cls.mooncake_store_port_base) + cls.mooncake_store_http_port = find_available_port( + cls.mooncake_store_http_port_base + ) + cls._start_mooncake_services() + try: + cls._start_mooncake_store_service() + cls._launch_spec_server() + except Exception: + cls._stop_mooncake_store_service() + cls._stop_mooncake_services() + raise + + @classmethod + def tearDownClass(cls): + try: + cls._stop_spec_server() + finally: + cls._stop_mooncake_store_service() + cls._stop_mooncake_services() + + # ---- server side ---- + + @classmethod + def _start_mooncake_store_service(cls): + print( + f"Starting Mooncake store service on rpc port {cls.mooncake_store_port}, " + f"http port {cls.mooncake_store_http_port}..." + ) + cls.store_service_process = subprocess.Popen( + [ + "mooncake_client", + "--host=127.0.0.1", + f"--port={cls.mooncake_store_port}", + f"--master_server_address=127.0.0.1:{cls.mooncake_master_port}", + f"--metadata_server=http://127.0.0.1:{cls.mooncake_metadata_port}/metadata", + "--protocol=tcp", + "--device_names=", + "--global_segment_size=4294967296", + "--enable_http_server=true", + f"--http_port={cls.mooncake_store_http_port}", + ], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + preexec_fn=os.setsid, + ) + cls._wait_for_mooncake_store_service_ready() + + @classmethod + def _wait_for_mooncake_store_service_ready(cls, timeout: int = 90): + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + process = getattr(cls, "store_service_process", None) + if process is not None and process.poll() is not None: + raise RuntimeError( + f"Mooncake store service exited with code {process.returncode}" + ) + try: + with socket.create_connection( + ("127.0.0.1", cls.mooncake_store_port), timeout=2 + ): + pass + print("Mooncake store service is ready") + return + except OSError: + pass + time.sleep(1) + raise TimeoutError("Timed out waiting for Mooncake store service readiness.") + + @classmethod + def _stop_mooncake_store_service(cls): + process = getattr(cls, "store_service_process", None) + if process is None: + return + print("Stopping Mooncake store service...") + try: + os.killpg(os.getpgid(process.pid), signal.SIGTERM) + process.wait(timeout=15) + print("Mooncake store service stopped") + except (ProcessLookupError, subprocess.TimeoutExpired, OSError): + try: + os.killpg(os.getpgid(process.pid), signal.SIGKILL) + process.wait(timeout=5) + except (ProcessLookupError, subprocess.TimeoutExpired, OSError) as e: + print(f"Warning: Could not stop Mooncake store service: {e}") + finally: + cls.store_service_process = None + + @classmethod + def _get_spec_server_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" + ), + # Keep storage capacity in the external mooncake_client so cached + # pages survive SGLang server restart. + "MOONCAKE_GLOBAL_SEGMENT_SIZE": "0", + } + + # ---- prompt / IO helpers ---- + + def _wait_for_storage_before_restart(self): + print( + f"[mooncake] draining backup queue " + f"({self.mooncake_backup_drain_seconds}s)..." + ) + time.sleep(self.mooncake_backup_drain_seconds) + + # ---- the test ---- + + def test_mooncake_storage_loadback_keeps_spec_accept_length(self): + self._run_storage_loadback_keeps_spec_accept_length() + + +if __name__ == "__main__": + unittest.main(verbosity=2)