chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 16 Packages Projects Releases Wiki Activity

Fix CP HiCache load_cp owner-pattern mismatch (cache-hit corruption)

Browse Source 
In CP=8 + NSA-shared-KV + HiCache disagg-prefill, cache-hit prefill produced
incoherent decode output. Cold prefill on CP was correct; pure CP without
HiCache was correct. The bug lived at the HiCache load_cp / device-alloc
interface.

Root cause: cache_controller.load_cp called the plain
mem_pool_device_allocator.alloc(logical_len), which returns logical pages
with no CP owner-pattern preservation. Cold prefill instead uses
alloc_extend_compute_owner with a zigzag owner pattern from
build_in_seq_page_compute_owners. The saved CpHiCacheNodeMetadata.owned_positions
records WHICH POSITIONS in the write-time alloc were owned by this rank. At
load time, those same positions are applied to a new alloc whose per-position
owner pattern is arbitrary -- each rank loads its host bytes into physical
slots whose corresponding logical page is owned by a DIFFERENT rank.
Attention's materialize_shared_token_kv_buffer reads from the owner's
physical slot, which was never loaded. Result: garbage.

Fix:
- CpHiCacheNodeMetadata gains two required fields: page_owners (int8 per
  logical page, identical on all CP ranks) and page_size. __post_init__
  validates; split() bisects page_owners by page index with a page-alignment
  check.
- _write_cp derives page_owners from device_indices (page-first slot of each
  page -> logical page id -> layout.owner_for_logical_pages) and stores in
  both metadata-construction sites (zero-owned and normal).
- New CPSharedPagedTokenToKVPoolAllocator.alloc_pages_with_owners() reuses
  _select_compute_owner_pages (with its tai-kernel fast path) and returns
  page-contiguous token locs whose per-page owner sequence equals the input.
- load_cp now concats page_owners across nodes_to_load and calls
  alloc_pages_with_owners. On None (lane exhausted) the caller hits the
  retry-with-eviction path; further failure returns None and degrades to
  cache miss. No silent fallback to plain alloc -- that recreated the bug.
- load_back retry path now calls _evict_for_compute_owner_lanes (module-top
  import) instead of plain evict(); this targets the deficit lane and gives
  the next alloc attempt a chance to satisfy it.
- envs import moved to module top in cache_controller.py per code-review
  feedback. Removed an over-defensive owned_check.all().item() in load_cp
  that would have re-introduced the host-sync anti-pattern 97a9f850c
  removed -- the invariant is already guaranteed by alloc_pages_with_owners.

