[UnifiedTree]: Support l3 storage for swa and deepseek v4 (#26881)
Co-authored-by: 晟海 <huangtingwei.htw@antgroup.com>
This commit is contained in:
Zhangheng
@@ -1,5 +1,9 @@
|
||||
"""Unit tests for UnifiedRadixCache"""
|
||||
|
||||
import json
|
||||
import shutil
|
||||
import tempfile
|
||||
import time
|
||||
import unittest
|
||||
from array import array
|
||||
from dataclasses import dataclass, replace
|
||||
@@ -1840,6 +1844,146 @@ class UnifiedRadixCacheSuite:
|
||||
# HiCache Unit Tests (real cache_controller D<->H backup/load)
|
||||
# ================================================================
|
||||
|
||||
# ---------- L3 storage (file backend) helpers ----------
|
||||
|
||||
def _path_chain(self, tree, node):
|
||||
"""Return root->node node chain (excluding root)."""
|
||||
chain = []
|
||||
cur = node
|
||||
while cur is not tree.root_node:
|
||||
chain.append(cur)
|
||||
cur = cur.parent
|
||||
chain.reverse()
|
||||
return chain
|
||||
|
||||
def _write_path_to_l3(self, tree, node):
|
||||
"""Offload every node on root->node path from host to L3 storage."""
|
||||
for n in self._path_chain(tree, node):
|
||||
tree.write_backup_storage(n)
|
||||
|
||||
def _flush_l3_backups(self, tree, timeout: float = 10.0):
|
||||
"""Wait for backup threads to finish, then drain acks (release locks)."""
|
||||
deadline = time.time() + timeout
|
||||
while tree.ongoing_backup and time.time() < deadline:
|
||||
tree.drain_storage_control_queues()
|
||||
if tree.ongoing_backup:
|
||||
time.sleep(0.01)
|
||||
tree.drain_storage_control_queues()
|
||||
self.assertFalse(tree.ongoing_backup, "L3 backups did not complete in time")
|
||||
|
||||
def _run_prefetch_to_completion(self, tree, req_id, timeout: float = 10.0):
|
||||
deadline = time.time() + timeout
|
||||
while time.time() < deadline:
|
||||
if tree.check_prefetch_progress(req_id):
|
||||
return
|
||||
time.sleep(0.01)
|
||||
self.fail(f"prefetch {req_id} did not complete in time")
|
||||
|
||||
def _all_page_hashes(self, tree, node):
|
||||
hashes = []
|
||||
for n in self._path_chain(tree, node):
|
||||
hashes.extend(list(n.hash_value))
|
||||
return hashes
|
||||
|
||||
def test_hicache_l3_write_storage(self):
|
||||
"""D->H->L3 offload: every KV page lands in the file storage backend."""
|
||||
if self._skip_unsupported_hicache_test():
|
||||
return
|
||||
if self.cfg.has_mamba:
|
||||
self.skipTest("mamba L3 offload is out of scope for this unit fixture")
|
||||
|
||||
storage_dir = tempfile.mkdtemp()
|
||||
self.addCleanup(shutil.rmtree, storage_dir, ignore_errors=True)
|
||||
|
||||
tree, allocator, req_to_token_pool = build_fixture(self.cfg)
|
||||
self._init_hicache(
|
||||
tree,
|
||||
storage_backend="file",
|
||||
storage_dir=storage_dir,
|
||||
prefetch_threshold=1,
|
||||
)
|
||||
|
||||
seq = self._make_seq(1, 4)
|
||||
self._insert(tree, allocator, req_to_token_pool, seq)
|
||||
m = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq))))
|
||||
leaf = m.last_device_node
|
||||
|
||||
# D->H first, then H->L3.
|
||||
self._backup_node(tree, leaf)
|
||||
self.assertTrue(leaf.hash_value)
|
||||
self._write_path_to_l3(tree, leaf)
|
||||
self._flush_l3_backups(tree)
|
||||
|
||||
# Every KV page hash on the path must now exist in storage.
|
||||
backend = tree.cache_controller.storage_backend
|
||||
page_hashes = self._all_page_hashes(tree, leaf)
|
||||
self.assertEqual(len(page_hashes), len(seq) // self.cfg.page_size)
|
||||
self.assertEqual(backend.batch_exists(page_hashes), len(page_hashes))
|
||||
tree.sanity_check()
|
||||
|
||||
def test_hicache_l3_prefetch(self):
|
||||
"""L3 round trip: write with one tree, prefetch into a fresh tree.
|
||||
|
||||
Uses two independent trees that share the same file storage dir so the
|
||||
prefetch path genuinely reloads from L3 (no host/device residue).
|
||||
"""
|
||||
if self._skip_unsupported_hicache_test():
|
||||
return
|
||||
if self.cfg.has_mamba:
|
||||
self.skipTest("mamba L3 prefetch is out of scope for this unit fixture")
|
||||
|
||||
storage_dir = tempfile.mkdtemp()
|
||||
self.addCleanup(shutil.rmtree, storage_dir, ignore_errors=True)
|
||||
seq = self._make_seq(1, 4)
|
||||
|
||||
# --- Producer tree: fill KV, backup D->H, offload H->L3. ---
|
||||
prod, prod_alloc, prod_rtp = build_fixture(self.cfg)
|
||||
self._init_hicache(
|
||||
prod,
|
||||
storage_backend="file",
|
||||
storage_dir=storage_dir,
|
||||
prefetch_threshold=1,
|
||||
)
|
||||
self._insert(prod, prod_alloc, prod_rtp, seq)
|
||||
mp = prod.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq))))
|
||||
prod_leaf = mp.last_device_node
|
||||
self._fill_full_kv(prod_alloc, mp.device_indices, marker=7)
|
||||
expected_k, expected_v = self._snapshot_full_kv(prod_alloc, mp.device_indices)
|
||||
self._backup_node(prod, prod_leaf)
|
||||
self._write_path_to_l3(prod, prod_leaf)
|
||||
self._flush_l3_backups(prod)
|
||||
|
||||
# --- Consumer tree: prefetch the same tokens straight from L3. ---
|
||||
cons, cons_alloc, cons_rtp = build_fixture(self.cfg)
|
||||
self._init_hicache(
|
||||
cons,
|
||||
storage_backend="file",
|
||||
storage_dir=storage_dir,
|
||||
prefetch_threshold=1,
|
||||
)
|
||||
req_id = "l3-prefetch-req"
|
||||
cons.prefetch_from_storage(req_id, cons.root_node, array("q", seq), None, None)
|
||||
self._run_prefetch_to_completion(cons, req_id)
|
||||
cons.drain_storage_control_queues()
|
||||
|
||||
# The full prefix must now be a host hit (loaded from L3).
|
||||
mc = cons.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq))))
|
||||
self.assertEqual(mc.host_hit_length, len(seq))
|
||||
host_node = mc.last_host_node
|
||||
self.assertIsNot(host_node, cons.root_node)
|
||||
self.assertTrue(host_node.evicted)
|
||||
|
||||
# Load the reloaded host prefix back to device and verify KV bytes.
|
||||
self._load_back_node(cons, host_node)
|
||||
loaded_indices = cons.match_prefix(
|
||||
MatchPrefixParams(key=RadixKey(array("q", seq)))
|
||||
).device_indices
|
||||
self.assertEqual(len(loaded_indices), len(seq))
|
||||
loaded_k, loaded_v = self._snapshot_full_kv(cons_alloc, loaded_indices)
|
||||
self.assertTrue(torch.equal(loaded_k, expected_k))
|
||||
self.assertTrue(torch.equal(loaded_v, expected_v))
|
||||
cons.sanity_check()
|
||||
|
||||
def _skip_unsupported_hicache_test(self):
|
||||
if self.cfg.has_swa and self.cfg.has_mamba:
|
||||
self.skipTest("HiCache unit fixture does not support SWA + Mamba stacks")
|
||||
@@ -1869,7 +2013,16 @@ class UnifiedRadixCacheSuite:
|
||||
self._simulate_backup(tree, node)
|
||||
stack.extend(node.children.values())
|
||||
|
||||
def _init_hicache(self, tree, *, write_policy: str = "write_through"):
|
||||
def _init_hicache(
|
||||
self,
|
||||
tree,
|
||||
*,
|
||||
write_policy: str = "write_through",
|
||||
storage_backend: Optional[str] = None,
|
||||
storage_dir: Optional[str] = None,
|
||||
prefetch_threshold: Optional[int] = None,
|
||||
prefetch_policy: str = "wait_complete",
|
||||
):
|
||||
import sglang.srt.mem_cache.hybrid_cache.hybrid_pool_assembler as assembler
|
||||
|
||||
orig_kv_host_pool = assembler.MHATokenToKVPoolHost
|
||||
@@ -1899,11 +2052,42 @@ class UnifiedRadixCacheSuite:
|
||||
patcher.start()
|
||||
self.addCleanup(patcher.stop)
|
||||
|
||||
storage_extra_config = None
|
||||
if storage_backend == "file":
|
||||
import sglang.srt.managers.cache_controller as cache_controller
|
||||
|
||||
# The file-backend storage config records TP rank/size. These unit
|
||||
# fixtures run without initializing distributed parallel state, so
|
||||
# provide the local single-rank values that the fixture represents.
|
||||
tp_rank_patcher = mock.patch.object(
|
||||
cache_controller, "get_tensor_model_parallel_rank", return_value=0
|
||||
)
|
||||
tp_size_patcher = mock.patch.object(
|
||||
cache_controller, "get_tensor_model_parallel_world_size", return_value=1
|
||||
)
|
||||
tp_rank_patcher.start()
|
||||
tp_size_patcher.start()
|
||||
self.addCleanup(tp_rank_patcher.stop)
|
||||
self.addCleanup(tp_size_patcher.stop)
|
||||
|
||||
assert storage_dir is not None, "file backend needs a storage_dir"
|
||||
# HiCacheFile reads the directory from this env var.
|
||||
cm = envs.SGLANG_HICACHE_FILE_BACKEND_STORAGE_DIR.override(storage_dir)
|
||||
cm.__enter__()
|
||||
self.addCleanup(cm.__exit__, None, None, None)
|
||||
extra = {}
|
||||
if prefetch_threshold is not None:
|
||||
extra["prefetch_threshold"] = prefetch_threshold
|
||||
storage_extra_config = json.dumps(extra) if extra else None
|
||||
|
||||
server_args = ServerArgs(
|
||||
model_path="dummy",
|
||||
page_size=self.cfg.page_size,
|
||||
hicache_io_backend="direct",
|
||||
hicache_write_policy=write_policy,
|
||||
hicache_storage_backend=storage_backend,
|
||||
hicache_storage_backend_extra_config=storage_extra_config,
|
||||
hicache_storage_prefetch_policy=prefetch_policy,
|
||||
)
|
||||
# See build_fixture for why _mamba_cache_chunk_size is preset.
|
||||
server_args._mamba_cache_chunk_size = max(FLA_CHUNK_SIZE, self.cfg.page_size)
|
||||
@@ -1911,6 +2095,17 @@ class UnifiedRadixCacheSuite:
|
||||
tree.init_hicache(server_args, tree.cache_init_params)
|
||||
tree.write_through_threshold = 1 << 30
|
||||
tree.load_back_threshold = 0
|
||||
if storage_backend is not None:
|
||||
# Unit fixtures size host/device pools equally, which makes the
|
||||
# production prefetch capacity limit (host - device) zero. Keep the
|
||||
# L3 tests focused on storage round trips by allowing one fixture
|
||||
# worth of prefetch tokens.
|
||||
tree.cache_controller.prefetch_capacity_limit = max(
|
||||
tree.cache_controller.prefetch_capacity_limit,
|
||||
tree.cache_controller.mem_pool_host.size,
|
||||
)
|
||||
# Background prefetch/backup threads are daemon; stop them per-test.
|
||||
self.addCleanup(tree.cache_controller._stop_storage_threads)
|
||||
|
||||
def _build_hicache_fixture(self):
|
||||
fixture = build_fixture(self.cfg)
|
||||
|
||||
Reference in New Issue
Block a user