chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity
Files
25f21476773fa89e1eec0770d4dec7c0af2df176
sglang/test/registered/unit/mem_cache/test_cp_hicache_metadata.py
laoyao0822 25f2147677 Reduce CP HiCache capacity synchronization to owner-lane logic 
CP shared KV and HiCache now use owner-lane metadata as the
authoritative capacity view for host write admission and GPU load-back
planning. This removes the debug scalar capacity env and keeps CP load-back
from relying on a rank-wide scalar collective when per-owner availability is
already known. The load-back planner also accounts for evicting child leaves
that unlock ancestor device residency, which fixes small lane deficits despite
large aggregate evictable capacity.

The commit also adds gated CPU timing logs for CP shared-KV MLA/index
prefetch and a CUDA microbenchmark for comparing dense all-reduce with
owner-packed all-gather layouts. The timing logs are intentionally behind the
existing MLA prefetch log env and should not be enabled for throughput
measurements.

Constraint: CP shared KV owner lanes require target/draft capacity decisions to preserve page_owners rather than total-token scalars
Constraint: CUDA collective benchmarks must run on target GPU hosts, not locally
Rejected: Keep SGLANG_CP_HICACHE_CAPACITY_DEBUG observer env | owner-lane admission now replaces that scalar debug path
Rejected: Add a silent scalar-allreduce fallback | unexpected owner-lane mismatch should fail fast or log loudly
Confidence: medium
Scope-risk: moderate
Directive: Do not reintroduce CP capacity collectives on the scheduler hot path without proving the owner-lane metadata is insufficient
Directive: Disable SGLANG_CP_SHARED_KV_LOG_MLA_PREFETCH for end-to-end performance runs; it is diagnostic and high-volume
Tested: git diff --check
Tested: python -m py_compile on changed runtime/test/benchmark Python files
Tested: remote pytest -q test/registered/unit/mem_cache/test_cp_hicache_load_back_owner_lanes.py test/registered/unit/mem_cache/test_cp_hicache_metadata.py (81 passed, 5 warnings)
Not-tested: CUDA benchmark benchmark/hicache/bench_cp_shared_kv_prefetch_collective.py
Not-tested: full GLM5 E2E throughput after this commit
2026-05-28 08:31:49 +08:00

2294 lines
88 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
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_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_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_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))
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_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