Tests: 40 existing CpHiCacheNodeMetadata constructions migrated to pass the
new required fields. 9 new metadata tests (validators + split page-alignment).
10 new allocator tests in test_alloc_pages_with_owners.py covering input-order
preservation, lane exhaustion, release_pages fallback, debug-mode invariant.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
leavelet
2026-05-27 00:01:39 +08:00
co-authored by Claude Opus 4.7
parent f2834b3403
commit 1d630def95
6 changed files with 531 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/unit/managers/test_hicache_controller_cp.py
+10
View File
@@ -425,6 +425,8 @@ class TestHiCacheControllerCPLoad(TestHiCacheControllerCPWrite):
logical_len=16,
owned_positions=torch.tensor([4, 5, 6, 7], dtype=torch.int64),
host_indices=torch.tensor([100, 101, 102, 103], dtype=torch.int64),
page_owners=torch.zeros(max(16, 0), dtype=torch.int8),
page_size=1,
)
device_indices = controller.load_cp([node], node_id=11)
@@ -455,6 +457,8 @@ class TestHiCacheControllerCPLoad(TestHiCacheControllerCPWrite):
owned_positions=torch.tensor([4, 5, 6, 7], dtype=torch.int64),
host_indices=torch.tensor([100, 101, 102, 103], dtype=torch.int64),
draft_host_indices=torch.tensor([200, 201, 202, 203], dtype=torch.int64),
page_owners=torch.zeros(max(16, 0), dtype=torch.int8),
page_size=1,
)
device_indices = controller.load_cp([node], node_id=14)
@@ -477,6 +481,8 @@ class TestHiCacheControllerCPLoad(TestHiCacheControllerCPWrite):
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,
)
device_indices = controller.load_cp([node], node_id=12)
@@ -524,6 +530,8 @@ class TestHiCacheControllerCPLoad(TestHiCacheControllerCPWrite):
logical_len=16,
owned_positions=torch.tensor([4, 5, 6, 7], dtype=torch.int64),
host_indices=torch.tensor([100, 101, 102, 103], dtype=torch.int64),
page_owners=torch.zeros(max(16, 0), dtype=torch.int8),
page_size=1,
)
with self.assertRaisesRegex(
@@ -541,6 +549,8 @@ class TestHiCacheControllerCPLoad(TestHiCacheControllerCPWrite):
logical_len=16,
owned_positions=torch.tensor([4, 5, 6, 7], dtype=torch.int64),
host_indices=torch.tensor([100, 101, 103, 102], dtype=torch.int64),
page_owners=torch.zeros(max(16, 0), dtype=torch.int8),
page_size=1,
)
with self.assertRaisesRegex(ValueError, "host_indices.*contiguous page spans"):
test/registered/unit/mem_cache/test_alloc_pages_with_owners.py
+169
View File
@@ -0,0 +1,169 @@
"""Unit tests for CPSharedPagedTokenToKVPoolAllocator.alloc_pages_with_owners.
This is the load-side allocator API that HiCache load_cp uses to reproduce the
CP owner pattern recorded at write time. Correctness here is load-bearing: a
returned token-loc tensor whose per-page owner sequence does not match the
request causes attention to read from physical slots whose data was never
loaded → garbage decode output.
Run:
pytest test/registered/unit/mem_cache/test_alloc_pages_with_owners.py -v
"""
import sys
import types
import unittest
import torch
# Minimal sgl_kernel stubs (mirrors the pattern used in other unit tests so
# this file runs on CPU-only dev hosts without the sgl_kernel C++ extension).
if "sgl_kernel" not in sys.modules:
sys.modules["sgl_kernel"] = types.ModuleType("sgl_kernel")
sys.modules["sgl_kernel"].__path__ = []
if "sgl_kernel.kvcacheio" not in sys.modules:
sys.modules["sgl_kernel.kvcacheio"] = types.ModuleType("sgl_kernel.kvcacheio")
from sglang.srt.mem_cache.allocator import CPSharedPagedTokenToKVPoolAllocator
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=1, suite="stage-a-test-cpu")
def _make_allocator(
*,
page_size: int = 4,
cp_size: int = 2,
cp_rank: int = 0,
num_pages: int = 8,
device: str = "cpu",
) -> CPSharedPagedTokenToKVPoolAllocator:
"""Construct a CPSharedPagedTokenToKVPoolAllocator without exercising
KVCache. Mirrors the __new__ + manual-attrs pattern used in the existing
HiRadixCache unit tests."""
alloc = CPSharedPagedTokenToKVPoolAllocator.__new__(
CPSharedPagedTokenToKVPoolAllocator
)
alloc.page_size = page_size
alloc.cp_size = cp_size
alloc.cp_rank = cp_rank
alloc.device = device
# Logical pages are 1..num_pages (page 0 is the dummy/padding page; see
# cp_shared_kv_layout.py:11).
alloc.free_pages = torch.arange(1, num_pages + 1, dtype=torch.int64, device=device)
alloc.release_pages = torch.empty((0,), dtype=torch.int64, device=device)
alloc.debug_mode = True
return alloc
def _owners_of(pages: torch.Tensor, cp_size: int) -> torch.Tensor:
return torch.remainder(pages - 1, cp_size)
class TestAllocPagesWithOwnersHappyPath(CustomTestCase):
def test_empty_input_returns_empty_tensor(self):
alloc = _make_allocator()
before_free = alloc.free_pages.clone()
out = alloc.alloc_pages_with_owners([])
self.assertEqual(out.numel(), 0)
# Free pool unchanged.
self.assertTrue(torch.equal(alloc.free_pages, before_free))
def test_basic_owner_pattern_matches_request(self):
# cp_size=2 → lanes: rank 0 owns pages 1,3,5,7; rank 1 owns 2,4,6,8.
alloc = _make_allocator(cp_size=2, num_pages=8, page_size=4)
out = alloc.alloc_pages_with_owners([0, 1, 0, 1])
self.assertEqual(out.numel(), 4 * 4) # 4 pages × page_size
# Per-page owners reconstructed from output token locs.
page_starts = out[::4]
pages = torch.div(page_starts, 4, rounding_mode="floor")
owners = _owners_of(pages, cp_size=2)
self.assertEqual(owners.tolist(), [0, 1, 0, 1])
def test_preserves_input_order_not_lane_grouped(self):
# Request a NON-grouped owner pattern. If the allocator silently
# groups by lane, pages 0/1/0/1 would yield [r0, r0, r1, r1] —
# the order test catches that.
alloc = _make_allocator(cp_size=2, num_pages=8, page_size=4)
seq = [1, 0, 1, 0, 1, 0]
out = alloc.alloc_pages_with_owners(seq)
pages = torch.div(out[::4], 4, rounding_mode="floor")
owners = _owners_of(pages, cp_size=2)
self.assertEqual(owners.tolist(), seq)
def test_zero_offset_within_page_at_first_token_of_each_page(self):
alloc = _make_allocator(cp_size=2, num_pages=8, page_size=4)
out = alloc.alloc_pages_with_owners([0, 1])
# Token 0 = page_start of page 0; token 4 = page_start of page 1.
self.assertEqual(int(out[0] % 4), 0)
self.assertEqual(int(out[4] % 4), 0)
# Within each page, offsets 0..3.
self.assertEqual(out[0:4].tolist(), [out[0].item() + i for i in range(4)])
def test_consumes_free_pages_in_lane_order(self):
alloc = _make_allocator(cp_size=2, num_pages=8, page_size=4)
out = alloc.alloc_pages_with_owners([0, 1])
# Expect page 1 (rank-0 first) and page 2 (rank-1 first) consumed.
consumed_pages = torch.div(out[::4], 4, rounding_mode="floor").tolist()
self.assertEqual(sorted(consumed_pages), [1, 2])
# Remaining free_pages should be 3..8 (any order).
self.assertEqual(sorted(alloc.free_pages.tolist()), [3, 4, 5, 6, 7, 8])
def test_uses_release_pages_when_free_pages_short(self):
# Push everything to release_pages so free_pages can't satisfy the
# request from the front. _select_compute_owner_pages must consult
# both lists.
alloc = _make_allocator(cp_size=2, num_pages=4, page_size=4)
alloc.release_pages = alloc.free_pages.clone()
alloc.free_pages = torch.empty((0,), dtype=torch.int64)
out = alloc.alloc_pages_with_owners([0, 1])
self.assertEqual(out.numel(), 2 * 4)
# release_pages should have shrunk.
self.assertEqual(alloc.release_pages.numel(), 2)
class TestAllocPagesWithOwnersExhaustion(CustomTestCase):
def test_returns_none_when_lane_exhausted(self):
# Only 2 pages, both rank-0 (pages 1, 3). Asking for 2 rank-1 pages
# must fail.
alloc = _make_allocator(cp_size=2, num_pages=3, page_size=4)
# Manually narrow the lanes: free_pages = [1, 3] are rank 0.
alloc.free_pages = torch.tensor([1, 3], dtype=torch.int64)
alloc.release_pages = torch.empty((0,), dtype=torch.int64)
out = alloc.alloc_pages_with_owners([1, 1])
self.assertIsNone(out)
# Free pages must be UNCHANGED on None return.
self.assertEqual(alloc.free_pages.tolist(), [1, 3])
def test_partial_exhaustion_returns_none_no_leftover(self):
# Lane 0 has 1 page (page 1), lane 1 has 1 (page 2). Request [0, 0]
# → only 1 lane-0 page available → None, no state change.
alloc = _make_allocator(cp_size=2, num_pages=2, page_size=4)
out = alloc.alloc_pages_with_owners([0, 0])
self.assertIsNone(out)
self.assertEqual(alloc.free_pages.tolist(), [1, 2])
self.assertEqual(alloc.release_pages.numel(), 0)
class TestAllocPagesWithOwnersDebugMode(CustomTestCase):
def test_debug_mode_invariant_holds(self):
# debug_mode=True (default in _make_allocator) — the internal assertion
# validates selected owners match input. Must NOT raise on a correct
# allocation.
alloc = _make_allocator(cp_size=4, num_pages=16, page_size=4)
out = alloc.alloc_pages_with_owners([0, 1, 2, 3, 0, 1, 2, 3])
self.assertEqual(out.numel(), 8 * 4)
owners = _owners_of(
torch.div(out[::4], 4, rounding_mode="floor"), cp_size=4
)
self.assertEqual(owners.tolist(), [0, 1, 2, 3, 0, 1, 2, 3])
def test_returned_locs_are_unique(self):
alloc = _make_allocator(cp_size=2, num_pages=8, page_size=4)
out = alloc.alloc_pages_with_owners([0, 1, 0, 1])
self.assertEqual(torch.unique(out).numel(), out.numel())
if __name__ == "__main__":
unittest.main()
test/registered/unit/mem_cache/test_cp_hicache_metadata.py
+175
View File
@@ -109,6 +109,8 @@ class TestCpHiCacheNodeMetadata(CustomTestCase):
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)
@@ -125,6 +127,8 @@ class TestCpHiCacheNodeMetadata(CustomTestCase):
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)
@@ -142,6 +146,8 @@ class TestCpHiCacheNodeMetadata(CustomTestCase):
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)
@@ -157,6 +163,8 @@ class TestCpHiCacheNodeMetadata(CustomTestCase):
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)
@@ -170,6 +178,8 @@ class TestCpHiCacheNodeMetadata(CustomTestCase):
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)
@@ -181,6 +191,8 @@ class TestCpHiCacheNodeMetadata(CustomTestCase):
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):
@@ -191,6 +203,8 @@ class TestCpHiCacheNodeMetadata(CustomTestCase):
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
@@ -203,6 +217,8 @@ class TestCpHiCacheNodeMetadata(CustomTestCase):
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"):
@@ -214,6 +230,8 @@ class TestCpHiCacheNodeMetadata(CustomTestCase):
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):
@@ -222,6 +240,8 @@ class TestCpHiCacheNodeMetadata(CustomTestCase):
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):
@@ -230,6 +250,8 @@ class TestCpHiCacheNodeMetadata(CustomTestCase):
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):
@@ -239,6 +261,8 @@ class TestCpHiCacheNodeMetadata(CustomTestCase):
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):
@@ -247,8 +271,129 @@ class TestCpHiCacheNodeMetadata(CustomTestCase):
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
@@ -280,6 +425,8 @@ class FakeWriteController:
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,
)
)
@@ -297,6 +444,8 @@ class FakeZeroOwnedWriteController:
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,
)
)
@@ -328,6 +477,8 @@ class TestHiRadixCacheCPBackup(CustomTestCase):
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))
@@ -348,6 +499,8 @@ class TestHiRadixCacheCPBackup(CustomTestCase):
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)
@@ -378,6 +531,8 @@ class TestHiRadixCacheCPBackup(CustomTestCase):
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] = evictable_node
cache.evictable_host_leaves.add(evictable_node)
@@ -457,6 +612,8 @@ class TestHiRadixCacheCPBackup(CustomTestCase):
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)
@@ -488,6 +645,8 @@ class TestHiRadixCacheCPSplitEvict(CustomTestCase):
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
@@ -534,6 +693,8 @@ class TestHiRadixCacheCPSplitEvict(CustomTestCase):
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)
@@ -576,6 +737,8 @@ class TestHiRadixCacheCPSplitEvict(CustomTestCase):
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
@@ -630,6 +793,8 @@ class TestHiRadixCacheCPSplitEvict(CustomTestCase):
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
@@ -696,6 +861,8 @@ class TestHiRadixCacheCPSplitEvict(CustomTestCase):
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)
@@ -751,6 +918,8 @@ class TestHiRadixCacheCPSplitEvict(CustomTestCase):
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)
@@ -804,6 +973,8 @@ class TestHiRadixCacheCPLoadBack(CustomTestCase):
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)
@@ -839,6 +1010,8 @@ class TestHiRadixCacheCPLoadBack(CustomTestCase):
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)
@@ -871,6 +1044,8 @@ class TestHiRadixCacheCPLoadBack(CustomTestCase):
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