chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity
Files
dd69c5970cc4ed3b163faffdb018ecefa4eae2ec
sglang/test/registered/unit/mem_cache/test_cp_hicache_metadata.py
laoyao0822 0043037f78 Keep CP HiCache owner lanes padded for valid-tail victims 
Load-back owner-lane eviction can need to evict a device-resident node whose radix value is a valid tail shorter than the physical tail page. The eviction planner now pads that value only to the page boundary before deriving owner counts, matching the existing write/load capacity contract without exposing padding to radix or scheduler lengths.

Constraint: CP HiCache capacity remains page-owner based while radix node values remain valid-token based.
Constraint: Avoid collectives; owner-lane capacity must be deterministic from local metadata and logical page ids.
Rejected: Require device-resident victim values to be page-aligned | valid-tail cache nodes are now an intentional supported state.
Rejected: Pad to cp_size pages | this would waste KV and violate the page-boundary-only contract.
Confidence: medium
Scope-risk: narrow
Directive: If split-inside-tail support is added later, preserve page ownership/refcount semantics before sharing one padded physical page across radix nodes.
Tested: local py_compile for hiradix_cache.py and touched CP HiCache tests.
Tested: remote g0034 new C8 exact tests: 3 passed, 3 warnings.
Tested: remote g0034 CP HiCache impacted suites: 146 passed, 5 warnings.
Tested: remote g0034 CP shared KV C1-C5 suite: 122 passed, 5 warnings.
Not-tested: full local pytest, blocked by missing runtime dependencies such as starlette.
Not-tested: CUDA E2E runtime for this commit.
Co-authored-by: OmX <omx@oh-my-codex.dev>
2026-05-29 05:39:25 +08:00

