[Unified Tree]fix compatibility with eagle key and l3 hicache (#27655)
This commit is contained in:
huangtingwei
@@ -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,
|
||||
|
||||
+149
-1
@@ -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()
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user