367dff06f3
CP HiCache now treats draft KV as a strict target-owned payload through pending write visibility, host eviction, and state-buffer registration. Host metadata created before async D2H ack is no longer request-visible, so match_prefix cannot select an in-flight host node. Draft host eviction now fails before target cleanup when draft metadata is missing, and prefill/decode share one helper for draft NSA state buffers so shared-KV mode cannot silently skip mismatched draft state. Constraint: CP shared KV + HiCache + EAGLE/MTP must not expose target-only host hits or skipped draft state as valid cache hits Rejected: Rely on event-loop ordering and lock_ref to hide in-flight writes | match_prefix does not consult lock_ref and can observe host_len/cp_hicache directly Rejected: Keep draft state mismatch as debug-only skip | it can poison speculative acceptance while looking like a successful cache hit Confidence: high Scope-risk: moderate Directive: Do not reintroduce silent draft/target fallback in CP shared-KV HiCache paths; malformed strong-sync metadata should fail fast Tested: python -m py_compile targeted modified files Tested: remote g0034 container pytest test/registered/unit/mem_cache/test_cp_hicache_metadata.py test/registered/unit/managers/test_hicache_controller_cp.py test/registered/unit/disaggregation/test_cp_shared_kv_transfer_mapping.py -q (91 passed) Not-tested: Full CP shared KV + HiCache + EAGLE/MTP ETE server run after this commit
767 lines
30 KiB
Python
767 lines
30 KiB
Python
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.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.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 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
|
|
|
|
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.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[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[0][1].tolist(), [4, 5, 6, 7])
|
|
self.assertEqual(draft_host_pool.backups[0][1].tolist(), [4, 5, 6, 7])
|
|
self.assertIs(draft_host_pool.backups[0][2], 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_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()
|