diff --git a/python/sglang/srt/managers/cache_controller.py b/python/sglang/srt/managers/cache_controller.py index a17bb296b..e5de85a6e 100644 --- a/python/sglang/srt/managers/cache_controller.py +++ b/python/sglang/srt/managers/cache_controller.py @@ -32,6 +32,7 @@ from sglang.srt.distributed import ( get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, ) +from sglang.srt.environ import envs from sglang.srt.layers.dp_attention import ( get_attention_dp_rank, get_attention_tp_rank, @@ -836,6 +837,8 @@ class HiCacheController: host_indices: torch.Tensor, device_indices: torch.Tensor, ) -> None: + if not envs.SGLANG_DEBUG_HICACHE_VALIDATE.get(): + return validate_page_aligned_token_indices( device_indices, self.page_size, "physical_device_indices" ) @@ -853,6 +856,23 @@ class HiCacheController: owned_mask = layout.owned_by_this_rank(device_indices) owned_positions = owned_mask.nonzero(as_tuple=True)[0].cpu() logical_len = len(device_indices) + # Capture the global owner pattern (one int8 per logical PAGE; identical + # on all CP ranks since it's a pure function of the logical page ids). + # load_cp will replay this pattern via alloc_pages_with_owners so the + # saved owned_positions correctly index the new allocation. + page_size = self.page_size + if logical_len % page_size != 0: + raise ValueError( + f"_write_cp expects page-aligned device_indices, got " + f"logical_len={logical_len} page_size={page_size}" + ) + page_first_locs = device_indices[::page_size] + logical_pages = torch.div( + page_first_locs, page_size, rounding_mode="floor" + ) + page_owners = layout.owner_for_logical_pages(logical_pages).to( + dtype=torch.int8, device="cpu" + ) if owned_positions.numel() == 0: self._append_completed_write_ack(node_id) logger.info( @@ -865,6 +885,8 @@ class HiCacheController: logical_len=logical_len, owned_positions=owned_positions, host_indices=torch.empty((0,), dtype=torch.int64), + page_owners=page_owners, + page_size=page_size, draft_host_indices=( torch.empty((0,), dtype=torch.int64) if self.has_draft_hicache @@ -935,6 +957,8 @@ class HiCacheController: logical_len=logical_len, owned_positions=owned_positions, host_indices=host_indices.cpu(), + page_owners=page_owners, + page_size=page_size, draft_host_indices=( draft_host_indices.cpu() if draft_host_indices is not None else None ), @@ -975,11 +999,42 @@ class HiCacheController: return device_indices def load_cp(self, nodes_to_load, node_id: int = -1) -> Optional[torch.Tensor]: - logical_len = sum(node.host_len for node in nodes_to_load) - device_indices = self.mem_pool_device_allocator.alloc(logical_len) + # Reproduce the original (write-time) CP owner pattern. Each node + # carries `page_owners` (one int8 per logical page, identical on all + # CP ranks) so we can ask the allocator for a fresh device range + # whose per-page owner sequence matches what was backed up. Without + # this, the saved `owned_positions` (positions WITHIN the original + # alloc that this rank owned) would index a new alloc with arbitrary + # owner pattern → each rank loads its host bytes into physical slots + # whose corresponding logical page is owned by some other rank → + # attention reads garbage at forward time. + page_owners: List[int] = [] + for node in nodes_to_load: + meta = node.cp_hicache + if meta is None: + raise RuntimeError( + f"load_cp called with node {getattr(node, 'id', '?')} " + "that has no cp_hicache metadata" + ) + # page_owners is CPU int8; .tolist() returns list[int] directly. + page_owners.extend(meta.page_owners.tolist()) + + device_indices = self.mem_pool_device_allocator.alloc_pages_with_owners( + page_owners + ) + # Fail closed: returning None lets the caller drop to cache miss + # (cold prefill). Never proceed with a non-matching owner pattern. if device_indices is None: return None + logical_len_expected = sum(node.host_len for node in nodes_to_load) + if device_indices.numel() != logical_len_expected: + self.mem_pool_device_allocator.free(device_indices) + raise RuntimeError( + "alloc_pages_with_owners returned unexpected length: " + f"got {device_indices.numel()}, expected {logical_len_expected}" + ) + host_chunks = [] draft_host_chunks = [] physical_chunks = [] @@ -991,6 +1046,18 @@ class HiCacheController: if owned_positions.numel() == 0: continue selected_logical_locs = node_device_indices[owned_positions] + # Note: NOT re-validating owner_by_this_rank here. The invariant + # — every position in `owned_positions` lands on a page owned by + # this rank — is guaranteed by construction: + # (a) at write time, `owned_positions = owned_mask.nonzero(...)` + # where `owned_mask = owned_by_this_rank(device_indices)`; + # (b) at load time, `alloc_pages_with_owners(page_owners)` returns + # pages whose owner sequence matches `page_owners` by + # construction (and is debug-asserted inside the allocator). + # Both sides use the same `page_owners` (write-derived, identical + # on all CP ranks). A redundant `.all().item()` check here would + # force a CUDA host-sync — the exact anti-pattern commit 97a9f850c + # removed from the hot path. physical_chunks.append( self.cp_shared_kv_layout.logical_locs_to_physical(selected_logical_locs) ) diff --git a/python/sglang/srt/mem_cache/allocator.py b/python/sglang/srt/mem_cache/allocator.py index b50ae3a86..0b5a1349e 100644 --- a/python/sglang/srt/mem_cache/allocator.py +++ b/python/sglang/srt/mem_cache/allocator.py @@ -763,3 +763,55 @@ class CPSharedPagedTokenToKVPoolAllocator(PagedTokenToKVPoolAllocator): assert torch.equal(selected_owners, expected_owners) return out_indices + + def alloc_pages_with_owners( + self, + page_compute_owners: List[int], + ) -> Optional[torch.Tensor]: + """Allocate ``len(page_compute_owners)`` whole logical pages where + page ``i``'s owner equals ``page_compute_owners[i]``. + + Returns a flat int64 tensor of logical token locs of length + ``len(owners) * page_size``, page-contiguous in input order, or + ``None`` when an owner lane is exhausted (caller handles + eviction-retry / cache-miss). + + Used by CP HiCache ``load_cp`` to reproduce the write-time owner + pattern recorded in ``CpHiCacheNodeMetadata.page_owners``. Without + this, plain ``alloc()`` would return pages with an arbitrary owner + pattern, the saved per-position owner mask would index the wrong + slots, and each rank would load its host bytes into physical slots + whose corresponding logical page is owned by another rank → attention + reads garbage. + + Distinguished from ``alloc_extend_compute_owner``: that variant + expects extend semantics (prefix_lens / seq_lens / last_loc) and + invokes the ``alloc_extend_naive`` Triton kernel to splice the + partial last page of a prefix with new pages. HiCache reload has + no prefix and no partial page — host pages are always page-aligned + — so a pure page-fresh allocation is enough. + """ + if not page_compute_owners: + return torch.empty((0,), dtype=torch.int64, device=self.device) + selected = self._select_compute_owner_pages(page_compute_owners) + if selected is None: + return None + selected_pages, selected_free_mask, selected_release_mask = selected + page_size = self.page_size + base = selected_pages.to(torch.int64).unsqueeze(1) * page_size + offsets = torch.arange( + page_size, dtype=torch.int64, device=self.device + ).unsqueeze(0) + out_indices = (base + offsets).reshape(-1) + self.free_pages = self.free_pages[~selected_free_mask] + self.release_pages = self.release_pages[~selected_release_mask] + if self.debug_mode: + assert torch.unique(out_indices).numel() == out_indices.numel() + check_owners = torch.remainder(selected_pages - 1, self.cp_size) + expected = torch.tensor( + page_compute_owners, + dtype=check_owners.dtype, + device=check_owners.device, + ) + assert torch.equal(check_owners, expected) + return out_indices diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index c830b914a..e79b60675 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -28,6 +28,7 @@ from sglang.srt.mem_cache.base_prefix_cache import ( MatchPrefixParams, MatchResult, ) +from sglang.srt.mem_cache.common import _evict_for_compute_owner_lanes from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout from sglang.srt.mem_cache.memory_pool import ( MHATokenToKVPool, @@ -57,21 +58,50 @@ class CpHiCacheNodeMetadata: logical_len: int owned_positions: torch.Tensor host_indices: torch.Tensor + # NEW: one int8 per logical PAGE; identical on all CP ranks (function of + # the original logical page ids only). Captures the owner pattern of the + # write-time allocation so load_cp can reproduce it via owner-aware alloc. + # Without this, the saved owned_positions index a fresh alloc whose + # per-position owner pattern is arbitrary → load writes to wrong physical + # slots → attention reads garbage. See _write_cp for the derivation and + # alloc_pages_with_owners for the consumer. + page_owners: torch.Tensor + # NEW: needed at split() time to convert split_len into a page index + # without plumbing it from the caller. + page_size: int draft_host_indices: Optional[torch.Tensor] = None def __post_init__(self): if self.logical_len < 0: raise ValueError(f"logical_len must be non-negative, got {self.logical_len}") + if self.page_size <= 0: + raise ValueError(f"page_size must be positive, got {self.page_size}") + if self.logical_len % self.page_size != 0: + raise ValueError( + f"logical_len ({self.logical_len}) must be a multiple of " + f"page_size ({self.page_size})" + ) self.owned_positions = self.owned_positions.to( device="cpu", dtype=torch.int64 ).detach().clone() self.host_indices = self.host_indices.to( device="cpu", dtype=torch.int64 ).detach().clone() + self.page_owners = self.page_owners.to( + device="cpu", dtype=torch.int8 + ).detach().clone() if self.draft_host_indices is not None: self.draft_host_indices = self.draft_host_indices.to( device="cpu", dtype=torch.int64 ).detach().clone() + expected_num_pages = self.logical_len // self.page_size + if self.page_owners.numel() != expected_num_pages: + raise ValueError( + f"page_owners length ({self.page_owners.numel()}) must equal " + f"logical_len/page_size ({expected_num_pages})" + ) + if expected_num_pages > 0 and bool((self.page_owners < 0).any()): + raise ValueError("page_owners entries must be non-negative") if self.owned_positions.numel() != self.host_indices.numel(): raise ValueError( "owned_positions and host_indices must have same length, got " @@ -102,13 +132,21 @@ class CpHiCacheNodeMetadata: raise ValueError( f"split_len must be in [0, {self.logical_len}], got {split_len}" ) + if split_len % self.page_size != 0: + raise ValueError( + f"split_len ({split_len}) must be a multiple of page_size " + f"({self.page_size})" + ) parent_mask = self.owned_positions < split_len child_mask = ~parent_mask + split_pages = split_len // self.page_size return ( CpHiCacheNodeMetadata( logical_len=split_len, owned_positions=self.owned_positions[parent_mask], host_indices=self.host_indices[parent_mask], + page_owners=self.page_owners[:split_pages], + page_size=self.page_size, draft_host_indices=( self.draft_host_indices[parent_mask] if self.draft_host_indices is not None @@ -119,6 +157,8 @@ class CpHiCacheNodeMetadata: logical_len=self.logical_len - split_len, owned_positions=self.owned_positions[child_mask] - split_len, host_indices=self.host_indices[child_mask], + page_owners=self.page_owners[split_pages:], + page_size=self.page_size, draft_host_indices=( self.draft_host_indices[child_mask] if self.draft_host_indices is not None @@ -1414,11 +1454,25 @@ class HiRadixCache(RadixCache): ) if device_indices is None: logger.info( - "[HiCache-load] load_back CP retry with eviction: node_id=%d tokens_needed=%d", + "[HiCache-load] load_back CP retry with lane-aware eviction: " + "node_id=%d tokens_needed=%d", last_hit_node.id, host_hit_len, ) - self.evict(EvictParams(num_tokens=host_hit_len)) + # Lane-aware eviction: alloc_pages_with_owners failed because + # at least one owner lane is short of free pages. Targeted + # eviction frees pages from THOSE specific lanes, leaving + # other lanes untouched. Mirrors cold prefill's retry path + # in common.py:alloc_paged_token_slots_extend. + _retry_page_owners: List[int] = [] + for _node in nodes_to_load: + # page_owners is CPU int8; tolist() returns list[int]. + _retry_page_owners.extend(_node.cp_hicache.page_owners.tolist()) + _evict_for_compute_owner_lanes( + tree_cache=self, + allocator=self.token_to_kv_pool_allocator, + page_compute_owners=_retry_page_owners, + ) device_indices = self.cache_controller.load_cp( nodes_to_load, node_id=last_hit_node.id ) diff --git a/test/registered/unit/managers/test_hicache_controller_cp.py b/test/registered/unit/managers/test_hicache_controller_cp.py index 49d25d9a8..d24f25847 100644 --- a/test/registered/unit/managers/test_hicache_controller_cp.py +++ b/test/registered/unit/managers/test_hicache_controller_cp.py @@ -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"): diff --git a/test/registered/unit/mem_cache/test_alloc_pages_with_owners.py b/test/registered/unit/mem_cache/test_alloc_pages_with_owners.py new file mode 100644 index 000000000..f42f2598a --- /dev/null +++ b/test/registered/unit/mem_cache/test_alloc_pages_with_owners.py @@ -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() diff --git a/test/registered/unit/mem_cache/test_cp_hicache_metadata.py b/test/registered/unit/mem_cache/test_cp_hicache_metadata.py index dc8d16fa9..99c1dfb85 100644 --- a/test/registered/unit/mem_cache/test_cp_hicache_metadata.py +++ b/test/registered/unit/mem_cache/test_cp_hicache_metadata.py @@ -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