Fix CP HiCache load_cp owner-pattern mismatch (cache-hit corruption)
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:
@@ -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)
|
||||
)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
)
|
||||
|
||||
@@ -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"):
|
||||
|
||||
169
test/registered/unit/mem_cache/test_alloc_pages_with_owners.py
Normal file
169
test/registered/unit/mem_cache/test_alloc_pages_with_owners.py
Normal 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()
|
||||
@@ -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
|
||||
|
||||
|
||||
Reference in New Issue
Block a user