chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity
Files
dd69c5970cc4ed3b163faffdb018ecefa4eae2ec
sglang/test/registered/unit/managers/test_hicache_controller_cp.py
laoyao0822 7e06eaebbf Prove draft HiCache mirrors target valid-tail pages 
Draft HiCache already piggybacks target CP reservations, but the controller coverage only exercised page-aligned nodes. The new tests pin the intended contract for valid-tail nodes: target and draft reserve/load the same padded physical page while scheduler-visible load results keep only valid logical locs.

Constraint: Draft KV is a target mirror; it must not choose an independent valid-only host path.
Rejected: Re-enable draft partial-current reuse in this slice | current safe contract keeps EAGLE/NextN cache-hit draft on full materialization until same-layer padded visibility is proven.
Confidence: medium
Scope-risk: narrow
Directive: Keep draft and target host metadata coupled; do not add draft-only capacity or prefetch decisions.
Tested: local py_compile for test_hicache_controller_cp.py.
Tested: remote g0034 test_hicache_controller_cp.py: 59 passed, 3 warnings.
Not-tested: CUDA E2E runtime for this commit.
Co-authored-by: OmX <omx@oh-my-codex.dev>
2026-05-29 05:43:49 +08:00

1414 lines
59 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 (
MHATokenToKVPoolHost,
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)
def test_mla_page_first_direct_per_layer_backup_uses_direct_lf_pf(self):
calls = []
def fake_direct(src_ptrs, dst_ptrs, src_indices, dst_indices, **kwargs):
calls.append(
(src_ptrs, dst_ptrs, src_indices.clone(), dst_indices.clone(), kwargs)
)
host_pool = MLATokenToKVPoolHost.__new__(MLATokenToKVPoolHost)
host_pool.layout = "page_first_direct"
host_pool.page_size = 4
host_pool.kv_buffer = torch.empty((8, 3, 4, 1, 16), dtype=torch.uint8)
device_pool = type("DevicePool", (), {})()
device_pool.kv_buffer = torch.empty((3, 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_direct_lf_pf",
return_value=fake_direct,
):
host_pool.backup_from_device_per_layer(
device_pool,
host_indices,
device_indices,
layer_id=2,
io_backend="direct",
)
self.assertEqual(len(calls), 1)
src_ptrs, dst_ptrs, src_indices, dst_indices, kwargs = calls[0]
self.assertEqual(len(src_ptrs), 1)
self.assertEqual(src_ptrs[0].data_ptr(), device_pool.kv_buffer[2].data_ptr())
self.assertEqual(len(dst_ptrs), 1)
self.assertEqual(dst_ptrs[0].data_ptr(), host_pool.kv_buffer.data_ptr())
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["page_size"], 4)
def test_mha_page_first_direct_per_layer_backup_uses_direct_lf_pf(self):
calls = []
def fake_direct(src_ptrs, dst_ptrs, src_indices, dst_indices, **kwargs):
calls.append(
(src_ptrs, dst_ptrs, src_indices.clone(), dst_indices.clone(), kwargs)
)
host_pool = MHATokenToKVPoolHost.__new__(MHATokenToKVPoolHost)
host_pool.layout = "page_first_direct"
host_pool.page_size = 4
host_pool.kv_buffer = torch.empty((2, 8, 3, 4, 2, 8), dtype=torch.uint8)
device_pool = type("DevicePool", (), {})()
device_pool.k_buffer = torch.empty((3, 32, 2, 8), dtype=torch.uint8)
device_pool.v_buffer = torch.empty((3, 32, 2, 8), 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_direct_lf_pf",
return_value=fake_direct,
):
host_pool.backup_from_device_per_layer(
device_pool,
host_indices,
device_indices,
layer_id=2,
io_backend="direct",
)
self.assertEqual(len(calls), 1)
src_ptrs, dst_ptrs, src_indices, dst_indices, kwargs = calls[0]
self.assertEqual(len(src_ptrs), 2)
self.assertEqual(src_ptrs[0].data_ptr(), device_pool.k_buffer[2].data_ptr())
self.assertEqual(src_ptrs[1].data_ptr(), device_pool.v_buffer[2].data_ptr())
self.assertEqual(len(dst_ptrs), 2)
self.assertEqual(dst_ptrs[0].data_ptr(), host_pool.k_buffer.data_ptr())
self.assertEqual(dst_ptrs[1].data_ptr(), host_pool.v_buffer.data_ptr())
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["page_size"], 4)
def test_nsa_indexer_page_first_direct_per_layer_backup_uses_direct_lf_pf(self):
calls = []
def fake_direct(src_ptrs, dst_ptrs, src_indices, dst_indices, **kwargs):
calls.append(
(src_ptrs, dst_ptrs, src_indices.clone(), dst_indices.clone(), kwargs)
)
host_pool = NSATokenToKVPoolHost.__new__(NSATokenToKVPoolHost)
host_pool.layout = "page_first_direct"
host_pool.page_size = 4
host_pool.index_k_with_scale_buffer = torch.empty(
(8, 3, 1, 32), dtype=torch.uint8
)
device_pool = type("DevicePool", (), {})()
device_pool.index_k_with_scale_buffer = torch.empty(
(3, 8, 32), 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_direct_lf_pf",
return_value=fake_direct,
):
host_pool._backup_indexer_from_device_per_layer(
device_pool,
host_indices,
device_indices,
layer_id=1,
io_backend="direct",
)
self.assertEqual(len(calls), 1)
src_ptrs, dst_ptrs, src_indices, dst_indices, kwargs = calls[0]
self.assertEqual(len(src_ptrs), 1)
self.assertEqual(
src_ptrs[0].data_ptr(),
device_pool.index_k_with_scale_buffer[1].data_ptr(),
)
self.assertEqual(len(dst_ptrs), 1)
self.assertEqual(
dst_ptrs[0].data_ptr(),
host_pool.index_k_with_scale_buffer.data_ptr(),
)
self.assertEqual(src_indices.tolist(), [3])
self.assertEqual(dst_indices.tolist(), [1])
self.assertEqual(kwargs["layer_id"], 1)
self.assertEqual(kwargs["page_size"], 1)
def test_mla_page_first_direct_per_layer_load_uses_tai_direct_pf_lf(self):
calls = []
def fake_direct(src_ptrs, dst_ptrs, src_indices, dst_indices, **kwargs):
calls.append(
(src_ptrs, dst_ptrs, src_indices.clone(), dst_indices.clone(), kwargs)
)
host_pool = MLATokenToKVPoolHost.__new__(MLATokenToKVPoolHost)
host_pool.layout = "page_first_direct"
host_pool.page_size = 4
host_pool.kv_buffer = torch.empty((8, 3, 4, 1, 16), dtype=torch.uint8)
device_pool = type("DevicePool", (), {})()
device_pool.kv_buffer = torch.empty((3, 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_direct_pf_lf",
return_value=fake_direct,
):
host_pool.load_to_device_per_layer(
device_pool,
host_indices,
device_indices,
layer_id=2,
io_backend="direct",
)
self.assertEqual(len(calls), 1)
src_ptrs, dst_ptrs, src_indices, dst_indices, kwargs = calls[0]
self.assertEqual(len(src_ptrs), 1)
self.assertEqual(src_ptrs[0].data_ptr(), host_pool.kv_buffer.data_ptr())
self.assertEqual(len(dst_ptrs), 1)
self.assertEqual(dst_ptrs[0].data_ptr(), device_pool.kv_buffer[2].data_ptr())
self.assertEqual(src_indices.tolist(), [4, 5, 6, 7])
self.assertEqual(dst_indices.tolist(), [12, 13, 14, 15])
self.assertEqual(kwargs["layer_id"], 2)
self.assertEqual(kwargs["page_size"], 4)
def test_mha_page_first_direct_per_layer_load_uses_tai_direct_pf_lf(self):
calls = []
def fake_direct(src_ptrs, dst_ptrs, src_indices, dst_indices, **kwargs):
calls.append(
(src_ptrs, dst_ptrs, src_indices.clone(), dst_indices.clone(), kwargs)
)
host_pool = MHATokenToKVPoolHost.__new__(MHATokenToKVPoolHost)
host_pool.layout = "page_first_direct"
host_pool.page_size = 4
host_pool.kv_buffer = torch.empty((2, 8, 3, 4, 2, 8), dtype=torch.uint8)
device_pool = type("DevicePool", (), {})()
device_pool.k_buffer = torch.empty((3, 32, 2, 8), dtype=torch.uint8)
device_pool.v_buffer = torch.empty((3, 32, 2, 8), 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_direct_pf_lf",
return_value=fake_direct,
):
host_pool.load_to_device_per_layer(
device_pool,
host_indices,
device_indices,
layer_id=2,
io_backend="direct",
)
self.assertEqual(len(calls), 1)
src_ptrs, dst_ptrs, src_indices, dst_indices, kwargs = calls[0]
self.assertEqual(len(src_ptrs), 2)
self.assertEqual(src_ptrs[0].data_ptr(), host_pool.k_buffer.data_ptr())
self.assertEqual(src_ptrs[1].data_ptr(), host_pool.v_buffer.data_ptr())
self.assertEqual(len(dst_ptrs), 2)
self.assertEqual(dst_ptrs[0].data_ptr(), device_pool.k_buffer[2].data_ptr())
self.assertEqual(dst_ptrs[1].data_ptr(), device_pool.v_buffer[2].data_ptr())
self.assertEqual(src_indices.tolist(), [4, 5, 6, 7])
self.assertEqual(dst_indices.tolist(), [12, 13, 14, 15])
self.assertEqual(kwargs["layer_id"], 2)
self.assertEqual(kwargs["page_size"], 4)
def test_nsa_indexer_page_first_direct_per_layer_load_uses_tai_direct_pf_lf(self):
calls = []
def fake_direct(src_ptrs, dst_ptrs, src_indices, dst_indices, **kwargs):
calls.append(
(src_ptrs, dst_ptrs, src_indices.clone(), dst_indices.clone(), kwargs)
)
host_pool = NSATokenToKVPoolHost.__new__(NSATokenToKVPoolHost)
host_pool.layout = "page_first_direct"
host_pool.page_size = 4
host_pool.index_k_with_scale_buffer = torch.empty(
(8, 3, 1, 32), dtype=torch.uint8
)
device_pool = type("DevicePool", (), {})()
device_pool.index_k_with_scale_buffer = torch.empty(
(3, 8, 32), 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_direct_pf_lf",
return_value=fake_direct,
):
host_pool._load_indexer_to_device_per_layer(
device_pool,
host_indices,
device_indices,
layer_id=1,
io_backend="direct",
)
self.assertEqual(len(calls), 1)
src_ptrs, dst_ptrs, src_indices, dst_indices, kwargs = calls[0]
self.assertEqual(len(src_ptrs), 1)
self.assertEqual(
src_ptrs[0].data_ptr(),
host_pool.index_k_with_scale_buffer.data_ptr(),
)
self.assertEqual(len(dst_ptrs), 1)
self.assertEqual(
dst_ptrs[0].data_ptr(),
device_pool.index_k_with_scale_buffer[1].data_ptr(),
)
self.assertEqual(src_indices.tolist(), [1])
self.assertEqual(dst_indices.tolist(), [3])
self.assertEqual(kwargs["layer_id"], 1)
self.assertEqual(kwargs["page_size"], 1)
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_accepts_valid_tail_and_pads_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)
result = controller.write(logical_locs, node_id=21)
self.assertEqual(result.metadata.logical_len, 3)
self.assertEqual(result.metadata.valid_len, 3)
self.assertEqual(result.metadata.padded_len, 4)
self.assertEqual(result.metadata.owned_positions.tolist(), [0, 1, 2, 3])
self.assertEqual(result.metadata.host_indices.tolist(), [100, 101, 102, 103])
self.assertEqual(host_pool.alloc_calls, [4])
self.assertEqual(host_pool.layer_backups[0][1].tolist(), [4, 5, 6, 7])
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_valid_tail_with_draft_mirrors_target_padded_page(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.tensor([8, 9, 10], dtype=torch.int64)
result = controller.write(logical_locs, node_id=177)
self.assertEqual(result.metadata.logical_len, 3)
self.assertEqual(result.metadata.padded_len, 4)
self.assertEqual(result.metadata.owned_positions.tolist(), [0, 1, 2, 3])
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.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)
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_returns_valid_locs_while_transferring_padded_tail_page(self):
host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
allocator = FakeAllocator(alloc_result=torch.arange(64, 72, dtype=torch.int64))
controller = self.make_controller(host_pool, allocator=allocator, cp_rank=1)
node = TreeNode()
node.host_len = 6
node.cp_hicache = CpHiCacheNodeMetadata(
logical_len=6,
padded_len=8,
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], dtype=torch.int8),
page_size=4,
)
device_indices = controller.load_cp([node], node_id=112)
controller.start_loading()
self.assertEqual(device_indices.tolist(), list(range(64, 70)))
self.assertEqual(allocator.owner_alloc_calls, [[3, 0]])
self.assertEqual(host_pool.loads[0][0].tolist(), [100, 101, 102, 103])
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_load_valid_tail_with_draft_returns_valid_locs_only(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(8, 12, 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 = 3
node.cp_hicache = CpHiCacheNodeMetadata(
logical_len=3,
padded_len=4,
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([1], dtype=torch.int8),
page_size=4,
)
device_indices = controller.load_cp([node], node_id=178)
controller.start_loading()
self.assertEqual(device_indices.tolist(), [8, 9, 10])
self.assertEqual(allocator.owner_alloc_calls, [[1]])
self.assertEqual(host_pool.loads[0][1].tolist(), [4, 5, 6, 7])
self.assertEqual(draft_host_pool.loads[0][1].tolist(), [4, 5, 6, 7])
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, [178])
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