chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

feat: track CP HiCache write-back state

Browse Source
This commit is contained in:
谢骞
2026-05-08 01:29:39 +08:00
parent a8816f18dd
commit 068aa2f910
2 changed files with 152 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

64
python/sglang/srt/mem_cache/hiradix_cache.py
View File

@@ -15,6 +15,7 @@ import torch
from sglang.srt.environ import envs
from sglang.srt.managers.cache_controller import HiCacheController, PrefetchOperation
from sglang.srt.mem_cache.allocator import CPSharedPagedTokenToKVPoolAllocator
from sglang.srt.mem_cache.base_prefix_cache import (
DecLockRefParams,
DecLockRefResult,
@@ -27,6 +28,7 @@ from sglang.srt.mem_cache.base_prefix_cache import (
MatchPrefixParams,
MatchResult,
)
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
from sglang.srt.mem_cache.memory_pool import (
MHATokenToKVPool,
MLATokenToKVPool,
@@ -117,6 +119,21 @@ class HiRadixCache(RadixCache):
bind_to_closest_numa_node_cuda()
self.page_size = params.page_size
self._uses_cp_hicache = isinstance(
params.token_to_kv_pool_allocator,
CPSharedPagedTokenToKVPoolAllocator,
)
if self._uses_cp_hicache:
if server_args.hicache_storage_backend is not None:
raise ValueError(
"CP shared KV HiCache host integration does not support storage backends."
)
if not isinstance(
params.token_to_kv_pool_allocator.get_kvcache(), NSATokenToKVPool
):
raise ValueError(
"CP shared KV HiCache host integration requires NSATokenToKVPool."
)
self.kv_cache = params.token_to_kv_pool_allocator.get_kvcache()
if isinstance(self.kv_cache, MHATokenToKVPool):
@@ -171,6 +188,13 @@ class HiRadixCache(RadixCache):
self.prefetch_stop_policy = server_args.hicache_storage_prefetch_policy
self.load_cache_event = threading.Event()
cp_shared_kv_layout = None
if self._uses_cp_hicache:
cp_shared_kv_layout = CpSharedKVLayout(
page_size=self.page_size,
cp_size=params.token_to_kv_pool_allocator.cp_size,
cp_rank=params.token_to_kv_pool_allocator.cp_rank,
)
self.cache_controller = HiCacheController(
params.token_to_kv_pool_allocator,
self.token_to_kv_pool_host,
@@ -186,6 +210,7 @@ class HiRadixCache(RadixCache):
pp_rank=self.pp_rank,
pp_size=self.pp_size,
enable_storage_metrics=self.enable_storage_metrics,
cp_shared_kv_layout=cp_shared_kv_layout,
)
self._apply_storage_runtime_config(
storage_backend=server_args.hicache_storage_backend,
@@ -679,7 +704,44 @@ class HiRadixCache(RadixCache):
logger.warning("Hierarchical cache storage backend is not enabled.")
return False
def _node_backuped(self, node: TreeNode) -> bool:
if self._uses_cp_hicache:
return node.host_len > 0 and node.cp_hicache is not None
return node.host_value is not None
def _node_host_len(self, node: TreeNode) -> int:
if self._uses_cp_hicache:
return node.host_len
return len(node.host_value)
def _node_host_evict_indices(self, node: TreeNode) -> torch.Tensor:
if self._uses_cp_hicache:
return node.cp_hicache.host_indices
return node.host_value
def write_backup(self, node: TreeNode, write_back=False):
if self._uses_cp_hicache:
result = self.cache_controller.write(
device_indices=node.value,
node_id=node.id,
)
if getattr(result, "metadata", None) is None:
self._evict_host_for_physical_slots(result.required_host_slots)
result = self.cache_controller.write(
device_indices=node.value,
node_id=node.id,
)
if getattr(result, "metadata", None) is None:
return 0
node.host_len = len(node.value)
node.cp_hicache = result.metadata
node.host_value = None
self.ongoing_write_through[node.id] = node
if not write_back:
self.inc_node_lock_ref(node)
return len(node.cp_hicache.host_indices)
host_indices = self.cache_controller.write(
device_indices=node.value,
node_id=node.id,
@@ -725,7 +787,7 @@ class HiRadixCache(RadixCache):
return
node.hit_count += 1
if not node.backuped:
if not self._node_backuped(node):
if node.hit_count >= self.write_through_threshold:
# write to host if the node is not backuped
self.write_backup(node)

90
test/registered/unit/mem_cache/test_cp_hicache_metadata.py
View File

@@ -10,7 +10,8 @@ for _mod in ("sgl_kernel", "sgl_kernel.kvcacheio"):
if _mod not in sys.modules:
sys.modules[_mod] = MagicMock()
from sglang.srt.mem_cache.hiradix_cache import CpHiCacheNodeMetadata
from sglang.srt.mem_cache.hiradix_cache import CpHiCacheNodeMetadata, HiRadixCache
from sglang.srt.mem_cache.radix_cache import TreeNode
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
@@ -139,5 +140,92 @@ class TestCpHiCacheNodeMetadata(CustomTestCase):
)
class FakeWriteFailure:
metadata = None
def __init__(self, required_host_slots):
self.required_host_slots = required_host_slots
class FakeWriteSuccess:
required_host_slots = 0
def __init__(self, metadata):
self.metadata = metadata
class FakeWriteController:
def __init__(self, required_host_slots):
self.required_host_slots = required_host_slots
self.calls = 0
self.write_policy = "write_through"
def write(self, device_indices, node_id=-1, priority=None):
self.calls += 1
if self.calls == 1:
return FakeWriteFailure(self.required_host_slots)
return FakeWriteSuccess(
CpHiCacheNodeMetadata(
logical_len=len(device_indices),
owned_positions=torch.tensor([0], dtype=torch.int64),
host_indices=torch.tensor([99], dtype=torch.int64),
)
)
class TestHiRadixCacheCPBackup(CustomTestCase):
def test_node_backuped_uses_cp_metadata(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache._uses_cp_hicache = True
node = TreeNode()
node.host_len = 8
node.cp_hicache = CpHiCacheNodeMetadata(
logical_len=8,
owned_positions=torch.tensor([1, 2], dtype=torch.int64),
host_indices=torch.tensor([10, 11], dtype=torch.int64),
)
self.assertTrue(cache._node_backuped(node))
def test_inc_hit_count_does_not_rewrite_cp_backed_node(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache._uses_cp_hicache = True
cache.write_through_threshold = 1
cache.cache_controller = type(
"Controller", (), {"write_policy": "write_through"}
)()
cache.write_backup = lambda node: (_ for _ in ()).throw(
AssertionError("must not rewrite")
)
node = TreeNode()
node.host_len = 4
node.cp_hicache = CpHiCacheNodeMetadata(
logical_len=4,
owned_positions=torch.tensor([], dtype=torch.int64),
host_indices=torch.tensor([], dtype=torch.int64),
)
cache._inc_hit_count(node)
self.assertEqual(node.hit_count, 1)
def test_write_backup_retries_by_required_physical_slots(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache._uses_cp_hicache = True
cache.cache_controller = FakeWriteController(required_host_slots=3)
cache._evict_host_for_physical_slots = lambda required: setattr(
cache, "evicted_required", required
)
cache.ongoing_write_through = {}
cache.inc_node_lock_ref = lambda node: None
node = TreeNode()
node.value = torch.arange(16, dtype=torch.int64)
cache.write_backup(node)
self.assertEqual(cache.evicted_required, 3)
self.assertEqual(node.host_len, 16)
if __name__ == "__main__":
unittest.main()