chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity
Files
f355fdd39e96b640a044f8033d46cd4632c5bc49
sglang/test/registered/unit/managers/test_hicache_controller_cp.py
laoyao0822 f355fdd39e Overlap CP HiCache backup without exposing partial host state 
CP shared KV with HiCache and EAGLE needs host backup to overlap forward while keeping radix visibility synchronous. The change reserves host slots before forward, drives target and draft backup from explicit layer-end hooks, and commits host visibility only after the final target/draft ack. It also probes the final insertion prefix before early reservation so repeated EAGLE prompts do not prepare duplicate suffix backups that later rollback as insert_miss.

Constraint: CP ranks use independent shared-KV pools, so target/draft host state must remain atomically visible at the radix boundary.

Constraint: Fused MLA and NSA store paths can bypass store-side notifier hooks, so layer end is the safer backup progress boundary.

Rejected: Store-side backup notifier as the primary trigger | fused store and zero-local paths made notifier coverage fragile.

Rejected: Reserve from cache_protected_len alone | EAGLE bigram/page alignment can make final insertion find a longer existing prefix and force duplicate rollback work.

Confidence: medium

Scope-risk: moderate

Directive: Do not add per-layer CP collectives here; keep radix state synchronous and data transfer asynchronous/local-event driven.

Tested: local git diff --check

Tested: local py_compile for touched CP HiCache/cache-controller/deepseek/test files

Tested: remote pytest test/registered/unit/mem_cache/test_cp_hicache_metadata.py test/registered/unit/managers/test_hicache_controller_cp.py -q (115 passed, 5 warnings)

Not-tested: full GLM5 ETE server rerun after this commit
2026-05-27 09:50:47 +08:00

1053 lines
43 KiB
Python
Raw Blame History

