From 301420f9763280641f0326c6c53de562b774ce60 Mon Sep 17 00:00:00 2001 From: ThomasX Date: Fri, 8 May 2026 01:14:58 +0800 Subject: [PATCH] feat: map CP HiCache writes to owned physical slots --- .../sglang/srt/managers/cache_controller.py | 80 ++++++- .../managers/test_hicache_controller_cp.py | 210 ++++++++++++++++++ 2 files changed, 288 insertions(+), 2 deletions(-) create mode 100644 test/registered/unit/managers/test_hicache_controller_cp.py diff --git a/python/sglang/srt/managers/cache_controller.py b/python/sglang/srt/managers/cache_controller.py index e4f8d0b70..50479b744 100644 --- a/python/sglang/srt/managers/cache_controller.py +++ b/python/sglang/srt/managers/cache_controller.py @@ -16,6 +16,7 @@ limitations under the License. import logging import threading import time +from dataclasses import dataclass from queue import Empty, Full, Queue from typing import TYPE_CHECKING, List, NamedTuple, Optional @@ -40,7 +41,8 @@ from sglang.srt.layers.dp_attention import ( get_attention_tp_size, is_dp_attention_enabled, ) -from sglang.srt.mem_cache.memory_pool import MLATokenToKVPool +from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout +from sglang.srt.mem_cache.memory_pool import MLATokenToKVPool, NSATokenToKVPool from sglang.srt.utils import get_device_module logger = logging.getLogger(__name__) @@ -244,6 +246,18 @@ class PrefetchOperation(StorageOperation): return self._terminated_flag +@dataclass +class HiCacheWriteResult: + metadata: object + required_host_slots: int = 0 + + +@dataclass +class HiCacheWriteFailure: + required_host_slots: int + metadata: object = None + + class HiCacheController: def __init__( @@ -262,6 +276,7 @@ class HiCacheController: pp_rank: int = 0, pp_size: int = 1, enable_storage_metrics: bool = False, + cp_shared_kv_layout: Optional[CpSharedKVLayout] = None, ): self.tp_group = tp_group self.mem_pool_device_allocator = token_to_kv_pool_allocator @@ -271,6 +286,14 @@ class HiCacheController: if isinstance(mem_pool_device, HybridLinearKVPool): mem_pool_device = mem_pool_device.full_kv_pool self.mem_pool_device = mem_pool_device + self.cp_shared_kv_layout = cp_shared_kv_layout + self.uses_cp_hicache = cp_shared_kv_layout is not None + if self.uses_cp_hicache and not isinstance( + self.mem_pool_device, NSATokenToKVPool + ): + raise ValueError( + "CP shared KV HiCache host integration requires NSATokenToKVPool." + ) self.mem_pool_host = mem_pool_host self.write_policy = write_policy self.page_size = page_size @@ -656,10 +679,13 @@ class HiCacheController: device_indices: torch.Tensor, priority: Optional[int] = None, node_id: int = -1, - ) -> Optional[torch.Tensor]: + ) -> Optional[torch.Tensor | HiCacheWriteResult | HiCacheWriteFailure]: """ Back up KV caches from device memory to host memory. """ + if self.uses_cp_hicache: + return self._write_cp(device_indices, priority, node_id) + host_indices = self.mem_pool_host.alloc(len(device_indices)) if host_indices is None: return None @@ -697,6 +723,56 @@ class HiCacheController: self.ack_write_queue.append(HiCacheAck(start_event, finish_event, op.node_ids)) + def _append_completed_write_ack(self, node_id: int) -> None: + event = device_module.Event() + event.record() + self.ack_write_queue.append(HiCacheAck(event, event, [node_id])) + + def _append_completed_load_ack(self, node_id: int) -> None: + event = device_module.Event() + event.record() + self.ack_load_queue.append(HiCacheAck(event, event, [node_id])) + + def _write_cp( + self, + device_indices: torch.Tensor, + priority: Optional[int] = None, + node_id: int = -1, + ) -> HiCacheWriteResult | HiCacheWriteFailure: + from sglang.srt.mem_cache.hiradix_cache import CpHiCacheNodeMetadata + + layout = self.cp_shared_kv_layout + owned_mask = layout.owned_by_this_rank(device_indices) + owned_positions = owned_mask.nonzero(as_tuple=True)[0].cpu() + logical_len = len(device_indices) + if owned_positions.numel() == 0: + self._append_completed_write_ack(node_id) + return HiCacheWriteResult( + metadata=CpHiCacheNodeMetadata( + logical_len=logical_len, + owned_positions=owned_positions, + host_indices=torch.empty((0,), dtype=torch.int64), + ) + ) + + owned_logical_indices = device_indices[owned_mask] + physical_device_indices = layout.logical_locs_to_physical(owned_logical_indices) + host_indices = self.mem_pool_host.alloc(len(physical_device_indices)) + if host_indices is None: + return HiCacheWriteFailure(required_host_slots=len(physical_device_indices)) + + self.write_queue.append( + CacheOperation(host_indices, physical_device_indices, node_id, priority) + ) + self.start_writing() + return HiCacheWriteResult( + metadata=CpHiCacheNodeMetadata( + logical_len=logical_len, + owned_positions=owned_positions, + host_indices=host_indices.cpu(), + ) + ) + def load( self, host_indices: torch.Tensor, diff --git a/test/registered/unit/managers/test_hicache_controller_cp.py b/test/registered/unit/managers/test_hicache_controller_cp.py new file mode 100644 index 000000000..34ffc5a30 --- /dev/null +++ b/test/registered/unit/managers/test_hicache_controller_cp.py @@ -0,0 +1,210 @@ +import sys +import types +import unittest +from unittest.mock import patch + +import torch + +if "sgl_kernel" not in sys.modules: + sys.modules["sgl_kernel"] = types.ModuleType("sgl_kernel") +if "sgl_kernel.kvcacheio" not in sys.modules: + kvcacheio_stub = types.ModuleType("sgl_kernel.kvcacheio") + for name in ( + "transfer_kv_all_layer", + "transfer_kv_all_layer_direct_lf_pf", + "transfer_kv_all_layer_lf_pf", + "transfer_kv_all_layer_lf_ph", + "transfer_kv_all_layer_mla", + "transfer_kv_all_layer_mla_lf_pf", + "transfer_kv_direct", + "transfer_kv_per_layer", + "transfer_kv_per_layer_direct_pf_lf", + "transfer_kv_per_layer_mla", + "transfer_kv_per_layer_mla_pf_lf", + "transfer_kv_per_layer_pf_lf", + "transfer_kv_per_layer_ph_lf", + ): + setattr(kvcacheio_stub, name, lambda *args, **kwargs: None) + sys.modules["sgl_kernel.kvcacheio"] = kvcacheio_stub + +from sglang.srt.managers.cache_controller import HiCacheController +from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=2, suite="stage-a-test-cpu") + + +class FakeHostPool: + def __init__(self, alloc_result): + self.alloc_result = alloc_result + self.alloc_calls = [] + self.backups = [] + self.loads = [] + self.frees = [] + self.page_size = 4 + self.layout = "page_first_direct" + + def alloc(self, need_size): + self.alloc_calls.append(need_size) + if self.alloc_result is None: + return None + return self.alloc_result[:need_size].clone() + + def backup_from_device_all_layer( + self, device_pool, host_indices, device_indices, io_backend + ): + self.backups.append((host_indices.clone(), device_indices.clone())) + + def load_to_device_per_layer( + self, device_pool, host_indices, device_indices, layer_id, io_backend + ): + self.loads.append((host_indices.clone(), device_indices.clone(), layer_id)) + + def free(self, indices): + self.frees.append(indices.clone()) + return len(indices) + + +class FakeDevicePool: + device = "cpu" + layer_num = 1 + + def register_layer_transfer_counter(self, counter): + self.counter = counter + + +class FakeAllocator: + def __init__(self, alloc_result=None): + self.alloc_result = alloc_result + self.alloc_calls = [] + self.cp_size = 4 + self.cp_rank = 1 + + def get_kvcache(self): + return FakeDevicePool() + + def alloc(self, need_size): + self.alloc_calls.append(need_size) + if self.alloc_result is None: + return None + return self.alloc_result[:need_size].clone() + + +class DummyEvent: + def record(self): + pass + + def wait(self, stream): + pass + + def query(self): + return True + + def synchronize(self): + pass + + +class DummyStream: + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, tb): + return False + + +class DummyDeviceModule: + Event = DummyEvent + Stream = DummyStream + + @staticmethod + def stream(stream): + return stream + + +class DummyLayerDoneCounter: + def __init__(self): + self.events = [ + type( + "ProducerEvent", + (), + { + "start_event": DummyEvent(), + "finish_event": DummyEvent(), + "complete": lambda self, layer_id: None, + }, + )() + ] + + def update_producer(self): + return 0 + + +class TestHiCacheControllerCPWrite(CustomTestCase): + def setUp(self): + self.device_module_patcher = patch( + "sglang.srt.managers.cache_controller.device_module", + DummyDeviceModule, + ) + self.nsa_pool_patcher = patch( + "sglang.srt.managers.cache_controller.NSATokenToKVPool", + FakeDevicePool, + ) + self.device_module_patcher.start() + self.nsa_pool_patcher.start() + self.addCleanup(self.device_module_patcher.stop) + self.addCleanup(self.nsa_pool_patcher.stop) + + def make_controller(self, host_pool, allocator=None, cp_rank=1): + allocator = allocator or FakeAllocator() + controller = HiCacheController( + token_to_kv_pool_allocator=allocator, + mem_pool_host=host_pool, + page_size=4, + tp_group=None, + load_cache_event=__import__("threading").Event(), + io_backend="direct", + cp_shared_kv_layout=CpSharedKVLayout( + page_size=4, cp_size=4, cp_rank=cp_rank + ), + ) + controller.layer_done_counter = DummyLayerDoneCounter() + return controller + + def test_cp_write_filters_to_owned_physical_locs(self): + host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64)) + controller = self.make_controller(host_pool, cp_rank=1) + logical_locs = torch.arange(4, 20, dtype=torch.int64) + + result = controller.write(logical_locs, node_id=7) + + self.assertEqual(result.metadata.logical_len, 16) + self.assertEqual(result.metadata.owned_positions.tolist(), [4, 5, 6, 7]) + self.assertEqual(result.metadata.host_indices.tolist(), [100, 101, 102, 103]) + self.assertEqual(host_pool.alloc_calls, [4]) + self.assertEqual(host_pool.backups[0][1].tolist(), [4, 5, 6, 7]) + + def test_cp_write_zero_owned_returns_metadata_and_noop_ack(self): + host_pool = FakeHostPool(torch.tensor([], dtype=torch.int64)) + controller = self.make_controller(host_pool, cp_rank=3) + logical_locs = torch.arange(4, 8, dtype=torch.int64) + + result = controller.write(logical_locs, node_id=8) + + self.assertEqual(result.metadata.logical_len, 4) + self.assertEqual(result.metadata.host_indices.tolist(), []) + self.assertEqual(host_pool.alloc_calls, []) + self.assertEqual(len(controller.ack_write_queue), 1) + + def test_cp_write_allocation_failure_reports_required_host_slots(self): + host_pool = FakeHostPool(None) + controller = self.make_controller(host_pool, cp_rank=1) + logical_locs = torch.arange(4, 20, dtype=torch.int64) + + result = controller.write(logical_locs, node_id=9) + + self.assertEqual(result.required_host_slots, 4) + + +if __name__ == "__main__": + unittest.main()