2497 lines
96 KiB
Python
Raw Blame History

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",
"moe_fused_gate(Tensor input, Tensor bias, int num_expert_group, int topk_group, int topk, int num_fused_shared_experts, float routed_scaling_factor, bool apply_routed_scaling_factor_on_output) -> 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")
_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",
"moe_fused_gate(Tensor input, Tensor bias, int num_expert_group, int topk_group, int topk, int num_fused_shared_experts, float routed_scaling_factor, bool apply_routed_scaling_factor_on_output) -> 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 (
HiCacheAck,
HiCacheWriteFailure,
HiCacheWriteReservation,
)
from sglang.srt.mem_cache.base_prefix_cache import (
EvictParams,
InsertParams,
MatchPrefixParams,
)
from sglang.srt.mem_cache.hiradix_cache import (
CpHiCacheNodeMetadata,
HiRadixCache,
PendingHiCacheBackup,
PreparedCpHiCacheBackup,
_compute_shared_hicache_token_capacities,
)
from sglang.srt.mem_cache.radix_cache import RadixKey, TreeNode, _key_match_paged
from sglang.srt.mem_cache.session_aware_cache import SessionAwareCache
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 TestHiRadixCacheCPDraftHostPool(CustomTestCase):
def test_shared_budget_keeps_draft_at_least_target_capacity(self):
target_tokens, draft_tokens = _compute_shared_hicache_token_capacities(
total_host_bytes=1000,
target_size_per_token=6,
draft_size_per_token=2,
page_size=10,
)
self.assertEqual(target_tokens, 120)
self.assertEqual(draft_tokens, 140)
self.assertGreaterEqual(draft_tokens, target_tokens)
self.assertLessEqual(target_tokens * 6 + draft_tokens * 2, 1000)
def test_shared_budget_handles_equal_target_and_draft_size(self):
target_tokens, draft_tokens = _compute_shared_hicache_token_capacities(
total_host_bytes=1000,
target_size_per_token=6,
draft_size_per_token=6,
page_size=10,
)
self.assertEqual(target_tokens, 80)
self.assertEqual(draft_tokens, 80)
self.assertLessEqual(target_tokens * 6 + draft_tokens * 6, 1000)
def test_reset_clears_target_and_draft_host_pools(self):
class ClearablePool:
def __init__(self):
self.clear_calls = 0
def clear(self):
self.clear_calls += 1
class Controller:
def __init__(self):
self.reset_calls = 0
self.clear_draft_calls = 0
def reset(self):
self.reset_calls += 1
def clear_draft_host_pool(self):
self.clear_draft_calls += 1
target_pool = ClearablePool()
controller = Controller()
cache = HiRadixCache.__new__(HiRadixCache)
cache.cache_controller = controller
cache.token_to_kv_pool_host = target_pool
cache.prefetch_loaded_tokens_by_reqid = {}
cache.evictable_host_leaves = set()
cache.pinned_size_ = 1
cache.evictable_leaves = set()
cache._record_all_cleared_event = lambda: None
cache.reset()
self.assertEqual(controller.reset_calls, 1)
self.assertEqual(target_pool.clear_calls, 1)
self.assertEqual(controller.clear_draft_calls, 1)
class TestCpHiCacheNodeMetadata(CustomTestCase):
def test_paged_key_match_returns_valid_tail_length_not_next_page(self):
key = RadixKey(list(range(6)))
self.assertEqual(_key_match_paged(key, key, page_size=4), 6)
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),
page_owners=torch.zeros(max(8, 0), dtype=torch.int8),
page_size=1,
)
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),
page_owners=torch.zeros(max(10, 0), dtype=torch.int8),
page_size=1,
)
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),
page_owners=torch.zeros(max(10, 0), dtype=torch.int8),
page_size=1,
)
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),
page_owners=torch.zeros(max(64, 0), dtype=torch.int8),
page_size=1,
)
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_valid_length_can_be_shorter_than_physical_padded_length(self):
metadata = CpHiCacheNodeMetadata(
logical_len=100,
padded_len=128,
owned_positions=torch.tensor([0, 63, 100, 127], dtype=torch.int64),
host_indices=torch.tensor([10, 11, 12, 13], dtype=torch.int64),
page_owners=torch.tensor([0, 1], dtype=torch.int8),
page_size=64,
)
self.assertEqual(metadata.logical_len, 100)
self.assertEqual(metadata.valid_len, 100)
self.assertEqual(metadata.padded_len, 128)
self.assertEqual(metadata.page_owners.tolist(), [0, 1])
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),
page_owners=torch.zeros(max(4, 0), dtype=torch.int8),
page_size=1,
)
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),
page_owners=torch.zeros(0, dtype=torch.int8),
page_size=1,
)
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,
page_owners=torch.zeros(max(4, 0), dtype=torch.int8),
page_size=1,
)
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),
page_owners=torch.zeros(max(4, 0), dtype=torch.int8),
page_size=1,
)
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),
page_owners=torch.zeros(max(4, 0), dtype=torch.int8),
page_size=1,
)
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),
page_owners=torch.zeros(max(4, 0), dtype=torch.int8),
page_size=1,
)
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),
page_owners=torch.zeros(max(4, 0), dtype=torch.int8),
page_size=1,
)
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),
page_owners=torch.zeros(max(4, 0), dtype=torch.int8),
page_size=1,
)
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),
page_owners=torch.zeros(max(4, 0), dtype=torch.int8),
page_size=1,
)
# ── New: validators for page_owners + page_size (the fields that carry the
# CP owner pattern across a HiCache write→load round-trip).
def test_zero_page_size_raises(self):
with self.assertRaisesRegex(ValueError, "page_size must be positive"):
CpHiCacheNodeMetadata(
logical_len=4,
owned_positions=torch.empty((0,), dtype=torch.int64),
host_indices=torch.empty((0,), dtype=torch.int64),
page_owners=torch.empty((0,), dtype=torch.int8),
page_size=0,
)
def test_logical_len_not_multiple_of_page_size_raises(self):
with self.assertRaisesRegex(ValueError, "multiple of"):
CpHiCacheNodeMetadata(
logical_len=10,
owned_positions=torch.empty((0,), dtype=torch.int64),
host_indices=torch.empty((0,), dtype=torch.int64),
page_owners=torch.zeros(2, dtype=torch.int8),
page_size=4, # 10 % 4 != 0
)
def test_page_owners_length_mismatch_raises(self):
with self.assertRaisesRegex(ValueError, "page_owners length"):
CpHiCacheNodeMetadata(
logical_len=8,
owned_positions=torch.empty((0,), dtype=torch.int64),
host_indices=torch.empty((0,), dtype=torch.int64),
page_owners=torch.zeros(3, dtype=torch.int8), # expected 8/4=2
page_size=4,
)
def test_negative_page_owners_raises(self):
with self.assertRaisesRegex(ValueError, "page_owners.*non-negative"):
CpHiCacheNodeMetadata(
logical_len=8,
owned_positions=torch.empty((0,), dtype=torch.int64),
host_indices=torch.empty((0,), dtype=torch.int64),
page_owners=torch.tensor([-1, 0], dtype=torch.int8),
page_size=4,
)
def test_page_owners_normalized_to_int8_cpu(self):
metadata = CpHiCacheNodeMetadata(
logical_len=8,
owned_positions=torch.empty((0,), dtype=torch.int64),
host_indices=torch.empty((0,), dtype=torch.int64),
# Pass int64 to test normalization.
page_owners=torch.tensor([0, 1], dtype=torch.int64),
page_size=4,
)
self.assertEqual(metadata.page_owners.dtype, torch.int8)
self.assertEqual(metadata.page_owners.device.type, "cpu")
self.assertEqual(metadata.page_owners.tolist(), [0, 1])
def test_page_owners_not_aliased_to_input(self):
page_owners = torch.tensor([0, 1, 0, 1], dtype=torch.int8)
metadata = CpHiCacheNodeMetadata(
logical_len=16,
owned_positions=torch.empty((0,), dtype=torch.int64),
host_indices=torch.empty((0,), dtype=torch.int64),
page_owners=page_owners,
page_size=4,
)
page_owners[0] = 1
self.assertEqual(metadata.page_owners.tolist(), [0, 1, 0, 1])
def test_split_non_page_aligned_split_len_raises(self):
metadata = CpHiCacheNodeMetadata(
logical_len=16,
owned_positions=torch.tensor([0, 8], dtype=torch.int64),
host_indices=torch.tensor([20, 21], dtype=torch.int64),
page_owners=torch.zeros(4, dtype=torch.int8),
page_size=4,
)
with self.assertRaisesRegex(ValueError, "must be a multiple of page_size"):
metadata.split(5) # 5 % 4 != 0
def test_split_preserves_page_owners_slice(self):
metadata = CpHiCacheNodeMetadata(
logical_len=32, # 8 pages of page_size=4
owned_positions=torch.tensor([0, 4, 16, 28], dtype=torch.int64),
host_indices=torch.tensor([100, 101, 102, 103], dtype=torch.int64),
page_owners=torch.tensor(
[0, 1, 0, 1, 0, 1, 0, 1], dtype=torch.int8
),
page_size=4,
)
parent, child = metadata.split(16) # split at page 4 of 8.
self.assertEqual(parent.page_owners.tolist(), [0, 1, 0, 1])
self.assertEqual(child.page_owners.tolist(), [0, 1, 0, 1])
self.assertEqual(parent.page_size, 4)
self.assertEqual(child.page_size, 4)
def test_split_at_zero_yields_empty_parent_page_owners(self):
metadata = CpHiCacheNodeMetadata(
logical_len=8,
owned_positions=torch.tensor([0], dtype=torch.int64),
host_indices=torch.tensor([50], dtype=torch.int64),
page_owners=torch.tensor([0, 1], dtype=torch.int8),
page_size=4,
)
parent, child = metadata.split(0)
self.assertEqual(parent.page_owners.numel(), 0)
self.assertEqual(child.page_owners.tolist(), [0, 1])
def test_split_at_logical_len_yields_empty_child_page_owners(self):
metadata = CpHiCacheNodeMetadata(
logical_len=8,
owned_positions=torch.tensor([0], dtype=torch.int64),
host_indices=torch.tensor([50], dtype=torch.int64),
page_owners=torch.tensor([0, 1], dtype=torch.int8),
page_size=4,
)
parent, child = metadata.split(8)
self.assertEqual(parent.page_owners.tolist(), [0, 1])
self.assertEqual(child.page_owners.numel(), 0)
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),
page_owners=torch.zeros(max(len(device_indices), 0), dtype=torch.int8),
page_size=1,
)
)
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),
page_owners=torch.zeros(max(len(device_indices), 0), dtype=torch.int8),
page_size=1,
)
)
class FakeReserveWriteController:
write_policy = "write_through"
has_draft_hicache = False
def __init__(self, results):
self.results = list(results)
self.reservations = []
self.submitted = []
self.submit_kwargs = []
self.evicted_host_indices = []
def reserve_write_cp(self, device_indices, priority=None, node_id=-1):
self.reservations.append((device_indices.clone(), node_id))
result = self.results.pop(0)
if not callable(result):
return result
try:
return result(device_indices, node_id=node_id)
except TypeError:
return result(device_indices)
def submit_write_cp_all_layer(self, reservation):
self.submitted.append(reservation)
def submit_write_cp_per_layer(self, reservation, **kwargs):
self.submitted.append(reservation)
self.submit_kwargs.append(kwargs)
def evict_cp_host(self, metadata):
self.evicted_host_indices.append(metadata.host_indices.clone())
return len(metadata.host_indices)
def make_write_reservation(device_indices, node_id=0, host_start=90):
host_indices = torch.arange(host_start, host_start + len(device_indices))
metadata = CpHiCacheNodeMetadata(
logical_len=len(device_indices),
owned_positions=torch.arange(len(device_indices), dtype=torch.int64),
host_indices=host_indices,
page_owners=torch.zeros(max(len(device_indices), 0), dtype=torch.int8),
page_size=1,
)
return HiCacheWriteReservation(
metadata=metadata,
host_indices=host_indices,
physical_device_indices=device_indices.clone(),
node_id=node_id,
)
class FakeEvictionStrategy:
def get_priority(self, node):
return 0
class FakeCpLayout:
def __init__(self, cp_size=4, cp_rank=0, page_size=1):
self.cp_size = cp_size
self.cp_rank = cp_rank
self.page_size = page_size
def owner_for_logical_pages(self, logical_pages):
owners = torch.remainder(logical_pages - 1, self.cp_size)
return torch.where(logical_pages <= 0, torch.full_like(owners, -1), owners)
def logical_locs_to_physical(self, logical_locs):
return logical_locs
def owned_by_this_rank(self, logical_locs):
logical_pages = torch.div(
logical_locs, self.page_size, rounding_mode="floor"
)
return self.owner_for_logical_pages(logical_pages) == self.cp_rank
class FakeEvictDeviceController:
write_policy = "write_through"
def evict_device(self, device_indices):
return len(device_indices)
class FakeTokenAllocator:
def available_size(self):
return 0
def compute_owner_lane_stats(self, page_owners):
if len(page_owners) == 0:
return [], [], []
cp_size = max(int(owner) for owner in page_owners) + 1
required = [0] * cp_size
for owner in page_owners:
required[int(owner)] += 1
available = list(required)
deficits = [0] * cp_size
return required, available, deficits
def allocator_state_str(self):
return "FakeTokenAllocator"
class TestHiRadixCacheCPBackup(CustomTestCase):
def test_session_aware_cache_forwards_cp_hicache_prepare(self):
calls = []
class Inner:
def prepare_write_backup_for_req(self, req):
calls.append(req)
wrapper = SessionAwareCache(Inner())
req = object()
wrapper.prepare_write_backup_for_req(req)
self.assertEqual(calls, [req])
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),
page_owners=torch.zeros(max(8, 0), dtype=torch.int8),
page_size=1,
)
self.assertTrue(cache._node_backuped(node))
def test_node_backuped_rejects_missing_draft_metadata_when_draft_attached(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache._uses_cp_hicache = True
cache.cache_controller = types.SimpleNamespace(
has_draft_hicache=True,
cp_shared_kv_layout=types.SimpleNamespace(cp_rank=2),
)
node = TreeNode()
node.id = 123
node.host_len = 64
node.cp_hicache = CpHiCacheNodeMetadata(
logical_len=64,
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=64,
)
with self.assertRaisesRegex(
RuntimeError, "node_id=123.*cp_rank=2.*missing draft_host_indices"
):
cache._node_backuped(node)
def test_node_backuped_rejects_cp_host_len_without_metadata(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache._uses_cp_hicache = True
cache.cache_controller = types.SimpleNamespace(
cp_shared_kv_layout=types.SimpleNamespace(cp_rank=1)
)
node = TreeNode()
node.id = 124
node.host_len = 64
node.cp_hicache = None
with self.assertRaisesRegex(
RuntimeError, "node_id=124.*host_len=64.*cp_rank=1.*without cp_hicache"
):
cache._node_backuped(node)
def test_node_backuped_accepts_empty_draft_metadata_for_zero_owned_rank(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache._uses_cp_hicache = True
cache.cache_controller = types.SimpleNamespace(has_draft_hicache=True)
node = TreeNode()
node.host_len = 64
node.cp_hicache = CpHiCacheNodeMetadata(
logical_len=64,
owned_positions=torch.empty((0,), dtype=torch.int64),
host_indices=torch.empty((0,), dtype=torch.int64),
draft_host_indices=torch.empty((0,), dtype=torch.int64),
page_owners=torch.tensor([0], dtype=torch.int8),
page_size=64,
)
self.assertTrue(cache._node_backuped(node))
def test_node_backuped_excludes_inflight_cp_write_until_ack(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache._uses_cp_hicache = True
cache.cache_controller = types.SimpleNamespace(has_draft_hicache=True)
cache.ongoing_write_through = {}
node = TreeNode()
node.id = 125
node.host_len = 64
node.cp_hicache = CpHiCacheNodeMetadata(
logical_len=64,
owned_positions=torch.empty((0,), dtype=torch.int64),
host_indices=torch.empty((0,), dtype=torch.int64),
draft_host_indices=torch.empty((0,), dtype=torch.int64),
page_owners=torch.tensor([0], dtype=torch.int8),
page_size=64,
)
cache.ongoing_write_through[node.id] = node
self.assertFalse(cache._node_backuped(node))
cache.ongoing_write_through.clear()
self.assertTrue(cache._node_backuped(node))
def test_node_backuped_excludes_explicit_pending_cp_backup_until_commit(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache._uses_cp_hicache = True
cache.cache_controller = types.SimpleNamespace(has_draft_hicache=False)
cache.ongoing_write_through = {}
cache.pending_host_backups = {}
dec_locked = []
cache.dec_node_lock_ref = lambda node: dec_locked.append(node)
node = TreeNode()
node.id = 128
node.host_len = 64
metadata = CpHiCacheNodeMetadata(
logical_len=64,
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=64,
)
node.cp_hicache = metadata
cache.pending_host_backups[node.id] = PendingHiCacheBackup(
node=node, metadata=metadata, logical_len=64
)
self.assertFalse(cache._node_backuped(node))
cache._commit_pending_backup(node.id)
self.assertTrue(cache._node_backuped(node))
self.assertEqual(dec_locked, [node])
def test_rollback_pending_cp_backup_frees_reserved_host_slots(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache._uses_cp_hicache = True
evicted = []
cache.cache_controller = types.SimpleNamespace(
evict_cp_host=lambda metadata: evicted.append(metadata) or 2
)
dec_locked = []
cache.dec_node_lock_ref = lambda node: dec_locked.append(node)
cache.pending_host_backups = {}
node = TreeNode()
node.id = 129
metadata = CpHiCacheNodeMetadata(
logical_len=64,
owned_positions=torch.tensor([0, 1], dtype=torch.int64),
host_indices=torch.tensor([10, 11], dtype=torch.int64),
page_owners=torch.tensor([0], dtype=torch.int8),
page_size=64,
)
cache.pending_host_backups[node.id] = PendingHiCacheBackup(
node=node, metadata=metadata, logical_len=64
)
cache._rollback_pending_backup(node.id)
self.assertEqual(evicted, [metadata])
self.assertEqual(dec_locked, [node])
self.assertNotIn(node.id, cache.pending_host_backups)
def test_single_node_write_lock_updates_device_evictable_leaf_set(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache.disable = False
cache.root_node = TreeNode()
cache.root_node.key = RadixKey([])
cache.evictable_leaves = set()
cache.evictable_size_ = 4
cache.protected_size_ = 0
node = TreeNode()
node.parent = cache.root_node
node.key = RadixKey([1, 2, 3, 4])
node.value = torch.arange(4, dtype=torch.int64)
cache.root_node.children[1] = node
cache.evictable_leaves.add(node)
cache.inc_node_lock_ref(node)
self.assertNotIn(node, cache.evictable_leaves)
self.assertEqual(cache.evictable_size(), 0)
self.assertEqual(cache.protected_size(), 4)
cache.dec_node_lock_ref(node)
self.assertIn(node, cache.evictable_leaves)
self.assertEqual(cache.evictable_size(), 4)
self.assertEqual(cache.protected_size(), 0)
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),
page_owners=torch.zeros(max(4, 0), dtype=torch.int8),
page_size=1,
)
cache._inc_hit_count(node)
self.assertEqual(node.hit_count, 1)
def test_inc_hit_count_does_not_duplicate_inflight_cp_write(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache._uses_cp_hicache = True
cache.write_through_threshold = 1
cache.ongoing_write_through = {}
cache.cache_controller = type(
"Controller", (), {"write_policy": "write_through"}
)()
cache.write_backup = lambda node: (_ for _ in ()).throw(
AssertionError("must not launch a second write")
)
node = TreeNode()
node.id = 127
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),
draft_host_indices=torch.tensor([], dtype=torch.int64),
page_owners=torch.zeros(max(4, 0), dtype=torch.int8),
page_size=1,
)
cache.ongoing_write_through[node.id] = node
cache._inc_hit_count(node)
self.assertEqual(node.hit_count, 1)
def test_write_backup_uses_deterministic_host_eviction_before_reserve(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache._uses_cp_hicache = True
cache.page_size = 1
reservation_factory = lambda device_indices: make_write_reservation(
device_indices, node_id=122
)
cache.cache_controller = FakeReserveWriteController([reservation_factory])
cache.cache_controller.cp_shared_kv_layout = FakeCpLayout(cp_size=1, cp_rank=0)
cache.token_to_kv_pool_host = types.SimpleNamespace(size=16)
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._update_host_leaf_status = lambda node: None
cache.ongoing_write_through = {}
cache.pending_host_backups = {}
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.arange(4, dtype=torch.int64),
host_indices=torch.arange(55, 59, dtype=torch.int64),
page_owners=torch.zeros(max(4, 0), dtype=torch.int8),
page_size=1,
)
root.children[1] = evictable_node
cache.evictable_host_leaves.add(evictable_node)
node = TreeNode()
node.id = 122
node.value = torch.arange(16, dtype=torch.int64)
cache.write_backup(node)
self.assertEqual(
cache.cache_controller.evicted_host_indices[0].tolist(),
list(range(55, 59)),
)
self.assertEqual(evictable_node.host_len, 0)
self.assertIsNone(evictable_node.cp_hicache)
self.assertNotIn(1, root.children)
self.assertEqual(node.host_len, 0)
self.assertIsNone(node.cp_hicache)
self.assertIn(node.id, cache.pending_host_backups)
self.assertEqual(len(cache.cache_controller.submitted), 1)
def test_write_backup_cp_post_forward_path_warns_fallback(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache._uses_cp_hicache = True
cache.cache_controller = FakeReserveWriteController(
[
lambda device_indices, node_id: make_write_reservation(
device_indices, node_id=node_id
)
]
)
cache.ongoing_write_through = {}
cache.pending_host_backups = {}
cache.inc_node_lock_ref = lambda node: None
node = TreeNode()
node.id = 133
node.value = torch.arange(16, dtype=torch.int64)
with self.assertLogs(
"sglang.srt.mem_cache.hiradix_cache", level="WARNING"
) as captured:
cache.write_backup(node)
self.assertTrue(
any(
"[CP_HICACHE_FALLBACK][post_forward_catch_up_backup]" in message
for message in captured.output
)
)
def test_prepare_write_backup_for_req_chunked_skip_warns_fallback(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache.disable = False
cache._uses_cp_hicache = True
cache.cache_controller = FakeReserveWriteController([])
req = types.SimpleNamespace(
rid="rid-chunked",
is_chunked=1,
cp_hicache_prepared_backup=None,
)
with self.assertLogs(
"sglang.srt.mem_cache.hiradix_cache", level="WARNING"
) as captured:
cache.prepare_write_backup_for_req(req)
self.assertTrue(
any(
"[CP_HICACHE_FALLBACK][prepare_write_backup_skipped]" in message
and "reason=chunked_req" in message
for message in captured.output
)
)
def test_write_backup_deterministic_eviction_avoids_reserve_all_reduce(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache._uses_cp_hicache = True
cache.page_size = 1
cache.tp_world_size = 2
cache.tp_group = object()
cache.cache_controller = FakeReserveWriteController(
[
lambda device_indices: make_write_reservation(
device_indices, node_id=130, host_start=100
),
]
)
cache.cache_controller.cp_shared_kv_layout = FakeCpLayout(cp_size=1, cp_rank=0)
cache.token_to_kv_pool_host = types.SimpleNamespace(size=16)
evictions = []
cache._evict_host_for_physical_slots = lambda required, synchronize_across_ranks=False: (
evictions.append((required, synchronize_across_ranks)) or required
)
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._update_host_leaf_status = lambda node: None
cache.ongoing_write_through = {}
cache.pending_host_backups = {}
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=[2], extra_key=None)
evictable_node.value = None
evictable_node.host_len = 16
evictable_node.cp_hicache = CpHiCacheNodeMetadata(
logical_len=16,
owned_positions=torch.arange(16, dtype=torch.int64),
host_indices=torch.arange(90, 106, dtype=torch.int64),
page_owners=torch.zeros(16, dtype=torch.int8),
page_size=1,
)
root.children[2] = evictable_node
cache.evictable_host_leaves.add(evictable_node)
node = TreeNode()
node.id = 130
node.value = torch.arange(16, dtype=torch.int64)
with patch(
"torch.distributed.all_reduce",
side_effect=AssertionError("reserve admission must not all_reduce"),
):
backed_len = cache.write_backup(node)
self.assertEqual(backed_len, 16)
self.assertEqual(
cache.cache_controller.evicted_host_indices[0].tolist(),
list(range(90, 106)),
)
self.assertEqual(evictions, [])
self.assertEqual(len(cache.cache_controller.submitted), 1)
self.assertEqual(
cache.cache_controller.submitted[0].metadata.host_indices.tolist(),
list(range(100, 116)),
)
def test_insert_attaches_prepared_cp_backup_without_catchup_copy(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache._uses_cp_hicache = True
cache.root_node = TreeNode()
cache.root_node.key = RadixKey([])
cache.evictable_size_ = 0
cache.get_child_key_fn = lambda key: key.token_ids[0]
cache.key_match_fn = lambda lhs, rhs: 0
cache.maybe_bigram_convert = lambda key, value: (key, value)
cache.is_eagle = False
cache.enable_storage = False
cache.enable_kv_cache_events = False
cache._update_leaf_status = lambda node: None
cache._update_host_leaf_status = lambda node: None
cache._record_store_event = lambda node: None
cache.ongoing_write_through = {}
cache.pending_host_backups = {}
locked_nodes = []
cache.inc_node_lock_ref = lambda node: locked_nodes.append(node)
cache.cache_controller = types.SimpleNamespace(write_policy="write_through")
value = torch.arange(16, dtype=torch.int64)
reservation = make_write_reservation(value, node_id=131, host_start=110)
prepared = PreparedCpHiCacheBackup(
node_id=131,
reservation=reservation,
metadata=reservation.metadata,
logical_len=16,
)
result = cache.insert(
InsertParams(
key=RadixKey(list(range(16))),
value=value,
cp_hicache_prepared_backup=prepared,
)
)
node = cache.root_node.children[0]
self.assertEqual(result.prefix_len, 0)
self.assertEqual(node.id, 131)
self.assertTrue(prepared.attached)
self.assertIs(cache.ongoing_write_through[131], node)
self.assertIs(cache.pending_host_backups[131].node, node)
self.assertIs(cache.pending_host_backups[131].metadata, reservation.metadata)
self.assertEqual(cache.pending_host_backups[131].logical_len, 16)
self.assertEqual(locked_nodes, [node])
def test_prepare_write_backup_for_req_registers_before_forward_without_catchup(
self,
):
cache = HiRadixCache.__new__(HiRadixCache)
cache.disable = False
cache._uses_cp_hicache = True
cache.is_eagle = False
cache.page_size = 1
cache.req_to_token_pool = types.SimpleNamespace(
req_to_token=torch.arange(16, dtype=torch.int64).view(1, 16)
)
cache.cache_controller = FakeReserveWriteController(
[
lambda device_indices, node_id: make_write_reservation(
device_indices, node_id=node_id, host_start=130
)
]
)
req = types.SimpleNamespace(
rid="rid-prepare",
fill_ids=list(range(16)),
cache_protected_len=0,
req_pool_idx=0,
is_chunked=0,
cp_hicache_prepared_backup=None,
)
cache.prepare_write_backup_for_req(req)
prepared = req.cp_hicache_prepared_backup
self.assertIsNotNone(prepared)
self.assertEqual(prepared.logical_len, 16)
self.assertEqual(len(cache.cache_controller.submitted), 1)
self.assertIs(cache.cache_controller.submitted[0], prepared.reservation)
self.assertEqual(
cache.cache_controller.submit_kwargs,
[{"catch_up_all_layers": False}],
)
def test_prepare_write_backup_for_req_keeps_valid_tail_length(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache.disable = False
cache._uses_cp_hicache = True
cache.is_eagle = False
cache.page_size = 4
cache.req_to_token_pool = types.SimpleNamespace(
req_to_token=torch.arange(8, dtype=torch.int64).view(1, 8)
)
cache.cache_controller = FakeReserveWriteController(
[
lambda device_indices, node_id: make_write_reservation(
device_indices, node_id=node_id, host_start=170
)
]
)
cache.maybe_bigram_convert = lambda key: (key, None)
cache.root_node = TreeNode()
cache.root_node.children = {}
cache.get_child_key_fn = lambda key: tuple(key.token_ids[:4])
req = types.SimpleNamespace(
rid="rid-tail-prepare",
fill_ids=list(range(6)),
cache_protected_len=0,
req_pool_idx=0,
is_chunked=0,
cp_hicache_prepared_backup=None,
)
cache.prepare_write_backup_for_req(req)
self.assertIsNotNone(req.cp_hicache_prepared_backup)
self.assertEqual(req.cp_hicache_prepared_backup.logical_len, 6)
self.assertEqual(cache.cache_controller.reservations[0][0].tolist(), list(range(6)))
def test_cache_finished_req_keeps_cp_valid_tail_insert_key(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache.disable_finished_insert = False
cache.disable = False
cache._uses_cp_hicache = True
cache.is_eagle = False
cache.page_size = 4
cache.req_to_token_pool = types.SimpleNamespace(
req_to_token=torch.arange(8, dtype=torch.int64).view(1, 8)
)
freed = []
cache.token_to_kv_pool_allocator = types.SimpleNamespace(
free=lambda indices: freed.append(indices.clone())
)
inserted = []
cache.insert = lambda params: inserted.append(params) or types.SimpleNamespace(
prefix_len=0
)
cache.dec_lock_ref = lambda node: None
req = types.SimpleNamespace(
origin_input_ids=list(range(6)),
output_ids=[],
req_pool_idx=0,
extra_key=None,
cache_protected_len=0,
last_node=TreeNode(),
cp_hicache_prepared_backup=None,
pop_committed_kv_cache=lambda: 6,
)
cache.cache_finished_req(req)
self.assertEqual(len(inserted), 1)
self.assertEqual(inserted[0].key.token_ids, list(range(6)))
self.assertEqual(inserted[0].value.tolist(), list(range(6)))
self.assertEqual([indices.tolist() for indices in freed], [[], []])
def test_prepare_write_backup_for_req_skips_existing_insert_prefix(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache.disable = False
cache._uses_cp_hicache = True
cache.is_eagle = False
cache.page_size = 1
cache.root_node = TreeNode()
cache.root_node.key = RadixKey([])
cache.root_node.children = {}
cache.get_child_key_fn = lambda key: key.token_ids[0]
def key_match(lhs, rhs):
matched = 0
for left, right in zip(lhs.token_ids, rhs.token_ids):
if left != right:
break
matched += 1
return matched
cache.key_match_fn = key_match
cache.maybe_bigram_convert = lambda key, value=None: (key, value)
cache.req_to_token_pool = types.SimpleNamespace(
req_to_token=torch.arange(16, dtype=torch.int64).view(1, 16)
)
cache.cache_controller = FakeReserveWriteController([])
existing = TreeNode()
existing.id = 201
existing.parent = cache.root_node
existing.key = RadixKey(list(range(16)))
existing.value = torch.arange(16, dtype=torch.int64)
cache.root_node.children[0] = existing
req = types.SimpleNamespace(
rid="rid-existing-prefix",
fill_ids=list(range(16)),
extra_key=None,
cache_protected_len=8,
req_pool_idx=0,
is_chunked=0,
cp_hicache_prepared_backup=None,
)
cache.prepare_write_backup_for_req(req)
self.assertIsNone(req.cp_hicache_prepared_backup)
self.assertEqual(cache.cache_controller.reservations, [])
self.assertEqual(cache.cache_controller.submitted, [])
def test_rollback_unattached_prepared_cp_backup_removes_orphan_ack(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache._uses_cp_hicache = True
value = torch.arange(4, dtype=torch.int64)
reservation = make_write_reservation(value, node_id=132, host_start=120)
prepared = PreparedCpHiCacheBackup(
node_id=132,
reservation=reservation,
metadata=reservation.metadata,
logical_len=4,
)
evicted = []
class ReadyEvent:
def synchronize(self):
pass
cache.cache_controller = types.SimpleNamespace(
pending_layer_writes={},
ack_write_queue=[HiCacheAck(ReadyEvent(), ReadyEvent(), [132])],
evict_cp_host=lambda metadata: evicted.append(metadata),
)
cache._rollback_prepared_cp_backup(prepared, "test")
self.assertEqual(cache.cache_controller.ack_write_queue, [])
self.assertEqual(evicted, [reservation.metadata])
def test_write_backup_cp_failfast_on_unplanned_reservation_failure(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache._uses_cp_hicache = True
cache.cache_controller = FakeReserveWriteController(
[
HiCacheWriteFailure(required_host_slots=2),
HiCacheWriteFailure(required_host_slots=2),
]
)
cache.evictable_host_leaves = set()
cache.eviction_strategy = FakeEvictionStrategy()
cache._record_remove_event = lambda node: None
cache.ongoing_write_through = {}
cache.pending_host_backups = {}
node = TreeNode()
node.id = 124
node.value = torch.arange(16, dtype=torch.int64)
with self.assertRaisesRegex(
RuntimeError,
"owner-lane admission predicted no host deficit",
):
cache.write_backup(node)
self.assertEqual(node.host_len, 0)
self.assertIsNone(node.cp_hicache)
self.assertEqual(cache.pending_host_backups, {})
self.assertEqual(cache.cache_controller.submitted, [])
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 = FakeReserveWriteController(
[lambda device_indices: HiCacheWriteReservation(
metadata=CpHiCacheNodeMetadata(
logical_len=len(device_indices),
owned_positions=torch.empty((0,), dtype=torch.int64),
host_indices=torch.empty((0,), dtype=torch.int64),
page_owners=torch.zeros(max(len(device_indices), 0), dtype=torch.int8),
page_size=1,
),
host_indices=torch.empty((0,), dtype=torch.int64),
physical_device_indices=torch.empty((0,), dtype=torch.int64),
node_id=123,
)]
)
cache.ongoing_write_through = {}
cache.pending_host_backups = {}
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, 0)
self.assertIsNone(node.cp_hicache)
self.assertIn(node.id, cache.pending_host_backups)
self.assertEqual(len(cache.cache_controller.submitted), 1)
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),
page_owners=torch.zeros(max(4, 0), dtype=torch.int8),
page_size=1,
)
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])
def test_cp_owner_token_counts_use_page_owners_and_page_size(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache.cache_controller = types.SimpleNamespace(
cp_shared_kv_layout=FakeCpLayout(cp_size=4, cp_rank=0)
)
counts = cache._cp_owner_token_counts(
torch.tensor([0, 1, 1, 2], dtype=torch.int8),
page_size=64,
)
self.assertEqual(counts, (64, 128, 64, 0))
def test_cp_metadata_count_asserts_local_owner_lane(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache.cache_controller = types.SimpleNamespace(
cp_shared_kv_layout=FakeCpLayout(cp_size=3, cp_rank=1),
has_draft_hicache=True,
)
metadata = CpHiCacheNodeMetadata(
logical_len=192,
owned_positions=torch.arange(64, dtype=torch.int64),
host_indices=torch.arange(100, 164, dtype=torch.int64),
draft_host_indices=torch.arange(200, 264, dtype=torch.int64),
page_owners=torch.tensor([0, 1, 2], dtype=torch.int8),
page_size=64,
)
counts = cache._cp_assert_metadata_counts(metadata, context="unit")
self.assertEqual(counts, (64, 64, 64))
def test_cp_metadata_count_asserts_mismatched_local_slots(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache.cache_controller = types.SimpleNamespace(
cp_shared_kv_layout=FakeCpLayout(cp_size=2, cp_rank=1),
has_draft_hicache=False,
)
metadata = CpHiCacheNodeMetadata(
logical_len=128,
owned_positions=torch.empty((0,), dtype=torch.int64),
host_indices=torch.empty((0,), dtype=torch.int64),
page_owners=torch.tensor([0, 1], dtype=torch.int8),
page_size=64,
)
with self.assertRaisesRegex(RuntimeError, "owner lane mismatch"):
cache._cp_assert_metadata_counts(metadata, context="unit")
def test_cp_host_capacity_snapshot_counts_committed_and_pending_draft(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache._uses_cp_hicache = True
cache.cache_controller = types.SimpleNamespace(
cp_shared_kv_layout=FakeCpLayout(cp_size=2, cp_rank=0),
has_draft_hicache=True,
draft_mem_pool_host=types.SimpleNamespace(size=512),
)
cache.token_to_kv_pool_host = types.SimpleNamespace(size=512)
cache.pending_host_backups = {}
root = TreeNode()
root.children = {}
cache.root_node = root
committed = TreeNode()
committed.id = 501
committed.parent = root
committed.key = RadixKey([1])
committed.host_len = 128
committed.cp_hicache = CpHiCacheNodeMetadata(
logical_len=128,
owned_positions=torch.arange(64, dtype=torch.int64),
host_indices=torch.arange(64, dtype=torch.int64),
draft_host_indices=torch.arange(100, 164, dtype=torch.int64),
page_owners=torch.tensor([0, 1], dtype=torch.int8),
page_size=64,
)
root.children[1] = committed
pending = TreeNode()
pending.id = 502
pending_metadata = CpHiCacheNodeMetadata(
logical_len=64,
owned_positions=torch.arange(64, dtype=torch.int64),
host_indices=torch.arange(200, 264, dtype=torch.int64),
draft_host_indices=torch.arange(300, 364, dtype=torch.int64),
page_owners=torch.tensor([0], dtype=torch.int8),
page_size=64,
)
cache.pending_host_backups[pending.id] = PendingHiCacheBackup(
node=pending,
metadata=pending_metadata,
logical_len=64,
)
snapshot = cache._cp_host_capacity_snapshot()
self.assertEqual(snapshot.target_capacity, (512, 512))
self.assertEqual(snapshot.draft_capacity, (512, 512))
self.assertEqual(snapshot.committed_target, (64, 64))
self.assertEqual(snapshot.committed_draft, (64, 64))
self.assertEqual(snapshot.pending_target, (64, 0))
self.assertEqual(snapshot.pending_draft, (64, 0))
def test_cp_write_admission_uses_padded_owner_vector_for_valid_tail(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache._uses_cp_hicache = True
cache.page_size = 64
cache.cache_controller = types.SimpleNamespace(
cp_shared_kv_layout=FakeCpLayout(cp_size=8, cp_rank=0, page_size=64),
has_draft_hicache=True,
draft_mem_pool_host=types.SimpleNamespace(size=1024),
)
cache.token_to_kv_pool_host = types.SimpleNamespace(size=1024)
cache.pending_host_backups = {}
cache.root_node = TreeNode()
cache.root_node.children = {}
cache.evictable_host_leaves = set()
admission = cache._cp_build_write_admission(
torch.arange(64, 164, dtype=torch.int64),
node_id=503,
phase="unit",
)
self.assertEqual(admission.required_by_owner, (64, 64, 0, 0, 0, 0, 0, 0))
self.assertEqual(
admission.target_available_by_owner,
(1024, 1024, 1024, 1024, 1024, 1024, 1024, 1024),
)
self.assertEqual(
admission.target_available_by_owner, admission.draft_available_by_owner
)
self.assertEqual(admission.deficit_by_owner, (0, 0, 0, 0, 0, 0, 0, 0))
def test_cp_capacity_snapshot_counts_padded_valid_tail_for_target_and_draft(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache._uses_cp_hicache = True
cache.cache_controller = types.SimpleNamespace(
cp_shared_kv_layout=FakeCpLayout(cp_size=2, cp_rank=0, page_size=64),
has_draft_hicache=True,
draft_mem_pool_host=types.SimpleNamespace(size=1024),
)
cache.token_to_kv_pool_host = types.SimpleNamespace(size=1024)
cache.pending_host_backups = {}
root = TreeNode()
root.children = {}
cache.root_node = root
committed = TreeNode()
committed.id = 504
committed.parent = root
committed.key = RadixKey(list(range(100)))
committed.host_len = 100
committed.cp_hicache = CpHiCacheNodeMetadata(
logical_len=100,
padded_len=128,
owned_positions=torch.arange(64, dtype=torch.int64),
host_indices=torch.arange(100, 164, dtype=torch.int64),
draft_host_indices=torch.arange(200, 264, dtype=torch.int64),
page_owners=torch.tensor([0, 1], dtype=torch.int8),
page_size=64,
)
root.children[0] = committed
snapshot = cache._cp_host_capacity_snapshot()
self.assertEqual(snapshot.committed_target, (64, 64))
self.assertEqual(snapshot.committed_draft, (64, 64))
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),
page_owners=torch.zeros(max(10, 0), dtype=torch.int8),
page_size=1,
)
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_leaf_status_skips_device_valid_host_backed_node(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache._uses_cp_hicache = True
cache.evictable_host_leaves = set()
node = TreeNode()
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),
page_owners=torch.zeros(max(4, 0), dtype=torch.int8),
page_size=1,
)
cache.evictable_host_leaves.add(node)
cache._update_host_leaf_status(node)
self.assertNotIn(node, cache.evictable_host_leaves)
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),
page_owners=torch.zeros(max(4, 0), dtype=torch.int8),
page_size=1,
)
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),
page_owners=torch.zeros(max(4, 0), dtype=torch.int8),
page_size=1,
)
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),
page_owners=torch.zeros(max(4, 0), dtype=torch.int8),
page_size=1,
)
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._tp_group_rank = 0
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),
page_owners=torch.zeros(max(4, 0), dtype=torch.int8),
page_size=1,
)
cache.root_node.children[1] = node
cache.evictable_host_leaves.add(node)
all_done_states = iter([0, 1])
def fake_all_reduce(done, op=None, group=None):
done.fill_(next(all_done_states, 1))
with patch("torch.distributed.all_reduce", side_effect=fake_all_reduce):
physical_freed = cache._evict_host_for_physical_slots(
0, synchronize_across_ranks=True
)
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),
page_owners=torch.zeros(max(4, 0), dtype=torch.int8),
page_size=1,
)
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)
def test_cp_host_eviction_plan_is_stable_across_leaf_insertion_order(self):
def make_cache(order):
cache = HiRadixCache.__new__(HiRadixCache)
cache._uses_cp_hicache = True
cache.cache_controller = types.SimpleNamespace(
cp_shared_kv_layout=FakeCpLayout(cp_size=2, cp_rank=0),
has_draft_hicache=False,
)
cache.eviction_strategy = FakeEvictionStrategy()
cache.pending_host_backups = {}
cache.root_node = TreeNode()
cache.root_node.children = {}
cache.get_child_key_fn = lambda key: key.token_ids[0]
cache._clear_pin = lambda node: None
nodes = []
for node_id in (42, 41):
node = TreeNode(id=node_id)
node.parent = cache.root_node
node.key = RadixKey([node_id])
node.value = None
node.lock_ref = 0
node.host_ref_counter = 0
node.host_len = 64
node.cp_hicache = CpHiCacheNodeMetadata(
logical_len=64,
owned_positions=torch.arange(64, dtype=torch.int64),
host_indices=torch.arange(node_id * 100, node_id * 100 + 64),
page_owners=torch.tensor([0], dtype=torch.int8),
page_size=64,
)
cache.root_node.children[node_id] = node
nodes.append(node)
cache.evictable_host_leaves = {nodes[i] for i in order}
return cache
first = make_cache([0, 1])._plan_cp_host_evictions((64, 0))
second = make_cache([1, 0])._plan_cp_host_evictions((64, 0))
self.assertEqual([node.id for node in first.victims], [41])
self.assertEqual([node.id for node in second.victims], [41])
self.assertEqual(first.planned_freed, (64, 0))
self.assertEqual(first.remaining_deficit, (0, 0))
def test_cp_host_eviction_plan_skips_pending_backup_nodes(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache._uses_cp_hicache = True
cache.cache_controller = types.SimpleNamespace(
cp_shared_kv_layout=FakeCpLayout(cp_size=1, cp_rank=0),
has_draft_hicache=False,
)
cache.eviction_strategy = FakeEvictionStrategy()
cache.root_node = TreeNode()
cache.root_node.children = {}
cache.get_child_key_fn = lambda key: key.token_ids[0]
cache.pending_host_backups = {}
cache.evictable_host_leaves = set()
for node_id in (10, 11):
node = TreeNode(id=node_id)
node.parent = cache.root_node
node.key = RadixKey([node_id])
node.value = None
node.lock_ref = 0
node.host_ref_counter = 0
node.host_len = 64
node.cp_hicache = CpHiCacheNodeMetadata(
logical_len=64,
owned_positions=torch.arange(64, dtype=torch.int64),
host_indices=torch.arange(node_id * 100, node_id * 100 + 64),
page_owners=torch.tensor([0], dtype=torch.int8),
page_size=64,
)
cache.root_node.children[node_id] = node
cache.evictable_host_leaves.add(node)
if node_id == 10:
cache.pending_host_backups[node.id] = PendingHiCacheBackup(
node=node,
metadata=node.cp_hicache,
logical_len=64,
)
plan = cache._plan_cp_host_evictions((64,))
self.assertEqual([node.id for node in plan.victims], [11])
self.assertEqual(plan.planned_freed, (64,))
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.page_size = 1
cache.token_to_kv_pool_allocator = FakeTokenAllocator()
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),
page_owners=torch.zeros(max(8, 0), dtype=torch.int8),
page_size=1,
)
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.page_size = 1
cache.token_to_kv_pool_allocator = FakeTokenAllocator()
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),
page_owners=torch.zeros(max(6, 0), dtype=torch.int8),
page_size=1,
)
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),
page_owners=torch.zeros(max(8, 0), dtype=torch.int8),
page_size=1,
)
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_device_valid_host_missing_remains_device_hit(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 = torch.arange(8, dtype=torch.int64)
node.host_value = None
node.host_len = 0
node.cp_hicache = None
root.children[0] = node
result = cache.match_prefix(MatchPrefixParams(key=RadixKey(list(range(8)))))
self.assertEqual(result.device_indices.tolist(), list(range(8)))
self.assertEqual(result.host_hit_length, 0)
self.assertIs(result.last_device_node, node)
self.assertIs(result.last_host_node, root)
def test_cp_match_prefix_does_not_admit_inflight_write_as_host_hit(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache._uses_cp_hicache = True
cache.device = "cpu"
cache.disable = False
cache.page_size = 1
cache.ongoing_write_through = {}
cache.cache_controller = types.SimpleNamespace(has_draft_hicache=True)
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.id = 126
node.parent = root
node.key = RadixKey(list(range(8)))
node.value = torch.arange(8, dtype=torch.int64)
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),
draft_host_indices=torch.tensor([150, 151], dtype=torch.int64),
page_owners=torch.zeros(max(8, 0), dtype=torch.int8),
page_size=1,
)
root.children[0] = node
cache.ongoing_write_through[node.id] = node
result = cache.match_prefix(MatchPrefixParams(key=RadixKey(list(range(8)))))
self.assertEqual(result.device_indices.tolist(), list(range(8)))
self.assertEqual(result.host_hit_length, 0)
self.assertIs(result.last_device_node, node)
self.assertIs(result.last_host_node, root)
def test_cp_match_prefix_defers_split_when_child_backup_pending(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache._uses_cp_hicache = True
cache.device = "cpu"
cache.disable = False
cache.page_size = 1
cache.ongoing_write_through = {}
cache.pending_host_backups = {}
cache.cache_controller = types.SimpleNamespace(has_draft_hicache=False)
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
child = TreeNode()
child.id = 130
child.parent = root
child.key = RadixKey(list(range(8)))
child.value = torch.arange(8, dtype=torch.int64)
child.host_len = 0
metadata = CpHiCacheNodeMetadata(
logical_len=8,
owned_positions=torch.arange(8, dtype=torch.int64),
host_indices=torch.arange(50, 58, dtype=torch.int64),
page_owners=torch.zeros(8, dtype=torch.int8),
page_size=1,
)
cache.pending_host_backups[child.id] = PendingHiCacheBackup(
node=child, metadata=metadata, logical_len=8
)
root.children[0] = child
result = cache.match_prefix(MatchPrefixParams(key=RadixKey([0, 1, 2, 9])))
self.assertIs(result.pending_backup_deferred_node, child)
self.assertIs(root.children[0], child)
self.assertEqual(child.key.token_ids, list(range(8)))
self.assertEqual(result.device_indices.tolist(), [])
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_cp_match_prefix_reports_valid_tail_host_hit(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache._uses_cp_hicache = True
cache.device = "cpu"
cache.disable = False
cache.page_size = 4
cache.ongoing_write_through = {}
cache.pending_host_backups = {}
cache.cache_controller = types.SimpleNamespace(has_draft_hicache=False)
cache.get_child_key_fn = lambda key: tuple(key.token_ids[:4])
cache.key_match_fn = lambda child_key, key: _key_match_paged(
child_key, key, page_size=4
)
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.id = 140
node.parent = root
node.key = RadixKey(list(range(6)))
node.value = None
node.host_value = None
node.host_len = 6
node.cp_hicache = CpHiCacheNodeMetadata(
logical_len=6,
padded_len=8,
owned_positions=torch.tensor([0, 1, 2, 3], dtype=torch.int64),
host_indices=torch.tensor([50, 51, 52, 53], dtype=torch.int64),
page_owners=torch.tensor([0, 1], dtype=torch.int8),
page_size=4,
)
root.children[(0, 1, 2, 3)] = node
result = cache.match_prefix(MatchPrefixParams(key=RadixKey(list(range(6)))))
self.assertEqual(result.device_indices.tolist(), [])
self.assertEqual(result.host_hit_length, 6)
self.assertIs(result.last_device_node, root)
self.assertIs(result.last_host_node, node)
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, {})
class TestCPHiCacheLayerBackupNotifications(CustomTestCase):
def _call_store_index_fast_path(self, out_loc):
from sglang.srt.layers.attention.nsa import nsa_indexer
notifications = []
fused_calls = []
class FakePool:
page_size = 64
start_layer = 0
def get_index_k_with_scale_buffer(self, layer_id):
return torch.empty((1,), dtype=torch.uint8)
def notify_layer_kv_stored_for_backup(self, layer_id, source="kv"):
notifications.append((layer_id, source))
forward_batch = types.SimpleNamespace(token_to_kv_pool=FakePool())
indexer = types.SimpleNamespace()
def fake_fused_store(key, buf, loc, page_size):
fused_calls.append((key, buf, loc.clone(), page_size))
with (
patch.object(nsa_indexer, "_is_cuda", True),
patch.object(nsa_indexer, "_is_fp8_fnuz", False),
patch.object(nsa_indexer, "can_use_nsa_fused_store", return_value=True),
patch.object(
nsa_indexer, "fused_store_index_k_cache", side_effect=fake_fused_store
),
):
nsa_indexer.Indexer._store_index_k_cache(
indexer,
forward_batch,
layer_id=5,
key=torch.empty((max(out_loc.numel(), 1), 128), dtype=torch.float32),
out_loc_override=out_loc,
)
return notifications, fused_calls
def test_nsa_indexer_fused_store_does_not_notify_cp_hicache_layer_backup(self):
notifications, fused_calls = self._call_store_index_fast_path(
torch.tensor([1, 2, 3], dtype=torch.int64)
)
self.assertEqual(notifications, [])
self.assertEqual(len(fused_calls), 1)
self.assertEqual(fused_calls[0][2].tolist(), [1, 2, 3])
def test_nsa_indexer_empty_store_does_not_notify_cp_hicache_layer_backup(self):
notifications, fused_calls = self._call_store_index_fast_path(
torch.empty((0,), dtype=torch.int64)
)
self.assertEqual(notifications, [])
self.assertEqual(fused_calls, [])
def test_cp_shared_zero_local_index_store_does_not_notify_layer_backup(self):
from sglang.srt.layers.attention.nsa import nsa_indexer
notifications = []
class FakePool:
page_size = 64
def notify_layer_kv_stored_for_backup(self, layer_id, source="kv"):
notifications.append((layer_id, source))
forward_batch = types.SimpleNamespace(token_to_kv_pool=FakePool())
indexer = types.SimpleNamespace(nsa_enable_prefill_cp=True)
with (
patch.object(nsa_indexer, "nsa_use_prefill_cp", return_value=True),
patch.object(
nsa_indexer,
"get_cp_shared_kv_local_out_cache_loc",
return_value=torch.empty((0,), dtype=torch.int64),
),
patch.object(
nsa_indexer,
"get_cp_shared_kv_local_physical_out_cache_loc",
side_effect=AssertionError("must not require physical locs"),
),
):
handled = nsa_indexer.Indexer._store_cp_shared_local_index_k_cache(
indexer,
forward_batch=forward_batch,
layer_id=6,
local_key=torch.empty((0, 128), dtype=torch.float32),
act_quant=None,
)
self.assertTrue(handled)
self.assertEqual(notifications, [])
if __name__ == "__main__":
unittest.main()
Reference in New Issue View Git Blame Copy Permalink