[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>

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

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
 

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

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}",
+        )

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__":

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)