diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index a9c7a0231..ba2364398 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -1868,7 +1868,7 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache): operation = self.cache_controller.prefetch( req_id, host_indices, - prefetch_key.token_ids, + prefetch_key, last_hash, prefix_keys, extra_pools=aux_xfers or None, diff --git a/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_dsv4.py b/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_dsv4.py index c19b7928f..617defda6 100644 --- a/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_dsv4.py +++ b/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_dsv4.py @@ -1,8 +1,10 @@ import os import shutil import tempfile +import time import unittest +import requests from test_unified_radix_cache_kl_nightly import AccuracyTwoPassMixin from sglang.srt.utils import kill_process_tree @@ -19,7 +21,7 @@ from sglang.test.test_utils import ( DSV4_FLASH_MODEL = "sgl-project/DeepSeek-V4-Flash-FP8" DSV4_FLASH_LAUNCH_TIMEOUT = 3600 -register_cuda_ci(est_time=768, stage="base-c", runner_config="4-gpu-h100") +register_cuda_ci(est_time=1200, stage="base-c", runner_config="4-gpu-h100") def _assert_dsv4_decode_cached_tokens(result, history_len, output_len, label): @@ -163,5 +165,151 @@ class TestUnifiedDeepSeekV4FlashHiCacheL3(AccuracyTwoPassMixin, CustomTestCase): shutil.rmtree(cls.hicache_dir, ignore_errors=True) +class TestUnifiedDeepSeekV4FlashEagleHiCacheL3(AccuracyTwoPassMixin, CustomTestCase): + """DeepSeek V4 Flash EAGLE + HiCache L3 should load from storage.""" + + page_size = 256 + input_ids = list(range(4000, 4300)) + storage_wait_timeout = 120 + num_gsm8k_questions = 100 + gsm8k_parallel = 4 + mmlu_num_threads = 4 + + @classmethod + def setUpClass(cls): + cls.model = DSV4_FLASH_MODEL + cls.base_url = DEFAULT_URL_FOR_TEST + cls.hicache_dir = tempfile.mkdtemp(prefix="hicache_l3_eagle_dsv4_") + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DSV4_FLASH_LAUNCH_TIMEOUT, + other_args=[ + "--trust-remote-code", + "--tp-size", + "4", + "--attention-backend", + "compressed", + "--page-size", + str(cls.page_size), + "--chunked-prefill-size", + "8192", + "--mem-fraction-static", + "0.9", + "--disable-shared-experts-fusion", + "--enable-hierarchical-cache", + "--hicache-ratio", + "2", + "--hicache-write-policy", + "write_through", + "--hicache-storage-prefetch-policy", + "wait_complete", + "--hicache-io-backend", + "direct", + "--hicache-mem-layout", + "page_first_direct", + "--hicache-storage-backend", + "file", + "--enable-cache-report", + "--swa-full-tokens-ratio", + "0.25", + "--max-total-tokens", + "20000", + "--max-running-requests", + "4", + "--speculative-algorithm", + "EAGLE", + "--speculative-num-steps", + "3", + "--speculative-eagle-topk", + "1", + "--speculative-num-draft-tokens", + "4", + ], + env={ + "SGLANG_DSV4_FP4_EXPERTS": "0", + "SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1", + "SGLANG_HICACHE_FILE_BACKEND_STORAGE_DIR": cls.hicache_dir, + }, + ) + + @classmethod + def tearDownClass(cls): + kill_process_tree(cls.process.pid) + if os.path.isdir(cls.hicache_dir): + shutil.rmtree(cls.hicache_dir, ignore_errors=True) + + @classmethod + def _count_file_storage_pages(cls): + try: + return sum( + 1 + for filename in os.listdir(cls.hicache_dir) + if filename.endswith(".bin") + ) + except FileNotFoundError: + return 0 + + @classmethod + def _wait_for_file_storage_pages(cls, min_pages: int): + deadline = time.monotonic() + cls.storage_wait_timeout + pages = 0 + while time.monotonic() < deadline: + pages = cls._count_file_storage_pages() + if pages >= min_pages: + return + time.sleep(0.2) + raise AssertionError( + f"Timed out waiting for HiCache file storage pages: {pages=}, {min_pages=}" + ) + + def _flush_cache(self): + response = requests.post( + self.base_url + "/flush_cache", + params={"timeout": 30}, + timeout=120, + ) + response.raise_for_status() + + def _generate(self): + response = requests.post( + self.base_url + "/generate", + json={ + "input_ids": self.input_ids, + "sampling_params": { + "temperature": 0, + "max_new_tokens": 1, + }, + }, + timeout=1200, + ) + self.assertEqual( + response.status_code, + 200, + f"Request failed: {response.status_code} - {response.text}", + ) + return response.json() + + def test_eagle_l3_storage_cache_hit(self): + self._flush_cache() + initial_pages = self._count_file_storage_pages() + + first = self._generate() + self.assertEqual(first["meta_info"]["cached_tokens"], 0) + + self._wait_for_file_storage_pages(initial_pages + 1) + self._flush_cache() + + second = self._generate() + cached_details = second["meta_info"].get("cached_tokens_details") or {} + storage_cached_tokens = int(cached_details.get("storage", 0)) + self.assertGreaterEqual( + storage_cached_tokens, + self.page_size, + f"Expected EAGLE request to load KV from HiCache file storage, got {cached_details=}", + ) + self.assertEqual(cached_details.get("storage_backend"), "HiCacheFile") + + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py index 2ecda3ed6..261073705 100644 --- a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py +++ b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py @@ -99,6 +99,7 @@ class CacheConfig: head_dim: int = 64 dtype: torch.dtype = torch.bfloat16 eviction_policy: str = "lru" + is_eagle: bool = False @property def has_mamba(self) -> bool: @@ -125,6 +126,8 @@ class CacheConfig: or self.num_layers != defaults["num_layers"].default ): parts.append(f"h{self.head_num}l{self.num_layers}") + if self.is_eagle: + parts.append("eagle") return "_".join(parts) @@ -324,6 +327,7 @@ def build_fixture(cfg: CacheConfig, *, enable_kv_cache_events: bool = False): enable_mamba_extra_buffer=cfg.enable_mamba_extra_buffer, enable_kv_cache_events=enable_kv_cache_events, eviction_policy=cfg.eviction_policy, + is_eagle=cfg.is_eagle, ) tree = UnifiedRadixCache(params=cache_init_params) tree.cache_init_params = cache_init_params @@ -331,6 +335,91 @@ def build_fixture(cfg: CacheConfig, *, enable_kv_cache_events: bool = False): return tree, allocator, req_to_token_pool +class TestUnifiedRadixCacheEagleHiCacheStorageKey(CustomTestCase): + cfg = CacheConfig( + page_size=4, + components=(ComponentType.FULL,), + is_eagle=True, + kv_size=64, + max_context_len=64, + ) + + def test_l3_prefetch_uses_bigram_radix_key(self): + from sglang.srt.mem_cache.utils import get_hash_str + + tree, allocator, _ = build_fixture(self.cfg) + tree.enable_storage = True + tree.prefetch_threshold = 1 + tokens = array("q", [1, 2, 3, 4, 5, 6, 7, 8, 9]) + + value = allocator.alloc(len(tokens) - 1) + self.assertIsNotNone(value) + tree.insert(InsertParams(key=RadixKey(tokens), value=value)) + match = tree.match_prefix(MatchPrefixParams(key=RadixKey(tokens))) + leaf = match.last_device_node + self.assertTrue(leaf.key.is_bigram) + self.assertEqual(len(leaf.hash_value), 2) + + class FakeHostPool: + def alloc(self, num_tokens): + return torch.arange(num_tokens, dtype=torch.int64) + + class FakeCacheController: + def __init__(self): + self.mem_pool_host = FakeHostPool() + self.prefetch_tokens_occupied = 0 + self.prefetch_args = None + + def prefetch_rate_limited(self): + return False + + def prefetch( + self, + request_id, + host_indices, + new_input_tokens, + last_hash=None, + prefix_keys=None, + extra_pools=None, + ): + self.prefetch_args = ( + request_id, + host_indices, + new_input_tokens, + last_hash, + prefix_keys, + extra_pools, + ) + return mock.Mock() + + controller = FakeCacheController() + tree.cache_controller = controller + tree.prefetch_from_storage("req", tree.root_node, tokens) + + _, _, storage_key, _, _, _ = controller.prefetch_args + self.assertIsInstance(storage_key, RadixKey) + self.assertTrue(storage_key.is_bigram) + self.assertEqual(len(storage_key), len(tokens) - 1) + + queried_hashes = [] + running_hash = None + for start in range(0, len(storage_key), tree.page_size): + running_hash = get_hash_str( + storage_key[start : start + tree.page_size], running_hash + ) + queried_hashes.append(running_hash) + self.assertEqual(queried_hashes, leaf.hash_value) + + canonical_hashes = [] + running_hash = None + for start in range(0, len(tokens) - 1, tree.page_size): + running_hash = get_hash_str( + tokens[start : start + tree.page_size], running_hash + ) + canonical_hashes.append(running_hash) + self.assertNotEqual(canonical_hashes, leaf.hash_value) + + class TestUnifiedRadixCacheKVEvents(CustomTestCase): cfg = CacheConfig(page_size=2, kv_size=64, max_context_len=64)