laoyao0822
315eaaff56
Cache-hit prefill can skip draft forward for the prefix while PD transfer still reads draft KV for that same prefix. CP HiCache therefore needs to persist draft/MTP KV alongside target KV instead of relying on whatever remains in the draft GPU pool. Constraint: CP HiCache is host-only here; storage backends remain unsupported for CP shared KV. Constraint: CP shared KV must keep owner-page semantics and avoid falling back to full KV on every rank. Rejected: Recompute cached-prefix draft KV during prefill | loses the HiCache benefit and reintroduces the large hidden/KV footprint. Rejected: Change PD transfer to skip draft prefix KV | decode still needs draft cache continuity for MTP acceptance. Confidence: medium Scope-risk: moderate Directive: Keep target and draft CP HiCache metadata/load/write/evict paths in lockstep; changing one without the other can silently reduce MTP accept length. Tested: Remote g0034 container /sgl-workspace/sglang-tai: python3 -m pytest -q test/registered/unit/managers/test_hicache_controller_cp.py test/registered/unit/mem_cache/test_cp_hicache_metadata.py => 58 passed, 3 warnings Not-tested: Full multi-node HiCache+MTP serving benchmark and accept-length recovery.
957 lines
36 KiB
Python
957 lines
36 KiB
Python
import sys
|
|
import types
|
|
import unittest
|
|
from unittest.mock import patch
|
|
|
|
import torch
|
|
|
|
# Prefer the real sgl_kernel package when the test image provides it so custom
|
|
# Torch operators are registered. Fall back to stubs on local CPU-only hosts.
|
|
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:
|
|
sys.modules["sgl_kernel.kvcacheio"] = types.ModuleType("sgl_kernel.kvcacheio")
|
|
|
|
from sglang.srt.mem_cache.base_prefix_cache import EvictParams, MatchPrefixParams
|
|
from sglang.srt.mem_cache.hiradix_cache import CpHiCacheNodeMetadata, HiRadixCache
|
|
from sglang.srt.mem_cache.radix_cache import RadixKey, 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 TestCpHiCacheImports(CustomTestCase):
|
|
def test_cp_hicache_public_imports_without_sgl_kernel(self):
|
|
import subprocess
|
|
|
|
subprocess.run(
|
|
[
|
|
sys.executable,
|
|
"-c",
|
|
"from sglang.srt.mem_cache.hiradix_cache import HiRadixCache, CpHiCacheNodeMetadata; "
|
|
"from sglang.srt.managers.cache_controller import HiCacheController; "
|
|
"print('OK')",
|
|
],
|
|
check=True,
|
|
capture_output=True,
|
|
text=True,
|
|
)
|
|
|
|
|
|
class TestCpHiCacheNodeMetadata(CustomTestCase):
|
|
def test_split_zero_len_moves_all_positions_to_child(self):
|
|
metadata = CpHiCacheNodeMetadata(
|
|
logical_len=8,
|
|
owned_positions=torch.tensor([1, 3, 7], dtype=torch.int64),
|
|
host_indices=torch.tensor([10, 11, 12], dtype=torch.int64),
|
|
)
|
|
|
|
parent, child = metadata.split(0)
|
|
|
|
self.assertEqual(parent.logical_len, 0)
|
|
self.assertEqual(parent.owned_positions.tolist(), [])
|
|
self.assertEqual(parent.host_indices.tolist(), [])
|
|
self.assertEqual(child.logical_len, 8)
|
|
self.assertEqual(child.owned_positions.tolist(), [1, 3, 7])
|
|
self.assertEqual(child.host_indices.tolist(), [10, 11, 12])
|
|
|
|
def test_split_rebases_child_positions(self):
|
|
metadata = CpHiCacheNodeMetadata(
|
|
logical_len=10,
|
|
owned_positions=torch.tensor([0, 2, 5, 9], dtype=torch.int64),
|
|
host_indices=torch.tensor([20, 21, 22, 23], dtype=torch.int64),
|
|
)
|
|
|
|
parent, child = metadata.split(5)
|
|
|
|
self.assertEqual(parent.logical_len, 5)
|
|
self.assertEqual(parent.owned_positions.tolist(), [0, 2])
|
|
self.assertEqual(parent.host_indices.tolist(), [20, 21])
|
|
self.assertEqual(child.logical_len, 5)
|
|
self.assertEqual(child.owned_positions.tolist(), [0, 4])
|
|
self.assertEqual(child.host_indices.tolist(), [22, 23])
|
|
|
|
def test_split_keeps_draft_host_indices_aligned_with_owned_positions(self):
|
|
metadata = CpHiCacheNodeMetadata(
|
|
logical_len=10,
|
|
owned_positions=torch.tensor([0, 2, 5, 9], dtype=torch.int64),
|
|
host_indices=torch.tensor([20, 21, 22, 23], dtype=torch.int64),
|
|
draft_host_indices=torch.tensor([120, 121, 122, 123], dtype=torch.int64),
|
|
)
|
|
|
|
parent, child = metadata.split(5)
|
|
|
|
self.assertEqual(parent.host_indices.tolist(), [20, 21])
|
|
self.assertEqual(parent.draft_host_indices.tolist(), [120, 121])
|
|
self.assertEqual(child.owned_positions.tolist(), [0, 4])
|
|
self.assertEqual(child.host_indices.tolist(), [22, 23])
|
|
self.assertEqual(child.draft_host_indices.tolist(), [122, 123])
|
|
|
|
def test_zero_owned_metadata_is_valid(self):
|
|
metadata = CpHiCacheNodeMetadata(
|
|
logical_len=64,
|
|
owned_positions=torch.empty((0,), dtype=torch.int32),
|
|
host_indices=torch.empty((0,), dtype=torch.int32),
|
|
)
|
|
|
|
self.assertEqual(metadata.logical_len, 64)
|
|
self.assertEqual(metadata.owned_positions.device.type, "cpu")
|
|
self.assertEqual(metadata.host_indices.device.type, "cpu")
|
|
self.assertEqual(metadata.owned_positions.dtype, torch.int64)
|
|
self.assertEqual(metadata.host_indices.dtype, torch.int64)
|
|
|
|
def test_non_int64_inputs_are_converted(self):
|
|
metadata = CpHiCacheNodeMetadata(
|
|
logical_len=4,
|
|
owned_positions=torch.tensor([1, 3], dtype=torch.int32),
|
|
host_indices=torch.tensor([10, 11], dtype=torch.int32),
|
|
)
|
|
|
|
self.assertEqual(metadata.owned_positions.dtype, torch.int64)
|
|
self.assertEqual(metadata.host_indices.dtype, torch.int64)
|
|
|
|
def test_negative_logical_len_raises(self):
|
|
with self.assertRaisesRegex(ValueError, "logical_len"):
|
|
CpHiCacheNodeMetadata(
|
|
logical_len=-1,
|
|
owned_positions=torch.empty((0,), dtype=torch.int64),
|
|
host_indices=torch.empty((0,), dtype=torch.int64),
|
|
)
|
|
|
|
def test_metadata_does_not_alias_input_tensors(self):
|
|
owned_positions = torch.tensor([1, 3], dtype=torch.int64)
|
|
host_indices = torch.tensor([10, 11], dtype=torch.int64)
|
|
|
|
metadata = CpHiCacheNodeMetadata(
|
|
logical_len=4,
|
|
owned_positions=owned_positions,
|
|
host_indices=host_indices,
|
|
)
|
|
owned_positions[0] = 2
|
|
host_indices[0] = 12
|
|
|
|
self.assertEqual(metadata.owned_positions.tolist(), [1, 3])
|
|
self.assertEqual(metadata.host_indices.tolist(), [10, 11])
|
|
|
|
def test_invalid_split_raises(self):
|
|
metadata = CpHiCacheNodeMetadata(
|
|
logical_len=4,
|
|
owned_positions=torch.tensor([1], dtype=torch.int64),
|
|
host_indices=torch.tensor([9], dtype=torch.int64),
|
|
)
|
|
|
|
with self.assertRaisesRegex(ValueError, "split_len"):
|
|
metadata.split(5)
|
|
|
|
def test_unsorted_positions_raise(self):
|
|
with self.assertRaisesRegex(ValueError, "sorted"):
|
|
CpHiCacheNodeMetadata(
|
|
logical_len=4,
|
|
owned_positions=torch.tensor([2, 1], dtype=torch.int64),
|
|
host_indices=torch.tensor([9, 10], dtype=torch.int64),
|
|
)
|
|
|
|
def test_duplicate_positions_raise(self):
|
|
with self.assertRaisesRegex(ValueError, "strictly increasing"):
|
|
CpHiCacheNodeMetadata(
|
|
logical_len=4,
|
|
owned_positions=torch.tensor([1, 1], dtype=torch.int64),
|
|
host_indices=torch.tensor([9, 10], dtype=torch.int64),
|
|
)
|
|
|
|
def test_length_mismatch_raises(self):
|
|
with self.assertRaisesRegex(ValueError, "same length"):
|
|
CpHiCacheNodeMetadata(
|
|
logical_len=4,
|
|
owned_positions=torch.tensor([1, 2], dtype=torch.int64),
|
|
host_indices=torch.tensor([9], dtype=torch.int64),
|
|
)
|
|
|
|
def test_draft_host_length_mismatch_raises(self):
|
|
with self.assertRaisesRegex(ValueError, "draft_host_indices.*same length"):
|
|
CpHiCacheNodeMetadata(
|
|
logical_len=4,
|
|
owned_positions=torch.tensor([1, 2], dtype=torch.int64),
|
|
host_indices=torch.tensor([9, 10], dtype=torch.int64),
|
|
draft_host_indices=torch.tensor([109], dtype=torch.int64),
|
|
)
|
|
|
|
def test_out_of_range_positions_raise(self):
|
|
with self.assertRaisesRegex(ValueError, r"\[0, logical_len\)"):
|
|
CpHiCacheNodeMetadata(
|
|
logical_len=4,
|
|
owned_positions=torch.tensor([4], dtype=torch.int64),
|
|
host_indices=torch.tensor([9], dtype=torch.int64),
|
|
)
|
|
|
|
|
|
class FakeWriteFailure:
|
|
metadata = None
|
|
|
|
def __init__(self, required_host_slots):
|
|
self.required_host_slots = required_host_slots
|
|
|
|
|
|
class FakeWriteSuccess:
|
|
required_host_slots = 0
|
|
|
|
def __init__(self, metadata):
|
|
self.metadata = metadata
|
|
|
|
|
|
class FakeWriteController:
|
|
def __init__(self, required_host_slots):
|
|
self.required_host_slots = required_host_slots
|
|
self.calls = 0
|
|
self.write_policy = "write_through"
|
|
self.evicted_host_indices = []
|
|
|
|
def write(self, device_indices, node_id=-1, priority=None):
|
|
self.calls += 1
|
|
if self.calls == 1:
|
|
return FakeWriteFailure(self.required_host_slots)
|
|
return FakeWriteSuccess(
|
|
CpHiCacheNodeMetadata(
|
|
logical_len=len(device_indices),
|
|
owned_positions=torch.tensor([0], dtype=torch.int64),
|
|
host_indices=torch.tensor([99], dtype=torch.int64),
|
|
)
|
|
)
|
|
|
|
def evict_host(self, host_indices):
|
|
self.evicted_host_indices.append(host_indices.clone())
|
|
return len(host_indices)
|
|
|
|
|
|
class FakeZeroOwnedWriteController:
|
|
write_policy = "write_through"
|
|
|
|
def write(self, device_indices, node_id=-1, priority=None):
|
|
return FakeWriteSuccess(
|
|
CpHiCacheNodeMetadata(
|
|
logical_len=len(device_indices),
|
|
owned_positions=torch.empty((0,), dtype=torch.int64),
|
|
host_indices=torch.empty((0,), dtype=torch.int64),
|
|
)
|
|
)
|
|
|
|
|
|
class FakeEvictionStrategy:
|
|
def get_priority(self, node):
|
|
return 0
|
|
|
|
|
|
class FakeEvictDeviceController:
|
|
write_policy = "write_through"
|
|
|
|
def evict_device(self, device_indices):
|
|
return len(device_indices)
|
|
|
|
|
|
class FakeTokenAllocator:
|
|
def available_size(self):
|
|
return 0
|
|
|
|
|
|
class TestHiRadixCacheCPBackup(CustomTestCase):
|
|
def test_node_backuped_uses_cp_metadata(self):
|
|
cache = HiRadixCache.__new__(HiRadixCache)
|
|
cache._uses_cp_hicache = True
|
|
node = TreeNode()
|
|
node.host_len = 8
|
|
node.cp_hicache = CpHiCacheNodeMetadata(
|
|
logical_len=8,
|
|
owned_positions=torch.tensor([1, 2], dtype=torch.int64),
|
|
host_indices=torch.tensor([10, 11], dtype=torch.int64),
|
|
)
|
|
|
|
self.assertTrue(cache._node_backuped(node))
|
|
|
|
def test_inc_hit_count_does_not_rewrite_cp_backed_node(self):
|
|
cache = HiRadixCache.__new__(HiRadixCache)
|
|
cache._uses_cp_hicache = True
|
|
cache.write_through_threshold = 1
|
|
cache.cache_controller = type(
|
|
"Controller", (), {"write_policy": "write_through"}
|
|
)()
|
|
cache.write_backup = lambda node: (_ for _ in ()).throw(
|
|
AssertionError("must not rewrite")
|
|
)
|
|
node = TreeNode()
|
|
node.host_len = 4
|
|
node.cp_hicache = CpHiCacheNodeMetadata(
|
|
logical_len=4,
|
|
owned_positions=torch.tensor([], dtype=torch.int64),
|
|
host_indices=torch.tensor([], dtype=torch.int64),
|
|
)
|
|
|
|
cache._inc_hit_count(node)
|
|
|
|
self.assertEqual(node.hit_count, 1)
|
|
|
|
def test_write_backup_retries_by_required_physical_slots(self):
|
|
cache = HiRadixCache.__new__(HiRadixCache)
|
|
cache._uses_cp_hicache = True
|
|
cache.cache_controller = FakeWriteController(required_host_slots=1)
|
|
cache.evictable_host_leaves = set()
|
|
cache.eviction_strategy = FakeEvictionStrategy()
|
|
cache.get_child_key_fn = lambda key: key.token_ids[0]
|
|
cache._record_remove_event = lambda node: None
|
|
cache.ongoing_write_through = {}
|
|
cache.inc_node_lock_ref = lambda node: None
|
|
|
|
root = TreeNode()
|
|
root.key = RadixKey(token_ids=[], extra_key=None)
|
|
root.value = []
|
|
cache.root_node = root
|
|
evictable_node = TreeNode()
|
|
evictable_node.parent = root
|
|
evictable_node.key = RadixKey(token_ids=[1], extra_key=None)
|
|
evictable_node.value = None
|
|
evictable_node.host_len = 4
|
|
evictable_node.cp_hicache = CpHiCacheNodeMetadata(
|
|
logical_len=4,
|
|
owned_positions=torch.tensor([0], dtype=torch.int64),
|
|
host_indices=torch.tensor([55], dtype=torch.int64),
|
|
)
|
|
root.children[1] = evictable_node
|
|
cache.evictable_host_leaves.add(evictable_node)
|
|
|
|
node = TreeNode()
|
|
node.value = torch.arange(16, dtype=torch.int64)
|
|
|
|
cache.write_backup(node)
|
|
|
|
self.assertEqual(
|
|
cache.cache_controller.evicted_host_indices[0].tolist(), [55]
|
|
)
|
|
self.assertEqual(evictable_node.host_len, 0)
|
|
self.assertIsNone(evictable_node.cp_hicache)
|
|
self.assertNotIn(1, root.children)
|
|
self.assertEqual(node.host_len, 16)
|
|
|
|
def test_write_backup_cp_success_returns_logical_length_for_zero_owned_rank(self):
|
|
cache = HiRadixCache.__new__(HiRadixCache)
|
|
cache._uses_cp_hicache = True
|
|
cache.cache_controller = FakeZeroOwnedWriteController()
|
|
cache.ongoing_write_through = {}
|
|
cache.inc_node_lock_ref = lambda node: None
|
|
node = TreeNode()
|
|
node.id = 123
|
|
node.value = torch.arange(16, dtype=torch.int64)
|
|
|
|
backed_len = cache.write_backup(node)
|
|
|
|
self.assertEqual(backed_len, 16)
|
|
self.assertEqual(node.host_len, 16)
|
|
self.assertEqual(node.cp_hicache.host_indices.tolist(), [])
|
|
|
|
def test_attach_storage_backend_rejects_cp_hicache_without_controller_call(self):
|
|
cache = HiRadixCache.__new__(HiRadixCache)
|
|
cache._uses_cp_hicache = True
|
|
cache.cache_controller = type(
|
|
"Controller",
|
|
(),
|
|
{
|
|
"attach_storage_backend": lambda *args, **kwargs: (_ for _ in ()).throw(
|
|
AssertionError("controller attach must not be called")
|
|
)
|
|
},
|
|
)()
|
|
|
|
ok, message = cache.attach_storage_backend("mooncake")
|
|
|
|
self.assertFalse(ok)
|
|
self.assertIn("CP shared KV", message)
|
|
self.assertIn("storage backend", message)
|
|
|
|
def test_evict_demotes_cp_backed_node_without_deleting_radix_child(self):
|
|
cache = HiRadixCache.__new__(HiRadixCache)
|
|
cache._uses_cp_hicache = True
|
|
cache.cache_controller = FakeEvictDeviceController()
|
|
cache.evictable_leaves = set()
|
|
cache.evictable_host_leaves = set()
|
|
cache.eviction_strategy = FakeEvictionStrategy()
|
|
cache.token_to_kv_pool_allocator = FakeTokenAllocator()
|
|
cache.evictable_size_ = 4
|
|
cache.protected_size_ = 0
|
|
cache._record_remove_event = lambda node: (_ for _ in ()).throw(
|
|
AssertionError("must not delete backed radix child")
|
|
)
|
|
|
|
root = TreeNode()
|
|
root.key = RadixKey(token_ids=[], extra_key=None)
|
|
root.value = []
|
|
cache.root_node = root
|
|
node = TreeNode()
|
|
node.parent = root
|
|
node.key = RadixKey(token_ids=[1, 2, 3, 4], extra_key=None)
|
|
node.value = torch.arange(4, dtype=torch.int64)
|
|
node.host_len = 4
|
|
node.cp_hicache = CpHiCacheNodeMetadata(
|
|
logical_len=4,
|
|
owned_positions=torch.tensor([0], dtype=torch.int64),
|
|
host_indices=torch.tensor([55], dtype=torch.int64),
|
|
)
|
|
root.children[1] = node
|
|
cache.evictable_leaves.add(node)
|
|
|
|
cache.evict(EvictParams(num_tokens=4))
|
|
|
|
self.assertIn(1, root.children)
|
|
self.assertIsNone(node.value)
|
|
self.assertIsNotNone(node.cp_hicache)
|
|
self.assertEqual(node.cp_hicache.host_indices.tolist(), [55])
|
|
|
|
|
|
class TestHiRadixCacheCPSplitEvict(CustomTestCase):
|
|
def test_split_node_splits_cp_metadata_by_owned_positions(self):
|
|
cache = HiRadixCache.__new__(HiRadixCache)
|
|
cache._uses_cp_hicache = True
|
|
cache.get_child_key_fn = lambda key: key.token_ids[0]
|
|
cache.page_size = 1
|
|
|
|
root = TreeNode()
|
|
root.key = RadixKey([])
|
|
child = TreeNode()
|
|
child.parent = root
|
|
child.key = RadixKey(list(range(10)))
|
|
child.value = None
|
|
child.hash_value = None
|
|
child.host_len = 10
|
|
child.cp_hicache = CpHiCacheNodeMetadata(
|
|
logical_len=10,
|
|
owned_positions=torch.tensor([0, 2, 5, 9], dtype=torch.int64),
|
|
host_indices=torch.tensor([20, 21, 22, 23], dtype=torch.int64),
|
|
)
|
|
root.children[0] = child
|
|
|
|
new_node = cache._split_node(child.key, child, 5)
|
|
|
|
self.assertEqual(new_node.host_len, 5)
|
|
self.assertEqual(child.host_len, 5)
|
|
self.assertEqual(new_node.cp_hicache.owned_positions.tolist(), [0, 2])
|
|
self.assertEqual(new_node.cp_hicache.host_indices.tolist(), [20, 21])
|
|
self.assertEqual(child.cp_hicache.owned_positions.tolist(), [0, 4])
|
|
self.assertEqual(child.cp_hicache.host_indices.tolist(), [22, 23])
|
|
|
|
def test_cp_host_eviction_uses_physical_freed_slots_for_progress(self):
|
|
cache = HiRadixCache.__new__(HiRadixCache)
|
|
cache._uses_cp_hicache = True
|
|
cache.root_node = TreeNode()
|
|
cache.root_node.key = RadixKey([])
|
|
cache.evictable_host_leaves = set()
|
|
cache.get_child_key_fn = lambda key: key.token_ids[0]
|
|
cache.eviction_strategy = type(
|
|
"Strategy", (), {"get_priority": lambda self, node: 0}
|
|
)()
|
|
cache._clear_pin = lambda node: None
|
|
cache._record_remove_event = lambda node: None
|
|
cache._update_host_leaf_status = lambda node: None
|
|
cache._node_host_evict_indices = lambda node: torch.tensor(
|
|
[99], dtype=torch.int64
|
|
)
|
|
freed = []
|
|
cache.cache_controller = type(
|
|
"Controller",
|
|
(),
|
|
{
|
|
"evict_host": lambda self, indices: freed.append(indices.clone())
|
|
or len(indices)
|
|
},
|
|
)()
|
|
node = TreeNode()
|
|
node.parent = cache.root_node
|
|
node.key = RadixKey([1, 2, 3, 4])
|
|
node.value = None
|
|
node.host_len = 4
|
|
node.cp_hicache = CpHiCacheNodeMetadata(
|
|
logical_len=4,
|
|
owned_positions=torch.tensor([1], dtype=torch.int64),
|
|
host_indices=torch.tensor([70], dtype=torch.int64),
|
|
)
|
|
cache.root_node.children[1] = node
|
|
cache.evictable_host_leaves.add(node)
|
|
|
|
physical_freed = cache._evict_host_for_physical_slots(1)
|
|
|
|
self.assertEqual(physical_freed, 1)
|
|
self.assertEqual(freed[0].tolist(), [99])
|
|
self.assertEqual(node.host_len, 0)
|
|
self.assertIsNone(node.cp_hicache)
|
|
|
|
def test_cp_host_eviction_unlinks_stale_leaf_to_free_parent_slots(self):
|
|
cache = HiRadixCache.__new__(HiRadixCache)
|
|
cache._uses_cp_hicache = True
|
|
cache.root_node = TreeNode()
|
|
cache.root_node.key = RadixKey([])
|
|
cache.evictable_host_leaves = set()
|
|
cache.get_child_key_fn = lambda key: key.token_ids[0]
|
|
cache.eviction_strategy = type(
|
|
"Strategy", (), {"get_priority": lambda self, node: 0}
|
|
)()
|
|
cache._clear_pin = lambda node: None
|
|
cache._record_remove_event = lambda node: None
|
|
freed = []
|
|
cache.cache_controller = type(
|
|
"Controller",
|
|
(),
|
|
{
|
|
"evict_host": lambda self, indices: freed.append(indices.clone())
|
|
or len(indices)
|
|
},
|
|
)()
|
|
|
|
parent = TreeNode()
|
|
parent.parent = cache.root_node
|
|
parent.key = RadixKey([1])
|
|
parent.value = None
|
|
parent.host_len = 4
|
|
parent.cp_hicache = CpHiCacheNodeMetadata(
|
|
logical_len=4,
|
|
owned_positions=torch.tensor([2], dtype=torch.int64),
|
|
host_indices=torch.tensor([80], dtype=torch.int64),
|
|
)
|
|
cache.root_node.children[1] = parent
|
|
|
|
stale_child = TreeNode()
|
|
stale_child.parent = parent
|
|
stale_child.key = RadixKey([2])
|
|
stale_child.value = None
|
|
stale_child.host_len = 0
|
|
stale_child.cp_hicache = None
|
|
parent.children[2] = stale_child
|
|
cache.evictable_host_leaves.add(stale_child)
|
|
|
|
physical_freed = cache._evict_host_for_physical_slots(1)
|
|
|
|
self.assertEqual(physical_freed, 1)
|
|
self.assertEqual(freed[0].tolist(), [80])
|
|
self.assertNotIn(2, parent.children)
|
|
self.assertEqual(parent.host_len, 0)
|
|
self.assertIsNone(parent.cp_hicache)
|
|
self.assertNotIn(1, cache.root_node.children)
|
|
|
|
def test_cp_host_eviction_preserves_parent_with_sibling_after_stale_cleanup(
|
|
self,
|
|
):
|
|
cache = HiRadixCache.__new__(HiRadixCache)
|
|
cache._uses_cp_hicache = True
|
|
cache.root_node = TreeNode()
|
|
cache.root_node.key = RadixKey([])
|
|
cache.evictable_host_leaves = set()
|
|
cache.get_child_key_fn = lambda key: key.token_ids[0]
|
|
cache.eviction_strategy = type(
|
|
"Strategy", (), {"get_priority": lambda self, node: 0}
|
|
)()
|
|
cache._clear_pin = lambda node: None
|
|
cache._record_remove_event = lambda node: None
|
|
freed = []
|
|
cache.cache_controller = type(
|
|
"Controller",
|
|
(),
|
|
{
|
|
"evict_host": lambda self, indices: freed.append(indices.clone())
|
|
or len(indices)
|
|
},
|
|
)()
|
|
|
|
parent = TreeNode()
|
|
parent.parent = cache.root_node
|
|
parent.key = RadixKey([1])
|
|
parent.value = None
|
|
parent.host_len = 4
|
|
parent.cp_hicache = CpHiCacheNodeMetadata(
|
|
logical_len=4,
|
|
owned_positions=torch.tensor([2], dtype=torch.int64),
|
|
host_indices=torch.tensor([80], dtype=torch.int64),
|
|
)
|
|
cache.root_node.children[1] = parent
|
|
|
|
stale_child = TreeNode()
|
|
stale_child.parent = parent
|
|
stale_child.key = RadixKey([2])
|
|
stale_child.value = None
|
|
stale_child.host_len = 0
|
|
stale_child.cp_hicache = None
|
|
parent.children[2] = stale_child
|
|
|
|
sibling = TreeNode()
|
|
sibling.parent = parent
|
|
sibling.key = RadixKey([3])
|
|
sibling.value = None
|
|
sibling.host_len = 0
|
|
sibling.cp_hicache = None
|
|
parent.children[3] = sibling
|
|
cache.evictable_host_leaves.add(stale_child)
|
|
|
|
physical_freed = cache._evict_host_for_physical_slots(1)
|
|
|
|
self.assertEqual(physical_freed, 0)
|
|
self.assertEqual(freed, [])
|
|
self.assertNotIn(2, parent.children)
|
|
self.assertIn(1, cache.root_node.children)
|
|
self.assertIs(cache.root_node.children[1], parent)
|
|
self.assertIn(3, parent.children)
|
|
self.assertIs(parent.children[3], sibling)
|
|
self.assertEqual(parent.host_len, 4)
|
|
self.assertIsNotNone(parent.cp_hicache)
|
|
self.assertEqual(parent.cp_hicache.host_indices.tolist(), [80])
|
|
|
|
def test_synchronized_cp_host_eviction_removes_zero_owned_logical_leaf(self):
|
|
cache = HiRadixCache.__new__(HiRadixCache)
|
|
cache._uses_cp_hicache = True
|
|
cache.tp_world_size = 2
|
|
cache.tp_group = object()
|
|
cache.root_node = TreeNode()
|
|
cache.root_node.key = RadixKey([])
|
|
cache.evictable_host_leaves = set()
|
|
cache.get_child_key_fn = lambda key: key.token_ids[0]
|
|
cache.eviction_strategy = type(
|
|
"Strategy", (), {"get_priority": lambda self, node: 0}
|
|
)()
|
|
cache._clear_pin = lambda node: None
|
|
cache._record_remove_event = lambda node: None
|
|
freed = []
|
|
cache.cache_controller = type(
|
|
"Controller",
|
|
(),
|
|
{
|
|
"evict_host": lambda self, indices: freed.append(indices.clone())
|
|
or len(indices)
|
|
},
|
|
)()
|
|
node = TreeNode()
|
|
node.parent = cache.root_node
|
|
node.key = RadixKey([1, 2, 3, 4])
|
|
node.value = None
|
|
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),
|
|
)
|
|
cache.root_node.children[1] = node
|
|
cache.evictable_host_leaves.add(node)
|
|
|
|
all_done_states = iter([False, True])
|
|
cache._cp_all_ranks_true = lambda done: next(all_done_states)
|
|
cache._cp_broadcast_node_ids = lambda node_ids, max_ids: node_ids[:max_ids]
|
|
cache._cp_filter_all_ranks_safe_node_ids = (
|
|
lambda node_ids, is_safe, **_kwargs: [
|
|
node_id
|
|
for node_id in node_ids
|
|
if is_safe(cache._cp_node_by_id(node_id))
|
|
]
|
|
)
|
|
|
|
physical_freed = cache._cp_evict_host_for_physical_slots(0)
|
|
|
|
self.assertEqual(physical_freed, 0)
|
|
self.assertEqual(freed, [])
|
|
self.assertNotIn(1, cache.root_node.children)
|
|
self.assertEqual(node.host_len, 0)
|
|
self.assertIsNone(node.cp_hicache)
|
|
|
|
def test_cp_host_eviction_skips_all_reduce_without_tp_group(self):
|
|
cache = HiRadixCache.__new__(HiRadixCache)
|
|
cache._uses_cp_hicache = True
|
|
cache.tp_world_size = 2
|
|
cache.tp_group = None
|
|
cache.root_node = TreeNode()
|
|
cache.root_node.key = RadixKey([])
|
|
cache.evictable_host_leaves = set()
|
|
cache.get_child_key_fn = lambda key: key.token_ids[0]
|
|
cache.eviction_strategy = type(
|
|
"Strategy", (), {"get_priority": lambda self, node: 0}
|
|
)()
|
|
cache._clear_pin = lambda node: None
|
|
cache._record_remove_event = lambda node: None
|
|
cache.cache_controller = type(
|
|
"Controller",
|
|
(),
|
|
{
|
|
"evict_host": lambda self, indices: (_ for _ in ()).throw(
|
|
AssertionError("must not evict host slots")
|
|
)
|
|
},
|
|
)()
|
|
node = TreeNode()
|
|
node.parent = cache.root_node
|
|
node.key = RadixKey([1, 2, 3, 4])
|
|
node.value = None
|
|
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),
|
|
)
|
|
cache.root_node.children[1] = node
|
|
cache.evictable_host_leaves.add(node)
|
|
|
|
with patch(
|
|
"torch.distributed.all_reduce",
|
|
side_effect=AssertionError("must not all_reduce without tp_group"),
|
|
):
|
|
physical_freed = cache._evict_host_for_physical_slots(
|
|
0, synchronize_across_ranks=True
|
|
)
|
|
|
|
self.assertEqual(physical_freed, 0)
|
|
self.assertIn(1, cache.root_node.children)
|
|
self.assertEqual(node.host_len, 4)
|
|
self.assertIsNotNone(node.cp_hicache)
|
|
|
|
|
|
class TestHiRadixCacheCPLoadBack(CustomTestCase):
|
|
def test_cp_load_back_uses_host_len_not_host_value(self):
|
|
cache = HiRadixCache.__new__(HiRadixCache)
|
|
cache._uses_cp_hicache = True
|
|
cache.root_node = TreeNode()
|
|
cache.device = "cpu"
|
|
cache.load_back_threshold = 1
|
|
cache.evictable_size_ = 0
|
|
cache.metrics_collector = None
|
|
cache.ongoing_load_back = {}
|
|
cache.cache_controller = type(
|
|
"Controller",
|
|
(),
|
|
{
|
|
"load_cp": lambda self, nodes, node_id=-1: torch.arange(
|
|
32, 40, dtype=torch.int64
|
|
)
|
|
},
|
|
)()
|
|
cache.inc_lock_ref = lambda node: type("Result", (), {"delta": 0})()
|
|
cache.dec_lock_ref = lambda node: None
|
|
cache.evict = lambda params: None
|
|
parent = cache.root_node
|
|
parent.key = RadixKey([])
|
|
parent.value = torch.empty((0,), dtype=torch.int64)
|
|
node = TreeNode()
|
|
node.parent = parent
|
|
node.key = RadixKey(list(range(8)))
|
|
node.value = None
|
|
node.host_value = None
|
|
node.host_len = 8
|
|
node.cp_hicache = CpHiCacheNodeMetadata(
|
|
logical_len=8,
|
|
owned_positions=torch.tensor([0, 1], dtype=torch.int64),
|
|
host_indices=torch.tensor([50, 51], dtype=torch.int64),
|
|
)
|
|
|
|
loaded = cache.load_back(node)
|
|
|
|
self.assertEqual(loaded.tolist(), list(range(32, 40)))
|
|
self.assertEqual(node.value.tolist(), list(range(32, 40)))
|
|
|
|
def test_cp_load_back_threshold_uses_logical_length(self):
|
|
cache = HiRadixCache.__new__(HiRadixCache)
|
|
cache._uses_cp_hicache = True
|
|
cache.root_node = TreeNode()
|
|
cache.root_node.value = torch.empty((0,), dtype=torch.int64)
|
|
cache.load_back_threshold = 5
|
|
cache.evictable_size_ = 0
|
|
cache.metrics_collector = None
|
|
cache.ongoing_load_back = {}
|
|
cache.inc_lock_ref = lambda node: type("Result", (), {"delta": 0})()
|
|
cache.dec_lock_ref = lambda node: None
|
|
cache.cache_controller = type(
|
|
"Controller",
|
|
(),
|
|
{
|
|
"load_cp": lambda self, nodes, node_id=-1: torch.arange(
|
|
10, 16, dtype=torch.int64
|
|
)
|
|
},
|
|
)()
|
|
node = TreeNode()
|
|
node.parent = cache.root_node
|
|
node.value = None
|
|
node.host_len = 6
|
|
node.cp_hicache = CpHiCacheNodeMetadata(
|
|
logical_len=6,
|
|
owned_positions=torch.empty((0,), dtype=torch.int64),
|
|
host_indices=torch.empty((0,), dtype=torch.int64),
|
|
)
|
|
|
|
loaded = cache.load_back(node)
|
|
|
|
self.assertEqual(loaded.tolist(), [10, 11, 12, 13, 14, 15])
|
|
|
|
def test_cp_match_prefix_counts_logical_host_hit_without_host_value(self):
|
|
cache = HiRadixCache.__new__(HiRadixCache)
|
|
cache._uses_cp_hicache = True
|
|
cache.device = "cpu"
|
|
cache.disable = False
|
|
cache.page_size = 1
|
|
cache.get_child_key_fn = lambda key: key.token_ids[0]
|
|
cache.key_match_fn = lambda child_key, key: sum(
|
|
1 for lhs, rhs in zip(child_key.token_ids, key.token_ids) if lhs == rhs
|
|
)
|
|
cache.maybe_bigram_convert = lambda key: (key, None)
|
|
root = TreeNode()
|
|
root.key = RadixKey([])
|
|
root.value = torch.empty((0,), dtype=torch.int64)
|
|
root.host_len = 0
|
|
cache.root_node = root
|
|
node = TreeNode()
|
|
node.parent = root
|
|
node.key = RadixKey(list(range(8)))
|
|
node.value = None
|
|
node.host_value = None
|
|
node.host_len = 8
|
|
node.cp_hicache = CpHiCacheNodeMetadata(
|
|
logical_len=8,
|
|
owned_positions=torch.tensor([0, 1], dtype=torch.int64),
|
|
host_indices=torch.tensor([50, 51], dtype=torch.int64),
|
|
)
|
|
root.children[0] = node
|
|
|
|
result = cache.match_prefix(MatchPrefixParams(key=RadixKey(list(range(8)))))
|
|
|
|
self.assertEqual(result.host_hit_length, 8)
|
|
self.assertIs(result.last_host_node, node)
|
|
self.assertIs(result.last_device_node, root)
|
|
|
|
def test_cp_match_prefix_shorter_than_page_returns_empty_root_match(self):
|
|
cache = HiRadixCache.__new__(HiRadixCache)
|
|
cache._uses_cp_hicache = True
|
|
cache.device = "cpu"
|
|
cache.disable = False
|
|
cache.page_size = 4
|
|
cache.maybe_bigram_convert = lambda key: (key, None)
|
|
cache.get_child_key_fn = lambda key: key.token_ids[0]
|
|
cache.root_node = TreeNode()
|
|
|
|
result = cache.match_prefix(MatchPrefixParams(key=RadixKey([1, 2, 3])))
|
|
|
|
self.assertEqual(result.device_indices.tolist(), [])
|
|
self.assertIs(result.last_device_node, cache.root_node)
|
|
self.assertIs(result.last_host_node, cache.root_node)
|
|
self.assertEqual(result.host_hit_length, 0)
|
|
|
|
def test_non_cp_match_prefix_uses_root_when_no_host_backup_exists(self):
|
|
cache = HiRadixCache.__new__(HiRadixCache)
|
|
cache._uses_cp_hicache = False
|
|
cache.device = "cpu"
|
|
cache.disable = False
|
|
cache.page_size = 1
|
|
cache.get_child_key_fn = lambda key: key.token_ids[0]
|
|
cache.key_match_fn = lambda child_key, key: sum(
|
|
1 for lhs, rhs in zip(child_key.token_ids, key.token_ids) if lhs == rhs
|
|
)
|
|
cache.maybe_bigram_convert = lambda key: (key, None)
|
|
root = TreeNode()
|
|
root.key = RadixKey([])
|
|
root.value = torch.empty((0,), dtype=torch.int64)
|
|
root.host_value = None
|
|
cache.root_node = root
|
|
node = TreeNode()
|
|
node.parent = root
|
|
node.key = RadixKey(list(range(4)))
|
|
node.value = torch.arange(4, dtype=torch.int64)
|
|
node.host_value = None
|
|
root.children[0] = node
|
|
|
|
result = cache.match_prefix(MatchPrefixParams(key=RadixKey(list(range(4)))))
|
|
|
|
self.assertEqual(result.host_hit_length, 0)
|
|
self.assertIs(result.last_host_node, root)
|
|
self.assertIs(result.last_device_node, node)
|
|
|
|
def test_cp_load_back_non_evicted_node_returns_none_without_loading(self):
|
|
cache = HiRadixCache.__new__(HiRadixCache)
|
|
cache._uses_cp_hicache = True
|
|
cache.load_back_threshold = 1
|
|
cache.evictable_size_ = 0
|
|
cache.metrics_collector = None
|
|
cache.ongoing_load_back = {}
|
|
cache.inc_lock_ref = lambda node: type("Result", (), {"delta": 0})()
|
|
cache.dec_lock_ref = lambda node: None
|
|
cache.cache_controller = type(
|
|
"Controller",
|
|
(),
|
|
{
|
|
"load_cp": lambda self, nodes, node_id=-1: (_ for _ in ()).throw(
|
|
AssertionError("must not load without evicted nodes")
|
|
)
|
|
},
|
|
)()
|
|
node = TreeNode()
|
|
node.value = torch.empty((0,), dtype=torch.int64)
|
|
|
|
loaded = cache.load_back(node)
|
|
|
|
self.assertIsNone(loaded)
|
|
self.assertEqual(cache.ongoing_load_back, {})
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|