import sys
import types
from unittest import main
from unittest.mock import patch
import torch
try:
import sgl_kernel # noqa: F401
import sgl_kernel.kvcacheio # noqa: F401
except (ImportError, RuntimeError):
if "sgl_kernel" not in sys.modules:
sys.modules["sgl_kernel"] = types.ModuleType("sgl_kernel")
sys.modules["sgl_kernel"].__file__ = "sgl_kernel_stub.py"
sys.modules["sgl_kernel"].__path__ = []
if not hasattr(sys.modules["sgl_kernel"], "__getattr__"):
def _sgl_kernel_getattr(name):
if name.startswith("__"):
raise AttributeError(name)
fn = lambda *args, **kwargs: None
setattr(sys.modules["sgl_kernel"], name, fn)
return fn
sys.modules["sgl_kernel"].__getattr__ = _sgl_kernel_getattr
if "sgl_kernel.quantization" not in sys.modules:
quantization_stub = types.ModuleType("sgl_kernel.quantization")
quantization_stub.__file__ = "sgl_kernel_quantization_stub.py"
def _quantization_getattr(name):
if name.startswith("__"):
raise AttributeError(name)
fn = lambda *args, **kwargs: None
setattr(quantization_stub, name, fn)
return fn
quantization_stub.__getattr__ = _quantization_getattr
for name in (
"ggml_dequantize",
"ggml_moe_a8",
"ggml_moe_a8_vec",
"ggml_moe_get_block_size",
"ggml_mul_mat_a8",
"ggml_mul_mat_vec_a8",
):
setattr(quantization_stub, name, lambda *args, **kwargs: None)
sys.modules["sgl_kernel.quantization"] = quantization_stub
for name in (
"sgl_per_token_group_quant_8bit",
"sgl_per_token_group_quant_fp8",
"sgl_per_token_quant_fp8",
"fp8_blockwise_scaled_mm",
"fp8_scaled_mm",
"silu_and_mul",
):
if not hasattr(sys.modules["sgl_kernel"], name):
setattr(sys.modules["sgl_kernel"], name, lambda *args, **kwargs: None)
_sgl_kernel_lib = torch.library.Library("sgl_kernel", "FRAGMENT")
for _schema in (
"sgl_per_token_group_quant_8bit(Tensor input, Tensor(a!) output_q, Tensor(b!) output_s, int group_size, float eps, float fp8_min, float fp8_max, bool scale_ue8m0) -> ()",
"sgl_per_token_group_quant_fp8(Tensor input, Tensor(a!) output_q, Tensor(b!) output_s, int group_size, float eps, float fp8_min, float fp8_max, bool scale_ue8m0) -> ()",
"sgl_per_token_quant_fp8(Tensor input, Tensor(a!) output_q, Tensor(b!) output_s) -> ()",
"fp8_scaled_mm(Tensor mat_a, Tensor mat_b, Tensor scales_a, Tensor scales_b, ScalarType out_dtype, Tensor? bias=None) -> Tensor",
"fp8_blockwise_scaled_mm(Tensor mat_a, Tensor mat_b, Tensor scales_a, Tensor scales_b, ScalarType out_dtype) -> Tensor",
):
try:
_sgl_kernel_lib.define(_schema)
except RuntimeError as exc:
if (
"already" not in str(exc).lower()
and "duplicate" not in str(exc).lower()
):
raise
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
_sgl_kernel_lib = torch.library.Library("sgl_kernel", "FRAGMENT")
for _schema in (
"sgl_per_token_group_quant_8bit(Tensor input, Tensor(a!) output_q, Tensor(b!) output_s, int group_size, float eps, float fp8_min, float fp8_max, bool scale_ue8m0) -> ()",
"sgl_per_token_group_quant_fp8(Tensor input, Tensor(a!) output_q, Tensor(b!) output_s, int group_size, float eps, float fp8_min, float fp8_max, bool scale_ue8m0) -> ()",
"sgl_per_token_quant_fp8(Tensor input, Tensor(a!) output_q, Tensor(b!) output_s) -> ()",
"fp8_scaled_mm(Tensor mat_a, Tensor mat_b, Tensor scales_a, Tensor scales_b, ScalarType out_dtype, Tensor? bias=None) -> Tensor",
"fp8_blockwise_scaled_mm(Tensor mat_a, Tensor mat_b, Tensor scales_a, Tensor scales_b, ScalarType out_dtype) -> Tensor",
):
try:
_sgl_kernel_lib.define(_schema)
except RuntimeError as exc:
if "already" not in str(exc).lower() and "duplicate" not in str(exc).lower():
raise
from sglang.srt.managers.cache_controller import HiCacheController
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
from sglang.srt.mem_cache.hiradix_cache import CpHiCacheNodeMetadata
from sglang.srt.mem_cache.memory_pool_host import (
MLATokenToKVPoolHost,
NSATokenToKVPoolHost,
)
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
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.layer_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(), device_pool))
def backup_from_device_per_layer(
self, device_pool, host_indices, device_indices, layer_id, io_backend
):
self.layer_backups.append(
(host_indices.clone(), device_indices.clone(), layer_id, device_pool)
)
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, device_pool)
)
def free(self, indices):
self.frees.append(indices.clone())
return len(indices)
class FakeDevicePool:
device = "cpu"
layer_num = 1
def __init__(self, name="target", layer_num=1):
self.name = name
self.layer_num = layer_num
self.layer_backup_notifiers = []
def register_layer_transfer_counter(self, counter):
self.counter = counter
def register_layer_backup_notifier(self, notifier):
self.layer_backup_notifiers.append(notifier)
def notify_layer_end_for_backup(self, layer_id):
for notifier in self.layer_backup_notifiers:
notifier(layer_id)
class TestPageFirstPerLayerBackupTaiKernel(CustomTestCase):
def test_mla_page_first_per_layer_backup_uses_tai_lf_pf_kernel(self):
calls = []
def fake_kernel(src, dst, src_indices, dst_indices, **kwargs):
calls.append((src, dst, src_indices.clone(), dst_indices.clone(), kwargs))
host_pool = MLATokenToKVPoolHost.__new__(MLATokenToKVPoolHost)
host_pool.layout = "page_first"
host_pool.token_stride_size = 16
host_pool.layout_dim = 64
host_pool.kv_buffer = torch.empty((32, 4, 1, 16), dtype=torch.uint8)
device_pool = type("DevicePool", (), {})()
device_pool.kv_buffer = torch.empty((4, 32, 1, 16), dtype=torch.uint8)
host_indices = torch.tensor([4, 5, 6, 7], dtype=torch.int64)
device_indices = torch.tensor([12, 13, 14, 15], dtype=torch.int64)
with patch(
"sglang.srt.mem_cache.memory_pool_host._load_tai_transfer_kv_per_layer_mla_lf_pf",
return_value=fake_kernel,
):
host_pool.backup_from_device_per_layer(
device_pool,
host_indices,
device_indices,
layer_id=2,
io_backend="kernel",
)
self.assertEqual(len(calls), 1)
src, dst, src_indices, dst_indices, kwargs = calls[0]
expected_src = device_pool.kv_buffer[2]
self.assertEqual(src.data_ptr(), expected_src.data_ptr())
self.assertEqual(src.shape, expected_src.shape)
self.assertEqual(src.stride(), expected_src.stride())
self.assertIs(dst, host_pool.kv_buffer)
self.assertEqual(src_indices.tolist(), [12, 13, 14, 15])
self.assertEqual(dst_indices.tolist(), [4, 5, 6, 7])
self.assertEqual(kwargs["layer_id"], 2)
self.assertEqual(kwargs["item_size"], 16)
self.assertEqual(kwargs["dst_layout_dim"], 64)
def test_nsa_indexer_page_first_per_layer_backup_uses_tai_lf_pf_kernel(self):
calls = []
def fake_kernel(src, dst, src_indices, dst_indices, **kwargs):
calls.append((src, dst, src_indices.clone(), dst_indices.clone(), kwargs))
host_pool = NSATokenToKVPoolHost.__new__(NSATokenToKVPoolHost)
host_pool.layout = "page_first"
host_pool.page_size = 4
host_pool.indexer_page_stride_size = 32
host_pool.indexer_layout_dim = 96
host_pool.index_k_with_scale_buffer = torch.empty(
(16, 3, 1, 32), dtype=torch.uint8
)
device_pool = type("DevicePool", (), {})()
device_pool.index_k_with_scale_buffer = torch.empty(
(3, 16, 32), dtype=torch.uint8
)
host_indices = torch.tensor([8, 9, 10, 11, 20, 21, 22, 23], dtype=torch.int64)
device_indices = torch.tensor(
[12, 13, 14, 15, 28, 29, 30, 31], dtype=torch.int64
)
with patch(
"sglang.srt.mem_cache.memory_pool_host._load_tai_transfer_kv_per_layer_mla_lf_pf",
return_value=fake_kernel,
):
host_pool._backup_indexer_from_device_per_layer(
device_pool,
host_indices,
device_indices,
layer_id=1,
io_backend="kernel",
)
self.assertEqual(len(calls), 1)
src, dst, src_indices, dst_indices, kwargs = calls[0]
expected_src = device_pool.index_k_with_scale_buffer[1]
self.assertEqual(src.data_ptr(), expected_src.data_ptr())
self.assertEqual(src.shape, expected_src.shape)
self.assertEqual(src.stride(), expected_src.stride())
self.assertIs(dst, host_pool.index_k_with_scale_buffer)
self.assertEqual(src_indices.tolist(), [3, 7])
self.assertEqual(dst_indices.tolist(), [2, 5])
self.assertEqual(kwargs["layer_id"], 1)
self.assertEqual(kwargs["item_size"], 32)
self.assertEqual(kwargs["dst_layout_dim"], 96)
class FakeAllocator:
def __init__(self, alloc_result=None):
self.alloc_result = alloc_result
self.alloc_calls = []
self.owner_alloc_calls = []
self.frees = []
self.cp_size = 4
self.cp_rank = 1
self.page_size = 4
self.device_pool = FakeDevicePool()
def get_kvcache(self):
return self.device_pool
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 alloc_pages_with_owners(self, page_owners):
owners = list(page_owners)
self.owner_alloc_calls.append(owners)
if self.alloc_result is None:
return None
need_size = len(owners) * self.page_size
return self.alloc_result[:need_size].clone()
def free(self, indices):
self.frees.append(indices.clone())
return len(indices)
class HostIndicesTensor(torch.Tensor):
@staticmethod
def __new__(cls, data):
return torch.Tensor._make_subclass(cls, data, require_grad=False)
def to(self, *args, **kwargs):
raise AssertionError("load_cp should not move host indices before queuing")
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 RecordingProducerEvent:
def __init__(self, order):
self.start_event = DummyEvent()
self.finish_event = DummyEvent()
self.order = order
def complete(self, layer_id):
self.order.append(("complete", layer_id))
class RecordingLayerDoneCounter:
def __init__(self, order):
self.events = [RecordingProducerEvent(order)]
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,
draft_host_pool=None,
draft_mem_pool_device=None,
):
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
),
draft_mem_pool_host=draft_host_pool,
draft_mem_pool_device=draft_mem_pool_device,
)
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, [])
self.assertEqual(host_pool.layer_backups[0][1].tolist(), [4, 5, 6, 7])
def test_cp_write_rejects_incomplete_owned_physical_page(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.tensor([8, 9, 10], dtype=torch.int64)
with self.assertRaisesRegex(
ValueError, "_write_cp expects page-aligned device_indices"
):
controller.write(logical_locs, node_id=21)
def test_cp_write_rejects_non_contiguous_owned_physical_page(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.tensor([8, 9, 11, 10], dtype=torch.int64)
with self.assertRaisesRegex(
ValueError, "physical_device_indices.*contiguous page spans"
):
controller.write(logical_locs, node_id=22)
def test_cp_write_rejects_non_contiguous_host_page(self):
host_pool = FakeHostPool(torch.tensor([100, 101, 103, 102], dtype=torch.int64))
controller = self.make_controller(host_pool, cp_rank=1)
logical_locs = torch.arange(8, 12, dtype=torch.int64)
with self.assertRaisesRegex(ValueError, "host_indices.*contiguous page spans"):
controller.write(logical_locs, node_id=23)
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_zero_owned_with_draft_returns_empty_draft_metadata(self):
host_pool = FakeHostPool(torch.tensor([], dtype=torch.int64))
draft_host_pool = FakeHostPool(torch.tensor([], dtype=torch.int64))
controller = self.make_controller(
host_pool,
cp_rank=3,
draft_host_pool=draft_host_pool,
draft_mem_pool_device=FakeDevicePool("draft"),
)
logical_locs = torch.arange(4, 8, dtype=torch.int64)
result = controller.write(logical_locs, node_id=18)
self.assertEqual(result.metadata.logical_len, 4)
self.assertEqual(result.metadata.host_indices.tolist(), [])
self.assertEqual(result.metadata.draft_host_indices.tolist(), [])
self.assertEqual(host_pool.alloc_calls, [])
self.assertEqual(draft_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)
def test_cp_write_with_draft_pool_backs_target_and_draft_locs(self):
host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
draft_host_pool = FakeHostPool(
torch.tensor([200, 201, 202, 203], dtype=torch.int64)
)
draft_device_pool = FakeDevicePool("draft")
controller = self.make_controller(
host_pool,
cp_rank=1,
draft_host_pool=draft_host_pool,
draft_mem_pool_device=draft_device_pool,
)
logical_locs = torch.arange(4, 20, dtype=torch.int64)
result = controller.write(logical_locs, node_id=77)
self.assertEqual(result.metadata.host_indices.tolist(), [100, 101, 102, 103])
self.assertEqual(
result.metadata.draft_host_indices.tolist(), [200, 201, 202, 203]
)
self.assertEqual(host_pool.alloc_calls, [4])
self.assertEqual(draft_host_pool.alloc_calls, [4])
self.assertEqual(host_pool.backups, [])
self.assertEqual(draft_host_pool.backups, [])
self.assertEqual(host_pool.layer_backups[0][1].tolist(), [4, 5, 6, 7])
self.assertEqual(draft_host_pool.layer_backups[0][1].tolist(), [4, 5, 6, 7])
self.assertIs(draft_host_pool.layer_backups[0][3], draft_device_pool)
self.assertEqual(len(controller.ack_write_queue), 1)
self.assertEqual(controller.ack_write_queue[0].node_ids, [77])
def test_cp_write_draft_allocation_failure_rolls_back_target_host(self):
host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
draft_host_pool = FakeHostPool(None)
controller = self.make_controller(
host_pool,
cp_rank=1,
draft_host_pool=draft_host_pool,
draft_mem_pool_device=FakeDevicePool("draft"),
)
logical_locs = torch.arange(4, 20, dtype=torch.int64)
result = controller.write(logical_locs, node_id=78)
self.assertIsNone(result.metadata)
self.assertEqual(result.required_host_slots, 4)
self.assertEqual(host_pool.frees[0].tolist(), [100, 101, 102, 103])
self.assertEqual(host_pool.backups, [])
self.assertEqual(draft_host_pool.backups, [])
def test_cp_reserve_write_queues_no_transfer_until_submit(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)
reservation = controller.reserve_write_cp(logical_locs, node_id=79)
self.assertEqual(reservation.metadata.logical_len, 16)
self.assertEqual(host_pool.alloc_calls, [4])
self.assertEqual(host_pool.backups, [])
self.assertEqual(controller.write_queue, [])
self.assertEqual(controller.ack_write_queue, [])
with self.assertLogs(
"sglang.srt.managers.cache_controller", level="WARNING"
) as logs:
controller.submit_write_cp_all_layer(reservation)
self.assertEqual(host_pool.backups[0][1].tolist(), [4, 5, 6, 7])
self.assertEqual(len(controller.ack_write_queue), 1)
self.assertEqual(controller.ack_write_queue[0].node_ids, [79])
self.assertIn("all-layer backup fallback", "\n".join(logs.output))
def test_cp_reserve_zero_owned_queues_no_ack_until_submit(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)
reservation = controller.reserve_write_cp(logical_locs, node_id=80)
self.assertEqual(reservation.metadata.host_indices.tolist(), [])
self.assertEqual(host_pool.alloc_calls, [])
self.assertEqual(controller.ack_write_queue, [])
with self.assertLogs(
"sglang.srt.managers.cache_controller", level="WARNING"
) as logs:
controller.submit_write_cp_all_layer(reservation)
self.assertEqual(len(controller.ack_write_queue), 1)
self.assertEqual(controller.ack_write_queue[0].node_ids, [80])
self.assertEqual(host_pool.backups, [])
self.assertIn("all-layer backup fallback", "\n".join(logs.output))
def test_cp_reserve_draft_allocation_failure_rolls_back_without_transfer(self):
host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
draft_host_pool = FakeHostPool(None)
controller = self.make_controller(
host_pool,
cp_rank=1,
draft_host_pool=draft_host_pool,
draft_mem_pool_device=FakeDevicePool("draft"),
)
logical_locs = torch.arange(4, 20, dtype=torch.int64)
result = controller.reserve_write_cp(logical_locs, node_id=81)
self.assertIsNone(result.metadata)
self.assertEqual(result.required_host_slots, 4)
self.assertEqual(host_pool.frees[0].tolist(), [100, 101, 102, 103])
self.assertEqual(host_pool.backups, [])
self.assertEqual(draft_host_pool.backups, [])
self.assertEqual(controller.write_queue, [])
self.assertEqual(controller.draft_write_queue, [])
def test_cp_submit_write_cp_layer_pairs_target_and_draft_with_single_final_ack(self):
host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
draft_host_pool = FakeHostPool(
torch.tensor([200, 201, 202, 203], dtype=torch.int64)
)
allocator = FakeAllocator()
allocator.device_pool = FakeDevicePool("target", layer_num=2)
draft_device_pool = FakeDevicePool("draft", layer_num=2)
controller = self.make_controller(
host_pool,
allocator=allocator,
cp_rank=1,
draft_host_pool=draft_host_pool,
draft_mem_pool_device=draft_device_pool,
)
logical_locs = torch.arange(4, 20, dtype=torch.int64)
reservation = controller.reserve_write_cp(logical_locs, node_id=82)
controller.submit_write_cp_layer(reservation, 0)
self.assertEqual(controller.ack_write_queue, [])
self.assertEqual(host_pool.layer_backups[0][1].tolist(), [4, 5, 6, 7])
self.assertEqual(draft_host_pool.layer_backups[0][1].tolist(), [4, 5, 6, 7])
self.assertEqual(host_pool.layer_backups[0][2], 0)
self.assertEqual(draft_host_pool.layer_backups[0][2], 0)
controller.submit_write_cp_layer(reservation, 1)
self.assertEqual(len(controller.ack_write_queue), 1)
self.assertEqual(controller.ack_write_queue[0].node_ids, [82])
self.assertEqual([x[2] for x in host_pool.layer_backups], [0, 1])
self.assertEqual([x[2] for x in draft_host_pool.layer_backups], [0, 1])
def test_cp_submit_write_cp_layer_zero_owned_final_ack_once(self):
host_pool = FakeHostPool(torch.tensor([], dtype=torch.int64))
allocator = FakeAllocator()
allocator.device_pool = FakeDevicePool("target", layer_num=2)
controller = self.make_controller(host_pool, allocator=allocator, cp_rank=3)
logical_locs = torch.arange(4, 8, dtype=torch.int64)
reservation = controller.reserve_write_cp(logical_locs, node_id=83)
controller.submit_write_cp_layer(reservation, 0)
self.assertEqual(controller.ack_write_queue, [])
controller.submit_write_cp_layer(reservation, 1)
self.assertEqual(len(controller.ack_write_queue), 1)
self.assertEqual(controller.ack_write_queue[0].node_ids, [83])
self.assertEqual(host_pool.layer_backups, [])
def test_cp_layer_hook_submits_registered_write_without_all_layer_backup(self):
host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
allocator = FakeAllocator()
allocator.device_pool = FakeDevicePool("target", layer_num=2)
controller = self.make_controller(host_pool, allocator=allocator, cp_rank=1)
logical_locs = torch.arange(4, 20, dtype=torch.int64)
reservation = controller.reserve_write_cp(logical_locs, node_id=84)
controller.submit_write_cp_per_layer(reservation, catch_up_all_layers=False)
self.assertEqual(host_pool.backups, [])
self.assertEqual(host_pool.layer_backups, [])
self.assertEqual(controller.ack_write_queue, [])
allocator.device_pool.notify_layer_end_for_backup(0)
self.assertEqual(host_pool.layer_backups[0][2], 0)
self.assertEqual(controller.ack_write_queue, [])
allocator.device_pool.notify_layer_end_for_backup(1)
self.assertEqual([x[2] for x in host_pool.layer_backups], [0, 1])
self.assertEqual(len(controller.ack_write_queue), 1)
self.assertEqual(controller.ack_write_queue[0].node_ids, [84])
self.assertEqual(host_pool.backups, [])
def test_cp_layer_hook_waits_for_draft_source_before_final_ack(self):
host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
draft_host_pool = FakeHostPool(
torch.tensor([200, 201, 202, 203], dtype=torch.int64)
)
allocator = FakeAllocator()
allocator.device_pool = FakeDevicePool("target", layer_num=1)
draft_device_pool = FakeDevicePool("draft", layer_num=1)
controller = self.make_controller(
host_pool,
allocator=allocator,
cp_rank=1,
draft_host_pool=draft_host_pool,
draft_mem_pool_device=draft_device_pool,
)
logical_locs = torch.arange(4, 20, dtype=torch.int64)
reservation = controller.reserve_write_cp(logical_locs, node_id=85)
controller.submit_write_cp_per_layer(reservation, catch_up_all_layers=False)
allocator.device_pool.notify_layer_end_for_backup(0)
self.assertEqual([x[2] for x in host_pool.layer_backups], [0])
self.assertEqual(draft_host_pool.layer_backups, [])
self.assertEqual(controller.ack_write_queue, [])
draft_device_pool.notify_layer_end_for_backup(0)
self.assertEqual([x[2] for x in draft_host_pool.layer_backups], [0])
self.assertEqual(len(controller.ack_write_queue), 1)
self.assertEqual(controller.ack_write_queue[0].node_ids, [85])
def test_generate_storage_config_constructs_config_at_runtime(self):
controller = HiCacheController.__new__(HiCacheController)
controller.mem_pool_device = FakeDevicePool()
controller.mem_pool_host = FakeHostPool(torch.tensor([], dtype=torch.int64))
controller.pp_rank = 1
controller.pp_size = 2
controller.enable_storage_metrics = True
with patch(
"sglang.srt.managers.cache_controller.is_dp_attention_enabled",
return_value=False,
), patch(
"sglang.srt.managers.cache_controller.get_tensor_model_parallel_rank",
return_value=3,
), patch(
"sglang.srt.managers.cache_controller.get_tensor_model_parallel_world_size",
return_value=4,
):
config = controller._generate_storage_config(
model_name="test-model",
storage_backend_extra_config={"tp_lcm_size": 8},
)
self.assertEqual(config.tp_rank, 3)
self.assertEqual(config.tp_size, 4)
self.assertEqual(config.pp_rank, 1)
self.assertEqual(config.pp_size, 2)
self.assertEqual(config.model_name, "test-model")
self.assertEqual(config.tp_lcm_size, 8)
def test_attach_storage_backend_rejects_cp_hicache(self):
host_pool = FakeHostPool(torch.tensor([], dtype=torch.int64))
controller = self.make_controller(host_pool)
with self.assertRaisesRegex(RuntimeError, "CP shared KV.*storage backend"):
controller.attach_storage_backend("mooncake")
class TestHiCacheControllerCPLoad(TestHiCacheControllerCPWrite):
def test_cp_load_allocates_full_logical_locs_and_transfers_owned_physical_locs(self):
host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
allocator = FakeAllocator(alloc_result=torch.arange(64, 80, dtype=torch.int64))
controller = self.make_controller(host_pool, allocator=allocator, cp_rank=1)
node = TreeNode()
node.host_len = 16
node.cp_hicache = CpHiCacheNodeMetadata(
logical_len=16,
owned_positions=torch.tensor([4, 5, 6, 7], dtype=torch.int64),
host_indices=torch.tensor([100, 101, 102, 103], dtype=torch.int64),
page_owners=torch.tensor([3, 0, 1, 2], dtype=torch.int8),
page_size=4,
)
device_indices = controller.load_cp([node], node_id=11)
controller.start_loading()
self.assertEqual(device_indices.tolist(), list(range(64, 80)))
self.assertEqual(allocator.alloc_calls, [])
self.assertEqual(allocator.owner_alloc_calls, [[3, 0, 1, 2]])
self.assertEqual(host_pool.loads[0][1].tolist(), [20, 21, 22, 23])
def test_cp_load_frees_unexpected_owner_allocator_length(self):
host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
allocator = FakeAllocator(alloc_result=torch.arange(64, 76, dtype=torch.int64))
controller = self.make_controller(host_pool, allocator=allocator, cp_rank=1)
node = TreeNode()
node.host_len = 16
node.cp_hicache = CpHiCacheNodeMetadata(
logical_len=16,
owned_positions=torch.tensor([4, 5, 6, 7], dtype=torch.int64),
host_indices=torch.tensor([100, 101, 102, 103], dtype=torch.int64),
page_owners=torch.tensor([3, 0, 1, 2], dtype=torch.int8),
page_size=4,
)
with self.assertRaisesRegex(
RuntimeError, "alloc_pages_with_owners returned unexpected length"
):
controller.load_cp([node], node_id=111)
self.assertEqual(allocator.owner_alloc_calls, [[3, 0, 1, 2]])
self.assertEqual(allocator.frees[0].tolist(), list(range(64, 76)))
def test_cp_load_with_draft_pool_restores_target_and_draft_locs(self):
host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
draft_host_pool = FakeHostPool(
torch.tensor([200, 201, 202, 203], dtype=torch.int64)
)
draft_device_pool = FakeDevicePool("draft")
allocator = FakeAllocator(alloc_result=torch.arange(64, 80, dtype=torch.int64))
controller = self.make_controller(
host_pool,
allocator=allocator,
cp_rank=1,
draft_host_pool=draft_host_pool,
draft_mem_pool_device=draft_device_pool,
)
node = TreeNode()
node.host_len = 16
node.cp_hicache = CpHiCacheNodeMetadata(
logical_len=16,
owned_positions=torch.tensor([4, 5, 6, 7], dtype=torch.int64),
host_indices=torch.tensor([100, 101, 102, 103], dtype=torch.int64),
draft_host_indices=torch.tensor([200, 201, 202, 203], dtype=torch.int64),
page_owners=torch.tensor([3, 0, 1, 2], dtype=torch.int8),
page_size=4,
)
device_indices = controller.load_cp([node], node_id=14)
controller.start_loading()
self.assertEqual(device_indices.tolist(), list(range(64, 80)))
self.assertEqual(allocator.owner_alloc_calls, [[3, 0, 1, 2]])
self.assertEqual(host_pool.loads[0][1].tolist(), [20, 21, 22, 23])
self.assertEqual(draft_host_pool.loads[0][1].tolist(), [20, 21, 22, 23])
self.assertIs(draft_host_pool.loads[0][3], draft_device_pool)
self.assertEqual(len(controller.ack_load_queue), 1)
self.assertEqual(controller.ack_load_queue[0].node_ids, [14])
def test_cp_start_loading_loads_draft_before_target_layer_ready(self):
order = []
host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
draft_host_pool = FakeHostPool(
torch.tensor([200, 201, 202, 203], dtype=torch.int64)
)
allocator = FakeAllocator(alloc_result=torch.arange(64, 80, dtype=torch.int64))
controller = self.make_controller(
host_pool,
allocator=allocator,
cp_rank=1,
draft_host_pool=draft_host_pool,
draft_mem_pool_device=FakeDevicePool("draft"),
)
controller.layer_done_counter = RecordingLayerDoneCounter(order)
def record_target_load(
device_pool, host_indices, device_indices, layer_id, io_backend
):
order.append(("target", layer_id))
def record_draft_load(
device_pool, host_indices, device_indices, layer_id, io_backend
):
order.append(("draft", layer_id))
host_pool.load_to_device_per_layer = record_target_load
draft_host_pool.load_to_device_per_layer = record_draft_load
node = TreeNode()
node.host_len = 16
node.cp_hicache = CpHiCacheNodeMetadata(
logical_len=16,
owned_positions=torch.tensor([4, 5, 6, 7], dtype=torch.int64),
host_indices=torch.tensor([100, 101, 102, 103], dtype=torch.int64),
draft_host_indices=torch.tensor([200, 201, 202, 203], dtype=torch.int64),
page_owners=torch.tensor([3, 0, 1, 2], dtype=torch.int8),
page_size=4,
)
controller.load_cp([node], node_id=114)
controller.start_loading()
self.assertEqual(order, [("draft", 0), ("target", 0), ("complete", 0)])
def test_cp_load_zero_owned_returns_full_logical_locs_and_noop_ack(self):
host_pool = FakeHostPool(torch.tensor([], dtype=torch.int64))
allocator = FakeAllocator(alloc_result=torch.arange(64, 68, dtype=torch.int64))
controller = self.make_controller(host_pool, allocator=allocator, cp_rank=3)
node = TreeNode()
node.host_len = 4
node.cp_hicache = CpHiCacheNodeMetadata(
logical_len=4,
owned_positions=torch.empty((0,), dtype=torch.int64),
host_indices=torch.empty((0,), dtype=torch.int64),
page_owners=torch.tensor([0], dtype=torch.int8),
page_size=4,
)
device_indices = controller.load_cp([node], node_id=12)
controller.start_loading()
self.assertEqual(device_indices.tolist(), [64, 65, 66, 67])
self.assertEqual(allocator.owner_alloc_calls, [[0]])
self.assertEqual(host_pool.loads, [])
self.assertEqual(len(controller.ack_load_queue), 1)
def test_cp_load_zero_owned_rejects_missing_draft_metadata_when_draft_attached(self):
host_pool = FakeHostPool(torch.tensor([], dtype=torch.int64))
allocator = FakeAllocator(alloc_result=torch.arange(64, 68, dtype=torch.int64))
controller = self.make_controller(
host_pool,
allocator=allocator,
cp_rank=3,
draft_host_pool=FakeHostPool(torch.tensor([], dtype=torch.int64)),
draft_mem_pool_device=FakeDevicePool("draft"),
)
node = TreeNode()
node.host_len = 4
node.cp_hicache = CpHiCacheNodeMetadata(
logical_len=4,
owned_positions=torch.empty((0,), dtype=torch.int64),
host_indices=torch.empty((0,), dtype=torch.int64),
page_owners=torch.tensor([0], dtype=torch.int8),
page_size=4,
)
with self.assertRaisesRegex(RuntimeError, "draft KV restore requested"):
controller.load_cp([node], node_id=33)
self.assertEqual(allocator.owner_alloc_calls, [[0]])
self.assertEqual(allocator.frees[0].tolist(), [64, 65, 66, 67])
self.assertEqual(controller.load_queue, [])
self.assertEqual(controller.draft_load_queue, [])
def test_cp_load_queues_cpu_host_indices_before_backend_moves(self):
host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
allocator = FakeAllocator(alloc_result=torch.arange(64, 80, dtype=torch.int64))
controller = self.make_controller(host_pool, allocator=allocator, cp_rank=1)
host_indices = HostIndicesTensor(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
node = TreeNode()
node.host_len = 16
metadata = CpHiCacheNodeMetadata(
logical_len=16,
owned_positions=torch.tensor([4, 5, 6, 7], dtype=torch.int64),
host_indices=torch.tensor([100, 101, 102, 103], dtype=torch.int64),
page_owners=torch.tensor([3, 0, 1, 2], dtype=torch.int8),
page_size=4,
)
metadata.host_indices = host_indices
node.cp_hicache = metadata
controller.load_cp([node], node_id=13)
queued_op = controller.load_queue[0]
self.assertEqual(queued_op.host_indices.device.type, "cpu")
self.assertEqual(queued_op.host_indices.tolist(), [100, 101, 102, 103])
self.assertEqual(queued_op.device_indices.tolist(), [20, 21, 22, 23])
def test_cp_load_rejects_non_contiguous_physical_device_page(self):
host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
allocator = FakeAllocator(
alloc_result=torch.tensor(
[64, 65, 66, 67, 68, 69, 71, 70, 72, 73, 74, 75, 76, 77, 78, 79],
dtype=torch.int64,
)
)
controller = self.make_controller(host_pool, allocator=allocator, cp_rank=1)
node = TreeNode()
node.host_len = 16
node.cp_hicache = CpHiCacheNodeMetadata(
logical_len=16,
owned_positions=torch.tensor([4, 5, 6, 7], dtype=torch.int64),
host_indices=torch.tensor([100, 101, 102, 103], dtype=torch.int64),
page_owners=torch.tensor([3, 0, 1, 2], dtype=torch.int8),
page_size=4,
)
with self.assertRaisesRegex(
ValueError, "physical_device_indices.*contiguous page spans"
):
controller.load_cp([node], node_id=31)
def test_cp_load_rejects_non_contiguous_host_page(self):
host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
allocator = FakeAllocator(alloc_result=torch.arange(64, 80, dtype=torch.int64))
controller = self.make_controller(host_pool, allocator=allocator, cp_rank=1)
node = TreeNode()
node.host_len = 16
node.cp_hicache = CpHiCacheNodeMetadata(
logical_len=16,
owned_positions=torch.tensor([4, 5, 6, 7], dtype=torch.int64),
host_indices=torch.tensor([100, 101, 103, 102], dtype=torch.int64),
page_owners=torch.tensor([3, 0, 1, 2], dtype=torch.int8),
page_size=4,
)
with self.assertRaisesRegex(ValueError, "host_indices.*contiguous page spans"):
controller.load_cp([node], node_id=32)
def test_cp_evict_host_frees_target_and_draft_host_indices(self):
host_pool = FakeHostPool(torch.tensor([], dtype=torch.int64))
draft_host_pool = FakeHostPool(torch.tensor([], dtype=torch.int64))
controller = self.make_controller(
host_pool,
draft_host_pool=draft_host_pool,
draft_mem_pool_device=FakeDevicePool("draft"),
)
metadata = CpHiCacheNodeMetadata(
logical_len=16,
owned_positions=torch.tensor([0, 1, 2, 3], dtype=torch.int64),
host_indices=torch.tensor([100, 101, 102, 103], dtype=torch.int64),
draft_host_indices=torch.tensor([200, 201, 202, 203], dtype=torch.int64),
page_owners=torch.tensor([0, 1, 2, 3], dtype=torch.int8),
page_size=4,
)
freed = controller.evict_cp_host(metadata)
self.assertEqual(freed, 4)
self.assertEqual(host_pool.frees[0].tolist(), [100, 101, 102, 103])
self.assertEqual(draft_host_pool.frees[0].tolist(), [200, 201, 202, 203])
def test_cp_evict_host_rejects_missing_draft_metadata_before_target_free(self):
host_pool = FakeHostPool(torch.tensor([], dtype=torch.int64))
draft_host_pool = FakeHostPool(torch.tensor([], dtype=torch.int64))
controller = self.make_controller(
host_pool,
draft_host_pool=draft_host_pool,
draft_mem_pool_device=FakeDevicePool("draft"),
)
metadata = CpHiCacheNodeMetadata(
logical_len=16,
owned_positions=torch.tensor([0, 1, 2, 3], dtype=torch.int64),
host_indices=torch.tensor([100, 101, 102, 103], dtype=torch.int64),
page_owners=torch.tensor([0, 1, 2, 3], dtype=torch.int8),
page_size=4,
)
with self.assertRaisesRegex(RuntimeError, "draft.*evict.*draft_host_indices"):
controller.evict_cp_host(metadata)
self.assertEqual(host_pool.frees, [])
self.assertEqual(draft_host_pool.frees, [])
if __name__ == "__main__":
main()
Reference in New Issue View Git Blame Copy Permalink