feat: map CP HiCache writes to owned physical slots
This commit is contained in:
@@ -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,
|
||||
|
||||
210
test/registered/unit/managers/test_hicache_controller_cp.py
Normal file
210
test/registered/unit/managers/test_hicache_controller_cp.py
Normal file
@@ -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()
|
||||
Reference in New Issue
Block a user