[UnifiedRadixCache] Fix HiCache load back start node (#25088)
This commit is contained in:
@@ -740,8 +740,10 @@ class Req(ReqDllmMixin):
|
|||||||
self.extend_input_len = 0
|
self.extend_input_len = 0
|
||||||
# The relative logprob_start_len in an extend batch
|
# The relative logprob_start_len in an extend batch
|
||||||
self.extend_logprob_start_len = 0
|
self.extend_logprob_start_len = 0
|
||||||
|
# TODO(ispobock): rename to last_device_node
|
||||||
self.last_node: Any = None
|
self.last_node: Any = None
|
||||||
self.last_host_node: Any = None
|
self.last_host_node: Any = None
|
||||||
|
self.best_match_node: Any = None
|
||||||
self.host_hit_length = 0
|
self.host_hit_length = 0
|
||||||
# Tokens loaded from storage backend (L3) during prefetch for this request
|
# Tokens loaded from storage backend (L3) during prefetch for this request
|
||||||
self.storage_hit_length = 0
|
self.storage_hit_length = 0
|
||||||
@@ -1041,12 +1043,14 @@ class Req(ReqDllmMixin):
|
|||||||
self.prefix_indices,
|
self.prefix_indices,
|
||||||
self.last_node,
|
self.last_node,
|
||||||
self.last_host_node,
|
self.last_host_node,
|
||||||
|
self.best_match_node,
|
||||||
self.host_hit_length,
|
self.host_hit_length,
|
||||||
self.mamba_branching_seqlen,
|
self.mamba_branching_seqlen,
|
||||||
) = (
|
) = (
|
||||||
match_result.device_indices,
|
match_result.device_indices,
|
||||||
match_result.last_device_node,
|
match_result.last_device_node,
|
||||||
match_result.last_host_node,
|
match_result.last_host_node,
|
||||||
|
match_result.best_match_node,
|
||||||
match_result.host_hit_length,
|
match_result.host_hit_length,
|
||||||
match_result.mamba_branching_seqlen,
|
match_result.mamba_branching_seqlen,
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -105,11 +105,13 @@ def match_prefix_for_req(
|
|||||||
req.prefix_indices,
|
req.prefix_indices,
|
||||||
req.last_node,
|
req.last_node,
|
||||||
req.last_host_node,
|
req.last_host_node,
|
||||||
|
req.best_match_node,
|
||||||
req.host_hit_length,
|
req.host_hit_length,
|
||||||
) = (
|
) = (
|
||||||
match_result.device_indices,
|
match_result.device_indices,
|
||||||
match_result.last_device_node,
|
match_result.last_device_node,
|
||||||
match_result.last_host_node,
|
match_result.last_host_node,
|
||||||
|
match_result.best_match_node,
|
||||||
match_result.host_hit_length,
|
match_result.host_hit_length,
|
||||||
)
|
)
|
||||||
if match_result.mamba_branching_seqlen is not None:
|
if match_result.mamba_branching_seqlen is not None:
|
||||||
@@ -877,7 +879,7 @@ class PrefillAdder:
|
|||||||
if req.host_hit_length > 0:
|
if req.host_hit_length > 0:
|
||||||
new_indices, req.last_node = self.tree_cache.init_load_back(
|
new_indices, req.last_node = self.tree_cache.init_load_back(
|
||||||
InitLoadBackParams(
|
InitLoadBackParams(
|
||||||
last_host_node=req.last_host_node,
|
best_match_node=req.best_match_node,
|
||||||
host_hit_length=req.host_hit_length,
|
host_hit_length=req.host_hit_length,
|
||||||
req=req,
|
req=req,
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -134,9 +134,9 @@ class DecLockRefResult:
|
|||||||
|
|
||||||
@dataclasses.dataclass
|
@dataclasses.dataclass
|
||||||
class InitLoadBackParams:
|
class InitLoadBackParams:
|
||||||
"""Unified parameters for init_load_back across different cache types"""
|
"""Unified parameters for init_load_back across different cache types."""
|
||||||
|
|
||||||
last_host_node: Any
|
best_match_node: Any
|
||||||
host_hit_length: int
|
host_hit_length: int
|
||||||
mem_quota: Optional[int] = None
|
mem_quota: Optional[int] = None
|
||||||
req: Optional[Req] = None
|
req: Optional[Req] = None
|
||||||
@@ -151,6 +151,13 @@ class MatchResult(NamedTuple):
|
|||||||
last_host_node : The last TreeNode on the host that was matched.
|
last_host_node : The last TreeNode on the host that was matched.
|
||||||
Note that if HiCache is not enabled,
|
Note that if HiCache is not enabled,
|
||||||
this **must** be the same as `last_device_node`.
|
this **must** be the same as `last_device_node`.
|
||||||
|
Reserved for L3 storage prefetch anchoring; L2 load_back
|
||||||
|
uses `best_match_node` instead.
|
||||||
|
best_match_node : Deepest node accepted by all component validators
|
||||||
|
during match_prefix. Anchor for every L2 host->device
|
||||||
|
load_back walk (FULL / SWA / ...). For legacy caches
|
||||||
|
that don't run multi-component validation, set this
|
||||||
|
equal to `last_host_node`.
|
||||||
host_hit_length : Length of the host cache hit. For pure-KV caches this is the
|
host_hit_length : Length of the host cache hit. For pure-KV caches this is the
|
||||||
number of evicted KV tokens on CPU. For hybrid Mamba models this
|
number of evicted KV tokens on CPU. For hybrid Mamba models this
|
||||||
is max(kv_host_tokens, 1-if-mamba-on-host) so that a mamba-only
|
is max(kv_host_tokens, 1-if-mamba-on-host) so that a mamba-only
|
||||||
@@ -164,6 +171,7 @@ class MatchResult(NamedTuple):
|
|||||||
device_indices: torch.Tensor
|
device_indices: torch.Tensor
|
||||||
last_device_node: Any
|
last_device_node: Any
|
||||||
last_host_node: Any
|
last_host_node: Any
|
||||||
|
best_match_node: Any
|
||||||
host_hit_length: int = 0
|
host_hit_length: int = 0
|
||||||
mamba_branching_seqlen: Optional[int] = None
|
mamba_branching_seqlen: Optional[int] = None
|
||||||
cache_protected_len: Optional[int] = None
|
cache_protected_len: Optional[int] = None
|
||||||
@@ -180,6 +188,7 @@ def zero_match_result(tree_cache, match_result: "MatchResult") -> "MatchResult":
|
|||||||
device_indices=match_result.device_indices[:0],
|
device_indices=match_result.device_indices[:0],
|
||||||
last_device_node=root,
|
last_device_node=root,
|
||||||
last_host_node=root,
|
last_host_node=root,
|
||||||
|
best_match_node=root,
|
||||||
host_hit_length=0,
|
host_hit_length=0,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|||||||
@@ -69,6 +69,7 @@ class ChunkCache(BasePrefixCache):
|
|||||||
device_indices=torch.empty((0,), dtype=torch.int64),
|
device_indices=torch.empty((0,), dtype=torch.int64),
|
||||||
last_device_node=None,
|
last_device_node=None,
|
||||||
last_host_node=None,
|
last_host_node=None,
|
||||||
|
best_match_node=None,
|
||||||
)
|
)
|
||||||
|
|
||||||
def insert(self, params: InsertParams) -> InsertResult:
|
def insert(self, params: InsertParams) -> InsertResult:
|
||||||
|
|||||||
@@ -345,7 +345,7 @@ class HiMambaRadixCache(MambaRadixCache):
|
|||||||
self,
|
self,
|
||||||
params: InitLoadBackParams,
|
params: InitLoadBackParams,
|
||||||
):
|
):
|
||||||
last_node = params.last_host_node
|
last_node = params.best_match_node
|
||||||
mem_quota = params.mem_quota
|
mem_quota = params.mem_quota
|
||||||
req = params.req
|
req = params.req
|
||||||
if last_node.evicted or (last_node.mamba_evicted and last_node.mamba_backuped):
|
if last_node.evicted or (last_node.mamba_evicted and last_node.mamba_backuped):
|
||||||
@@ -932,6 +932,7 @@ class HiMambaRadixCache(MambaRadixCache):
|
|||||||
device_indices=torch.empty((0,), dtype=torch.int64, device=self.device),
|
device_indices=torch.empty((0,), dtype=torch.int64, device=self.device),
|
||||||
last_device_node=self.root_node,
|
last_device_node=self.root_node,
|
||||||
last_host_node=self.root_node,
|
last_host_node=self.root_node,
|
||||||
|
best_match_node=self.root_node,
|
||||||
host_hit_length=0,
|
host_hit_length=0,
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -1066,6 +1067,8 @@ class HiMambaRadixCache(MambaRadixCache):
|
|||||||
device_indices=value,
|
device_indices=value,
|
||||||
last_device_node=last_device_node,
|
last_device_node=last_device_node,
|
||||||
last_host_node=last_host_node,
|
last_host_node=last_host_node,
|
||||||
|
# TODO(ispobock): use best_match_node as start node for load_back
|
||||||
|
best_match_node=last_host_node,
|
||||||
host_hit_length=host_hit_length,
|
host_hit_length=host_hit_length,
|
||||||
mamba_branching_seqlen=mamba_branching_seqlen,
|
mamba_branching_seqlen=mamba_branching_seqlen,
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -1046,7 +1046,7 @@ class HiRadixCache(RadixCache):
|
|||||||
self,
|
self,
|
||||||
params: InitLoadBackParams,
|
params: InitLoadBackParams,
|
||||||
):
|
):
|
||||||
last_node = params.last_host_node
|
last_node = params.best_match_node
|
||||||
mem_quota = params.mem_quota
|
mem_quota = params.mem_quota
|
||||||
if last_node.evicted:
|
if last_node.evicted:
|
||||||
loading_values = self.load_back(last_node, mem_quota)
|
loading_values = self.load_back(last_node, mem_quota)
|
||||||
@@ -1251,6 +1251,8 @@ class HiRadixCache(RadixCache):
|
|||||||
device_indices=value,
|
device_indices=value,
|
||||||
last_device_node=last_node,
|
last_device_node=last_node,
|
||||||
last_host_node=last_host_node,
|
last_host_node=last_host_node,
|
||||||
|
# TODO(ispobock): use best_match_node as start node for load_back
|
||||||
|
best_match_node=last_host_node,
|
||||||
host_hit_length=host_hit_length,
|
host_hit_length=host_hit_length,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|||||||
@@ -488,6 +488,7 @@ class MambaRadixCache(KVCacheEventMixin, BasePrefixCache):
|
|||||||
),
|
),
|
||||||
last_device_node=self.root_node,
|
last_device_node=self.root_node,
|
||||||
last_host_node=self.root_node,
|
last_host_node=self.root_node,
|
||||||
|
best_match_node=self.root_node,
|
||||||
)
|
)
|
||||||
|
|
||||||
value, last_node, best_value_len = self._match_prefix_helper(key)
|
value, last_node, best_value_len = self._match_prefix_helper(key)
|
||||||
@@ -1072,6 +1073,7 @@ class MambaRadixCache(KVCacheEventMixin, BasePrefixCache):
|
|||||||
device_indices=value,
|
device_indices=value,
|
||||||
last_device_node=last_node,
|
last_device_node=last_node,
|
||||||
last_host_node=last_node,
|
last_host_node=last_node,
|
||||||
|
best_match_node=last_node,
|
||||||
mamba_branching_seqlen=mamba_branching_seqlen,
|
mamba_branching_seqlen=mamba_branching_seqlen,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|||||||
@@ -353,6 +353,7 @@ class RadixCache(KVCacheEventMixin, BasePrefixCache):
|
|||||||
),
|
),
|
||||||
last_device_node=self.root_node,
|
last_device_node=self.root_node,
|
||||||
last_host_node=self.root_node,
|
last_host_node=self.root_node,
|
||||||
|
best_match_node=self.root_node,
|
||||||
)
|
)
|
||||||
self._record_all_cleared_event()
|
self._record_all_cleared_event()
|
||||||
|
|
||||||
@@ -413,6 +414,7 @@ class RadixCache(KVCacheEventMixin, BasePrefixCache):
|
|||||||
device_indices=value,
|
device_indices=value,
|
||||||
last_device_node=last_node,
|
last_device_node=last_node,
|
||||||
last_host_node=last_node,
|
last_host_node=last_node,
|
||||||
|
best_match_node=last_node,
|
||||||
)
|
)
|
||||||
|
|
||||||
def insert(self, params: InsertParams) -> InsertResult:
|
def insert(self, params: InsertParams) -> InsertResult:
|
||||||
|
|||||||
@@ -107,6 +107,7 @@ class RadixCacheCpp(BasePrefixCache):
|
|||||||
device_indices=self._merge_tensor(device_indices_vec),
|
device_indices=self._merge_tensor(device_indices_vec),
|
||||||
last_device_node=node_gpu,
|
last_device_node=node_gpu,
|
||||||
last_host_node=node_cpu,
|
last_host_node=node_cpu,
|
||||||
|
best_match_node=node_cpu,
|
||||||
host_hit_length=host_indices_length,
|
host_hit_length=host_indices_length,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|||||||
@@ -207,6 +207,7 @@ class LMCRadixCache(RadixCache):
|
|||||||
device_indices=value,
|
device_indices=value,
|
||||||
last_device_node=last_node,
|
last_device_node=last_node,
|
||||||
last_host_node=last_node,
|
last_host_node=last_node,
|
||||||
|
best_match_node=last_node,
|
||||||
)
|
)
|
||||||
|
|
||||||
return base_res
|
return base_res
|
||||||
|
|||||||
@@ -405,6 +405,7 @@ class SWARadixCache(KVCacheEventMixin, BasePrefixCache):
|
|||||||
),
|
),
|
||||||
last_device_node=self.root_node,
|
last_device_node=self.root_node,
|
||||||
last_host_node=self.root_node,
|
last_host_node=self.root_node,
|
||||||
|
best_match_node=self.root_node,
|
||||||
)
|
)
|
||||||
|
|
||||||
value, last_node, best_value_len = self._match_prefix_helper(key)
|
value, last_node, best_value_len = self._match_prefix_helper(key)
|
||||||
@@ -960,6 +961,7 @@ class SWARadixCache(KVCacheEventMixin, BasePrefixCache):
|
|||||||
device_indices=value,
|
device_indices=value,
|
||||||
last_device_node=last_node,
|
last_device_node=last_node,
|
||||||
last_host_node=last_node,
|
last_host_node=last_node,
|
||||||
|
best_match_node=last_node,
|
||||||
)
|
)
|
||||||
|
|
||||||
def _compact_single_child_chain(self, node: TreeNode) -> None:
|
def _compact_single_child_chain(self, node: TreeNode) -> None:
|
||||||
|
|||||||
@@ -152,12 +152,20 @@ class FullComponent(TreeComponent):
|
|||||||
) -> IncLockRefResult:
|
) -> IncLockRefResult:
|
||||||
ct = self.component_type
|
ct = self.component_type
|
||||||
root = self.cache.root_node
|
root = self.cache.root_node
|
||||||
delta = 0
|
|
||||||
cur = node
|
cur = node
|
||||||
while cur != root:
|
|
||||||
cd = cur.component_data[ct]
|
|
||||||
assert cd.value is not None
|
|
||||||
|
|
||||||
|
# Skip the bottom evicted segment
|
||||||
|
while cur is not root and cur.component_data[ct].value is None:
|
||||||
|
result.skip_lock_node_ids.setdefault(ct, set()).add(cur.id)
|
||||||
|
cur = cur.parent
|
||||||
|
|
||||||
|
# Lock the device-on segment up to root
|
||||||
|
delta = 0
|
||||||
|
while cur is not root:
|
||||||
|
cd = cur.component_data[ct]
|
||||||
|
assert (
|
||||||
|
cd.value is not None
|
||||||
|
), f"FULL invariant broken: evicted ancestor {cur.id} above device-on segment"
|
||||||
if cd.lock_ref == 0:
|
if cd.lock_ref == 0:
|
||||||
key_len = len(cd.value)
|
key_len = len(cd.value)
|
||||||
self.cache.component_evictable_size_[ct] -= key_len
|
self.cache.component_evictable_size_[ct] -= key_len
|
||||||
@@ -174,8 +182,12 @@ class FullComponent(TreeComponent):
|
|||||||
) -> None:
|
) -> None:
|
||||||
ct = self.component_type
|
ct = self.component_type
|
||||||
root = self.cache.root_node
|
root = self.cache.root_node
|
||||||
|
skip_lock_node_ids = params.skip_lock_node_ids.get(ct, ()) if params else ()
|
||||||
cur = node
|
cur = node
|
||||||
while cur != root:
|
while cur != root:
|
||||||
|
if cur.id in skip_lock_node_ids:
|
||||||
|
cur = cur.parent
|
||||||
|
continue
|
||||||
cd = cur.component_data[ct]
|
cd = cur.component_data[ct]
|
||||||
assert cd.value is not None
|
assert cd.value is not None
|
||||||
assert cd.lock_ref > 0
|
assert cd.lock_ref > 0
|
||||||
@@ -203,13 +215,14 @@ class FullComponent(TreeComponent):
|
|||||||
return None
|
return None
|
||||||
|
|
||||||
if phase == CacheTransferPhase.LOAD_BACK:
|
if phase == CacheTransferPhase.LOAD_BACK:
|
||||||
# Walk evicted chain, collect host_values and nodes
|
# `node` is best_match_node. FULL device evict only from leaves,
|
||||||
|
# so once we hit a device-on node, everything above is also device-on
|
||||||
backed_up: list[torch.Tensor] = []
|
backed_up: list[torch.Tensor] = []
|
||||||
nodes: list = []
|
nodes: list = []
|
||||||
cur = node
|
cur = node
|
||||||
while cur.evicted:
|
while cur.evicted:
|
||||||
cd = cur.component_data[ct]
|
cd = cur.component_data[ct]
|
||||||
if cd.host_value is not None:
|
assert cd.host_value is not None
|
||||||
backed_up.append(cd.host_value)
|
backed_up.append(cd.host_value)
|
||||||
nodes.append(cur)
|
nodes.append(cur)
|
||||||
cur = cur.parent
|
cur = cur.parent
|
||||||
|
|||||||
@@ -67,7 +67,7 @@ class MambaComponent(TreeComponent):
|
|||||||
) -> MatchResult:
|
) -> MatchResult:
|
||||||
cow_mamba = params.cow_mamba
|
cow_mamba = params.cow_mamba
|
||||||
req = params.req
|
req = params.req
|
||||||
last_node = result.last_device_node
|
last_node = result.best_match_node
|
||||||
|
|
||||||
if len(value_chunks) > best_value_len:
|
if len(value_chunks) > best_value_len:
|
||||||
chunk_size = get_global_server_args().mamba_cache_chunk_size
|
chunk_size = get_global_server_args().mamba_cache_chunk_size
|
||||||
@@ -101,8 +101,7 @@ class MambaComponent(TreeComponent):
|
|||||||
|
|
||||||
# HiCache: if mamba was evicted from device but has host backup,
|
# HiCache: if mamba was evicted from device but has host backup,
|
||||||
# ensure host_hit_length >= 1 so load_back is triggered.
|
# ensure host_hit_length >= 1 so load_back is triggered.
|
||||||
host_node = result.last_host_node
|
cd = last_node.component_data[self.component_type]
|
||||||
cd = host_node.component_data[self.component_type]
|
|
||||||
if cd.value is None and cd.host_value is not None:
|
if cd.value is None and cd.host_value is not None:
|
||||||
result = result._replace(host_hit_length=max(result.host_hit_length, 1))
|
result = result._replace(host_hit_length=max(result.host_hit_length, 1))
|
||||||
|
|
||||||
|
|||||||
@@ -95,11 +95,12 @@ class SWAComponent(TreeComponent):
|
|||||||
) -> MatchResult:
|
) -> MatchResult:
|
||||||
ct = self.component_type
|
ct = self.component_type
|
||||||
n_swa = 0
|
n_swa = 0
|
||||||
node = result.last_host_node
|
node = result.best_match_node
|
||||||
root = self.cache.root_node
|
root = self.cache.root_node
|
||||||
while node is not root and n_swa < self.sliding_window_size:
|
while node is not root and n_swa < self.sliding_window_size:
|
||||||
cd = node.component_data[ct]
|
cd = node.component_data[ct]
|
||||||
if cd.value is None and cd.host_value is not None:
|
if cd.value is None and cd.host_value is not None:
|
||||||
|
# TODO(ispobock): refactor host_hit_length usage
|
||||||
return result._replace(host_hit_length=max(result.host_hit_length, 1))
|
return result._replace(host_hit_length=max(result.host_hit_length, 1))
|
||||||
if cd.value is not None:
|
if cd.value is not None:
|
||||||
n_swa += len(cd.value)
|
n_swa += len(cd.value)
|
||||||
@@ -440,11 +441,14 @@ class SWAComponent(TreeComponent):
|
|||||||
]
|
]
|
||||||
|
|
||||||
if phase == CacheTransferPhase.LOAD_BACK:
|
if phase == CacheTransferPhase.LOAD_BACK:
|
||||||
|
# `node` is best_match_node; the SWA validator guarantees every
|
||||||
|
# ancestor within `sliding_window_size` has value or host_value.
|
||||||
n_swa = 0
|
n_swa = 0
|
||||||
backed_up: list[torch.Tensor] = []
|
backed_up: list[torch.Tensor] = []
|
||||||
nodes: list = []
|
nodes: list = []
|
||||||
while node is not self.cache.root_node and n_swa < self.sliding_window_size:
|
cur = node
|
||||||
cd = node.component_data[ct]
|
while cur is not self.cache.root_node and n_swa < self.sliding_window_size:
|
||||||
|
cd = cur.component_data[ct]
|
||||||
assert cd.host_value is not None or cd.value is not None
|
assert cd.host_value is not None or cd.value is not None
|
||||||
if cd.value is not None:
|
if cd.value is not None:
|
||||||
# device exists, skip it
|
# device exists, skip it
|
||||||
@@ -452,9 +456,9 @@ class SWAComponent(TreeComponent):
|
|||||||
else:
|
else:
|
||||||
# host only, collect it
|
# host only, collect it
|
||||||
backed_up.append(cd.host_value)
|
backed_up.append(cd.host_value)
|
||||||
nodes.append(node)
|
nodes.append(cur)
|
||||||
n_swa += len(cd.host_value)
|
n_swa += len(cd.host_value)
|
||||||
node = node.parent
|
cur = cur.parent
|
||||||
|
|
||||||
if not backed_up:
|
if not backed_up:
|
||||||
return None
|
return None
|
||||||
|
|||||||
@@ -290,6 +290,7 @@ class UnifiedRadixCache(BasePrefixCache):
|
|||||||
),
|
),
|
||||||
last_device_node=self.root_node,
|
last_device_node=self.root_node,
|
||||||
last_host_node=self.root_node,
|
last_host_node=self.root_node,
|
||||||
|
best_match_node=self.root_node,
|
||||||
)
|
)
|
||||||
|
|
||||||
def init_hicache(self, server_args: ServerArgs, params: CacheInitParams) -> None:
|
def init_hicache(self, server_args: ServerArgs, params: CacheInitParams) -> None:
|
||||||
@@ -352,8 +353,10 @@ class UnifiedRadixCache(BasePrefixCache):
|
|||||||
if len(key) == 0:
|
if len(key) == 0:
|
||||||
return self._empty_match_result
|
return self._empty_match_result
|
||||||
|
|
||||||
value, last_node, best_value_len = self._match_prefix_helper(key)
|
value, best_match_node, best_value_len = self._match_prefix_helper(key)
|
||||||
return self._match_post_processor(params, value, last_node, best_value_len)
|
return self._match_post_processor(
|
||||||
|
params, value, best_match_node, best_value_len
|
||||||
|
)
|
||||||
|
|
||||||
def insert(self, params: InsertParams) -> InsertResult:
|
def insert(self, params: InsertParams) -> InsertResult:
|
||||||
if self.disable:
|
if self.disable:
|
||||||
@@ -595,16 +598,16 @@ class UnifiedRadixCache(BasePrefixCache):
|
|||||||
child_key = key.child_key(self.page_size)
|
child_key = key.child_key(self.page_size)
|
||||||
value: list[torch.Tensor] = []
|
value: list[torch.Tensor] = []
|
||||||
best_value_len = 0
|
best_value_len = 0
|
||||||
best_node = node
|
best_match_node = node
|
||||||
validators = tuple(
|
validators = tuple(
|
||||||
comp.create_match_validator() for comp in self._components_tuple
|
comp.create_match_validator() for comp in self._components_tuple
|
||||||
)
|
)
|
||||||
|
|
||||||
def _update_best_if_valid(node):
|
def _update_best_if_valid(node):
|
||||||
nonlocal best_value_len, best_node
|
nonlocal best_value_len, best_match_node
|
||||||
if all(v(node) for v in validators):
|
if all(v(node) for v in validators):
|
||||||
best_value_len = len(value)
|
best_value_len = len(value)
|
||||||
best_node = node
|
best_match_node = node
|
||||||
|
|
||||||
while len(key) > 0 and child_key in node.children:
|
while len(key) > 0 and child_key in node.children:
|
||||||
child = node.children[child_key]
|
child = node.children[child_key]
|
||||||
@@ -625,7 +628,7 @@ class UnifiedRadixCache(BasePrefixCache):
|
|||||||
key = key[prefix_len:]
|
key = key[prefix_len:]
|
||||||
if len(key):
|
if len(key):
|
||||||
child_key = key.child_key(self.page_size)
|
child_key = key.child_key(self.page_size)
|
||||||
return value, best_node, best_value_len
|
return value, best_match_node, best_value_len
|
||||||
|
|
||||||
def _match_prefix_helper(
|
def _match_prefix_helper(
|
||||||
self, key: RadixKey
|
self, key: RadixKey
|
||||||
@@ -634,16 +637,16 @@ class UnifiedRadixCache(BasePrefixCache):
|
|||||||
child_key = key.child_key(self.page_size)
|
child_key = key.child_key(self.page_size)
|
||||||
value: list[torch.Tensor] = []
|
value: list[torch.Tensor] = []
|
||||||
best_value_len = 0
|
best_value_len = 0
|
||||||
best_node = node
|
best_match_node = node
|
||||||
validators = tuple(
|
validators = tuple(
|
||||||
comp.create_match_validator() for comp in self._components_tuple
|
comp.create_match_validator() for comp in self._components_tuple
|
||||||
)
|
)
|
||||||
|
|
||||||
def _update_best_if_valid(node):
|
def _update_best_if_valid(node):
|
||||||
nonlocal best_value_len, best_node
|
nonlocal best_value_len, best_match_node
|
||||||
if all(v(node) for v in validators):
|
if all(v(node) for v in validators):
|
||||||
best_value_len = len(value)
|
best_value_len = len(value)
|
||||||
best_node = node
|
best_match_node = node
|
||||||
|
|
||||||
while len(key) > 0 and child_key in node.children:
|
while len(key) > 0 and child_key in node.children:
|
||||||
child = node.children[child_key]
|
child = node.children[child_key]
|
||||||
@@ -667,16 +670,16 @@ class UnifiedRadixCache(BasePrefixCache):
|
|||||||
key = key[prefix_len:]
|
key = key[prefix_len:]
|
||||||
if len(key):
|
if len(key):
|
||||||
child_key = key.child_key(self.page_size)
|
child_key = key.child_key(self.page_size)
|
||||||
return value, best_node, best_value_len
|
return value, best_match_node, best_value_len
|
||||||
|
|
||||||
def _match_post_processor(
|
def _match_post_processor(
|
||||||
self,
|
self,
|
||||||
params: MatchPrefixParams,
|
params: MatchPrefixParams,
|
||||||
value: list[torch.Tensor],
|
value: list[torch.Tensor],
|
||||||
last_node: UnifiedTreeNode,
|
best_match_node: UnifiedTreeNode,
|
||||||
best_value_len: int,
|
best_value_len: int,
|
||||||
) -> MatchResult:
|
) -> MatchResult:
|
||||||
node_update = last_node
|
node_update = best_match_node
|
||||||
for comp in self._components_tuple:
|
for comp in self._components_tuple:
|
||||||
if comp.component_type == BASE_COMPONENT_TYPE:
|
if comp.component_type == BASE_COMPONENT_TYPE:
|
||||||
continue # Full uses last_access_time, not LRU
|
continue # Full uses last_access_time, not LRU
|
||||||
@@ -691,12 +694,12 @@ class UnifiedRadixCache(BasePrefixCache):
|
|||||||
node_update = node_update.parent
|
node_update = node_update.parent
|
||||||
|
|
||||||
# Walk up to find last_device_node
|
# Walk up to find last_device_node
|
||||||
last_device_node = last_node
|
last_device_node = best_match_node
|
||||||
while last_device_node is not self.root_node and last_device_node.evicted:
|
while last_device_node is not self.root_node and last_device_node.evicted:
|
||||||
last_device_node = last_device_node.parent
|
last_device_node = last_device_node.parent
|
||||||
|
|
||||||
# Walk up to find last_host_node
|
# Walk up to find last_host_node
|
||||||
last_host_node = last_node
|
last_host_node = best_match_node
|
||||||
while last_host_node is not self.root_node and not last_host_node.backuped:
|
while last_host_node is not self.root_node and not last_host_node.backuped:
|
||||||
last_host_node = last_host_node.parent
|
last_host_node = last_host_node.parent
|
||||||
|
|
||||||
@@ -708,6 +711,7 @@ class UnifiedRadixCache(BasePrefixCache):
|
|||||||
device_indices=device_indices,
|
device_indices=device_indices,
|
||||||
last_device_node=last_device_node,
|
last_device_node=last_device_node,
|
||||||
last_host_node=last_host_node,
|
last_host_node=last_host_node,
|
||||||
|
best_match_node=best_match_node,
|
||||||
host_hit_length=0,
|
host_hit_length=0,
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -1213,7 +1217,7 @@ class UnifiedRadixCache(BasePrefixCache):
|
|||||||
|
|
||||||
def load_back(
|
def load_back(
|
||||||
self,
|
self,
|
||||||
node: UnifiedTreeNode,
|
best_match_node: UnifiedTreeNode,
|
||||||
mem_quota: Optional[int] = None,
|
mem_quota: Optional[int] = None,
|
||||||
req=None,
|
req=None,
|
||||||
) -> Optional[torch.Tensor]:
|
) -> Optional[torch.Tensor]:
|
||||||
@@ -1222,15 +1226,12 @@ class UnifiedRadixCache(BasePrefixCache):
|
|||||||
return None
|
return None
|
||||||
|
|
||||||
# Build KV transfer
|
# Build KV transfer
|
||||||
last_hit_node = node
|
|
||||||
kv_xfer = self.components[BASE_COMPONENT_TYPE].build_hicache_transfers(
|
kv_xfer = self.components[BASE_COMPONENT_TYPE].build_hicache_transfers(
|
||||||
last_hit_node, CacheTransferPhase.LOAD_BACK
|
best_match_node, CacheTransferPhase.LOAD_BACK
|
||||||
)[0]
|
)[0]
|
||||||
|
|
||||||
# Lock path & pre-evict if device pool is insufficient
|
# Lock path & pre-evict if device pool is insufficient
|
||||||
nodes_to_load = kv_xfer.nodes_to_load
|
result = self.inc_lock_ref(best_match_node)
|
||||||
ancestor_node = nodes_to_load[0].parent if nodes_to_load else last_hit_node
|
|
||||||
result = self.inc_lock_ref(ancestor_node)
|
|
||||||
ancestor_lock_params = result.to_dec_params()
|
ancestor_lock_params = result.to_dec_params()
|
||||||
kv_tokens = len(kv_xfer.host_indices)
|
kv_tokens = len(kv_xfer.host_indices)
|
||||||
|
|
||||||
@@ -1240,7 +1241,7 @@ class UnifiedRadixCache(BasePrefixCache):
|
|||||||
if comp.component_type == BASE_COMPONENT_TYPE:
|
if comp.component_type == BASE_COMPONENT_TYPE:
|
||||||
continue
|
continue
|
||||||
t = comp.build_hicache_transfers(
|
t = comp.build_hicache_transfers(
|
||||||
last_hit_node, CacheTransferPhase.LOAD_BACK, req=req
|
best_match_node, CacheTransferPhase.LOAD_BACK, req=req
|
||||||
)
|
)
|
||||||
if t:
|
if t:
|
||||||
comp_xfers[comp.component_type] = t
|
comp_xfers[comp.component_type] = t
|
||||||
@@ -1255,7 +1256,7 @@ class UnifiedRadixCache(BasePrefixCache):
|
|||||||
if (kv_tokens < self.load_back_threshold and not comp_xfers) or (
|
if (kv_tokens < self.load_back_threshold and not comp_xfers) or (
|
||||||
mem_quota is not None and kv_tokens > mem_quota + result.delta
|
mem_quota is not None and kv_tokens > mem_quota + result.delta
|
||||||
):
|
):
|
||||||
self.dec_lock_ref(ancestor_node, ancestor_lock_params)
|
self.dec_lock_ref(best_match_node, ancestor_lock_params)
|
||||||
return None
|
return None
|
||||||
|
|
||||||
avail = self.token_to_kv_pool_allocator.available_size()
|
avail = self.token_to_kv_pool_allocator.available_size()
|
||||||
@@ -1263,7 +1264,7 @@ class UnifiedRadixCache(BasePrefixCache):
|
|||||||
needed = kv_tokens - avail
|
needed = kv_tokens - avail
|
||||||
result = self.evict(EvictParams(num_tokens=needed))
|
result = self.evict(EvictParams(num_tokens=needed))
|
||||||
if result.num_tokens_evicted < needed:
|
if result.num_tokens_evicted < needed:
|
||||||
self.dec_lock_ref(ancestor_node, ancestor_lock_params)
|
self.dec_lock_ref(best_match_node, ancestor_lock_params)
|
||||||
return None
|
return None
|
||||||
|
|
||||||
# Load H→D
|
# Load H→D
|
||||||
@@ -1271,32 +1272,32 @@ class UnifiedRadixCache(BasePrefixCache):
|
|||||||
aux_xfers.extend(anchor_kv_shared_indices_xfers)
|
aux_xfers.extend(anchor_kv_shared_indices_xfers)
|
||||||
device_indices = self.cache_controller.load(
|
device_indices = self.cache_controller.load(
|
||||||
host_indices=kv_xfer.host_indices,
|
host_indices=kv_xfer.host_indices,
|
||||||
node_id=last_hit_node.id,
|
node_id=best_match_node.id,
|
||||||
extra_pools=aux_xfers or None,
|
extra_pools=aux_xfers or None,
|
||||||
)
|
)
|
||||||
|
|
||||||
self.dec_lock_ref(ancestor_node, ancestor_lock_params)
|
self.dec_lock_ref(best_match_node, ancestor_lock_params)
|
||||||
if device_indices is None:
|
if device_indices is None:
|
||||||
return None
|
return None
|
||||||
|
|
||||||
# Commit: each component gets only its own transfers
|
# Commit: each component gets only its own transfers
|
||||||
kv_xfer.device_indices = device_indices
|
kv_xfer.device_indices = device_indices
|
||||||
self.components[BASE_COMPONENT_TYPE].commit_hicache_transfer(
|
self.components[BASE_COMPONENT_TYPE].commit_hicache_transfer(
|
||||||
last_hit_node,
|
best_match_node,
|
||||||
CacheTransferPhase.LOAD_BACK,
|
CacheTransferPhase.LOAD_BACK,
|
||||||
[kv_xfer],
|
[kv_xfer],
|
||||||
)
|
)
|
||||||
for ct, xfers in comp_xfers.items():
|
for ct, xfers in comp_xfers.items():
|
||||||
self.components[ct].commit_hicache_transfer(
|
self.components[ct].commit_hicache_transfer(
|
||||||
last_hit_node,
|
best_match_node,
|
||||||
CacheTransferPhase.LOAD_BACK,
|
CacheTransferPhase.LOAD_BACK,
|
||||||
xfers,
|
xfers,
|
||||||
)
|
)
|
||||||
|
|
||||||
self._update_evictable_leaf_sets(ancestor_node)
|
self._update_evictable_leaf_sets(best_match_node)
|
||||||
self.ongoing_load_back[last_hit_node.id] = (
|
self.ongoing_load_back[best_match_node.id] = (
|
||||||
last_hit_node,
|
best_match_node,
|
||||||
self.inc_lock_ref(last_hit_node).to_dec_params(),
|
self.inc_lock_ref(best_match_node).to_dec_params(),
|
||||||
)
|
)
|
||||||
return device_indices
|
return device_indices
|
||||||
|
|
||||||
@@ -1384,27 +1385,28 @@ class UnifiedRadixCache(BasePrefixCache):
|
|||||||
) -> tuple[torch.Tensor, UnifiedTreeNode]:
|
) -> tuple[torch.Tensor, UnifiedTreeNode]:
|
||||||
"""Prepare KV cache loading from host to device.
|
"""Prepare KV cache loading from host to device.
|
||||||
Returns (device_indices, last_node) tuple."""
|
Returns (device_indices, last_node) tuple."""
|
||||||
last_node = params.last_host_node
|
best_match_node = params.best_match_node
|
||||||
mem_quota = params.mem_quota
|
mem_quota = params.mem_quota
|
||||||
req = params.req
|
req = params.req
|
||||||
|
|
||||||
if last_node.evicted or params.host_hit_length > 0:
|
if best_match_node.evicted or params.host_hit_length > 0:
|
||||||
loading_values = self.load_back(last_node, mem_quota, req=req)
|
loading_values = self.load_back(best_match_node, mem_quota, req=req)
|
||||||
if loading_values is not None:
|
if loading_values is not None:
|
||||||
logger.debug(
|
logger.debug(
|
||||||
"init_load_back success: loaded %d tokens for node %d",
|
"init_load_back success: loaded %d tokens for node %d",
|
||||||
len(loading_values),
|
len(loading_values),
|
||||||
last_node.id,
|
best_match_node.id,
|
||||||
)
|
)
|
||||||
return loading_values, last_node
|
return loading_values, best_match_node
|
||||||
|
|
||||||
# Fallback: walk up to non-evicted ancestor
|
# Fallback: walk up to non-evicted ancestor
|
||||||
while last_node is not self.root_node and last_node.evicted:
|
# TODO(ispobock): The fallback path is not correct. The last_device_node should consider all the components.
|
||||||
last_node = last_node.parent
|
while best_match_node is not self.root_node and best_match_node.evicted:
|
||||||
|
best_match_node = best_match_node.parent
|
||||||
|
|
||||||
return (
|
return (
|
||||||
self._empty_match_result.device_indices,
|
self._empty_match_result.device_indices,
|
||||||
last_node,
|
best_match_node,
|
||||||
)
|
)
|
||||||
|
|
||||||
def check_hicache_events(self) -> None:
|
def check_hicache_events(self) -> None:
|
||||||
|
|||||||
@@ -259,6 +259,7 @@ class StreamingSession(BasePrefixCache):
|
|||||||
device_indices=device_indices,
|
device_indices=device_indices,
|
||||||
last_device_node=slot.virtual_node,
|
last_device_node=slot.virtual_node,
|
||||||
last_host_node=slot.virtual_node,
|
last_host_node=slot.virtual_node,
|
||||||
|
best_match_node=slot.virtual_node,
|
||||||
cache_protected_len=slot.cache_protected_len,
|
cache_protected_len=slot.cache_protected_len,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|||||||
@@ -33,6 +33,7 @@ class _StubReq:
|
|||||||
self.prefix_indices = None
|
self.prefix_indices = None
|
||||||
self.last_node = None
|
self.last_node = None
|
||||||
self.last_host_node = None
|
self.last_host_node = None
|
||||||
|
self.best_match_node = None
|
||||||
self.host_hit_length = None
|
self.host_hit_length = None
|
||||||
self.mamba_branching_seqlen = None
|
self.mamba_branching_seqlen = None
|
||||||
self.cache_protected_len = None
|
self.cache_protected_len = None
|
||||||
@@ -50,6 +51,7 @@ class TestZeroMatchResult(unittest.TestCase):
|
|||||||
self.assertEqual(int(zeroed.device_indices.numel()), 0)
|
self.assertEqual(int(zeroed.device_indices.numel()), 0)
|
||||||
self.assertIs(zeroed.last_device_node, tree.root_node)
|
self.assertIs(zeroed.last_device_node, tree.root_node)
|
||||||
self.assertIs(zeroed.last_host_node, tree.root_node)
|
self.assertIs(zeroed.last_host_node, tree.root_node)
|
||||||
|
self.assertIs(zeroed.best_match_node, tree.root_node)
|
||||||
self.assertEqual(zeroed.host_hit_length, 0)
|
self.assertEqual(zeroed.host_hit_length, 0)
|
||||||
# dtype/device preserved (slice-not-allocate).
|
# dtype/device preserved (slice-not-allocate).
|
||||||
self.assertEqual(zeroed.device_indices.dtype, match.device_indices.dtype)
|
self.assertEqual(zeroed.device_indices.dtype, match.device_indices.dtype)
|
||||||
@@ -64,6 +66,7 @@ class TestZeroMatchResult(unittest.TestCase):
|
|||||||
device_indices=torch.empty((0,), dtype=torch.int64),
|
device_indices=torch.empty((0,), dtype=torch.int64),
|
||||||
last_device_node=None,
|
last_device_node=None,
|
||||||
last_host_node=None,
|
last_host_node=None,
|
||||||
|
best_match_node=None,
|
||||||
host_hit_length=0,
|
host_hit_length=0,
|
||||||
)
|
)
|
||||||
self.assertIs(zero_match_result(_StubChunkCache(), original), original)
|
self.assertIs(zero_match_result(_StubChunkCache(), original), original)
|
||||||
|
|||||||
@@ -104,6 +104,7 @@ def test_preabort_detaches_session_and_preserves_slot():
|
|||||||
device_indices=torch.tensor([], dtype=torch.int64),
|
device_indices=torch.tensor([], dtype=torch.int64),
|
||||||
last_device_node=None,
|
last_device_node=None,
|
||||||
last_host_node=None,
|
last_host_node=None,
|
||||||
|
best_match_node=None,
|
||||||
)
|
)
|
||||||
],
|
],
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -922,19 +922,19 @@ class UnifiedRadixCacheSuite:
|
|||||||
self.assertEqual(node_count_before, 2)
|
self.assertEqual(node_count_before, 2)
|
||||||
|
|
||||||
tree._match_prefix_helper(RadixKey([1, 2]))
|
tree._match_prefix_helper(RadixKey([1, 2]))
|
||||||
value, best_node, best_value_len = tree._match_prefix_helper(
|
value, best_match_node, best_value_len = tree._match_prefix_helper(
|
||||||
RadixKey([1, 2, 3, 4])
|
RadixKey([1, 2, 3, 4])
|
||||||
)
|
)
|
||||||
self.assertEqual(best_value_len, 2)
|
self.assertEqual(best_value_len, 2)
|
||||||
self.assertEqual(best_node.key.token_ids, [3, 4])
|
self.assertEqual(best_match_node.key.token_ids, [3, 4])
|
||||||
node_count_after_regular = count_nodes(tree.root_node)
|
node_count_after_regular = count_nodes(tree.root_node)
|
||||||
self.assertEqual(node_count_after_regular, node_count_before + 2)
|
self.assertEqual(node_count_after_regular, node_count_before + 2)
|
||||||
|
|
||||||
value, best_node, best_value_len = tree._match_prefix_helper_readonly(
|
value, best_match_node, best_value_len = tree._match_prefix_helper_readonly(
|
||||||
RadixKey([1, 2, 3])
|
RadixKey([1, 2, 3])
|
||||||
)
|
)
|
||||||
self.assertEqual(best_value_len, 1)
|
self.assertEqual(best_value_len, 1)
|
||||||
self.assertEqual(best_node.key.token_ids, [1, 2])
|
self.assertEqual(best_match_node.key.token_ids, [1, 2])
|
||||||
node_count_after_readonly = count_nodes(tree.root_node)
|
node_count_after_readonly = count_nodes(tree.root_node)
|
||||||
self.assertEqual(node_count_after_readonly, node_count_after_regular)
|
self.assertEqual(node_count_after_readonly, node_count_after_regular)
|
||||||
|
|
||||||
@@ -1821,6 +1821,7 @@ class UnifiedRadixCacheSuite:
|
|||||||
),
|
),
|
||||||
last_device_node=leaf,
|
last_device_node=leaf,
|
||||||
last_host_node=leaf,
|
last_host_node=leaf,
|
||||||
|
best_match_node=leaf,
|
||||||
host_hit_length=0,
|
host_hit_length=0,
|
||||||
)
|
)
|
||||||
result = swa_comp.finalize_match_result(
|
result = swa_comp.finalize_match_result(
|
||||||
@@ -1905,6 +1906,102 @@ class UnifiedRadixCacheSuite:
|
|||||||
n_swa,
|
n_swa,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
def _swa_anchor_chain_tokens(self, num_pages: int) -> list[int]:
|
||||||
|
"""Reproduce the token sequence used by _build_chain_pages."""
|
||||||
|
tokens: list[int] = []
|
||||||
|
for i in range(num_pages):
|
||||||
|
tokens.extend(self._make_seq(1000 * (i + 1), 1))
|
||||||
|
return tokens
|
||||||
|
|
||||||
|
def _swa_anchor_setup(self):
|
||||||
|
"""Chain layout (root-to-leaf):
|
||||||
|
|
||||||
|
... top padding ...
|
||||||
|
chain[-(window_pages + 1)] = N : SWA tombstone
|
||||||
|
chain[-(window_pages)] = Y : SWA host-only, FULL host-backed
|
||||||
|
chain[-(window_pages - 1)..-2] : SWA device, FULL.host=None
|
||||||
|
chain[-1] = X : SWA device, FULL.host=None
|
||||||
|
|
||||||
|
Anchor=X has window_pages of device pages below N, so the SWA
|
||||||
|
load_back walker accumulates n_swa >= sliding_window_size and
|
||||||
|
exits before reaching N. Anchor=Y has only 1 page and walks into
|
||||||
|
N's tombstone.
|
||||||
|
"""
|
||||||
|
if not self.cfg.has_swa:
|
||||||
|
self.skipTest("requires SWA")
|
||||||
|
if self.cfg.has_mamba:
|
||||||
|
self.skipTest("SWA-only path keeps the chain construction simple")
|
||||||
|
ps = self.cfg.page_size
|
||||||
|
sw = self.cfg.sliding_window_size
|
||||||
|
if ps >= sw:
|
||||||
|
self.skipTest("test scenario requires ps < sw")
|
||||||
|
window_pages = (sw + ps - 1) // ps
|
||||||
|
chain_pages = window_pages + 3
|
||||||
|
if chain_pages * ps > self.cfg.kv_size // 2:
|
||||||
|
self.skipTest("kv_size too small for the desired chain")
|
||||||
|
|
||||||
|
tree, allocator, req_to_token_pool = build_fixture(self.cfg)
|
||||||
|
chain = self._build_chain_pages(tree, allocator, req_to_token_pool, chain_pages)
|
||||||
|
if len(chain) < chain_pages:
|
||||||
|
self.skipTest("chain too short")
|
||||||
|
self._simulate_backup_tree(tree)
|
||||||
|
|
||||||
|
x = chain[-1]
|
||||||
|
y = chain[-window_pages]
|
||||||
|
n = chain[-(window_pages + 1)]
|
||||||
|
|
||||||
|
n.component_data[ComponentType.SWA].value = None
|
||||||
|
n.component_data[ComponentType.SWA].host_value = None
|
||||||
|
y.component_data[ComponentType.SWA].value = None
|
||||||
|
# Strip FULL.host on X + intermediates so last_host_node walks past
|
||||||
|
# them to Y. Y.FULL untouched preserves the leaf-up evict invariant.
|
||||||
|
for node in chain[-(window_pages - 1) :]:
|
||||||
|
node.component_data[ComponentType.FULL].host_value = None
|
||||||
|
|
||||||
|
tokens = self._swa_anchor_chain_tokens(len(chain))
|
||||||
|
return tree, chain, n, y, x, tokens
|
||||||
|
|
||||||
|
def test_hicache_swa_match_prefix_picks_best_match_node_above_last_host(self):
|
||||||
|
tree, _, _, y, x, tokens = self._swa_anchor_setup()
|
||||||
|
result = tree.match_prefix(MatchPrefixParams(key=RadixKey(tokens)))
|
||||||
|
self.assertIs(result.best_match_node, x)
|
||||||
|
self.assertIs(result.last_device_node, x)
|
||||||
|
self.assertIs(result.last_host_node, y)
|
||||||
|
|
||||||
|
def test_hicache_swa_load_back_anchored_on_best_match_node(self):
|
||||||
|
tree, _, _, y, x, _ = self._swa_anchor_setup()
|
||||||
|
ps = self.cfg.page_size
|
||||||
|
swa_comp = tree.components[ComponentType.SWA]
|
||||||
|
|
||||||
|
transfers = swa_comp.build_hicache_transfers(x, CacheTransferPhase.LOAD_BACK)
|
||||||
|
self.assertEqual(len(transfers), 1)
|
||||||
|
xfer = transfers[0]
|
||||||
|
self.assertEqual(xfer.name, PoolName.SWA)
|
||||||
|
self.assertEqual(xfer.nodes_to_load, [y])
|
||||||
|
self.assertEqual(int(xfer.host_indices.numel()), ps)
|
||||||
|
|
||||||
|
with self.assertRaises(AssertionError):
|
||||||
|
swa_comp.build_hicache_transfers(y, CacheTransferPhase.LOAD_BACK)
|
||||||
|
|
||||||
|
def test_hicache_swa_finalize_anchored_on_best_match_node(self):
|
||||||
|
tree, _, _, y, x, _ = self._swa_anchor_setup()
|
||||||
|
swa_comp = tree.components[ComponentType.SWA]
|
||||||
|
|
||||||
|
base = MatchResult(
|
||||||
|
device_indices=torch.empty((0,), dtype=torch.int64, device=tree.device),
|
||||||
|
last_device_node=x,
|
||||||
|
last_host_node=y,
|
||||||
|
best_match_node=x,
|
||||||
|
host_hit_length=0,
|
||||||
|
)
|
||||||
|
result = swa_comp.finalize_match_result(
|
||||||
|
result=base,
|
||||||
|
params=MatchPrefixParams(key=RadixKey(self._make_seq(1, 1))),
|
||||||
|
value_chunks=[],
|
||||||
|
best_value_len=0,
|
||||||
|
)
|
||||||
|
self.assertEqual(result.host_hit_length, 1)
|
||||||
|
|
||||||
def test_hicache_swa_temp_lock_does_not_release_restored_tombstone(self):
|
def test_hicache_swa_temp_lock_does_not_release_restored_tombstone(self):
|
||||||
"""A temporary scheduler lock that skipped a SWA tombstone must not
|
"""A temporary scheduler lock that skipped a SWA tombstone must not
|
||||||
release later load-back/request locks after the tombstone is restored.
|
release later load-back/request locks after the tombstone is restored.
|
||||||
@@ -1950,6 +2047,61 @@ class UnifiedRadixCacheSuite:
|
|||||||
tree.dec_lock_ref(leaf, request_lock.to_dec_params())
|
tree.dec_lock_ref(leaf, request_lock.to_dec_params())
|
||||||
self.assertEqual(cd.lock_ref, 0)
|
self.assertEqual(cd.lock_ref, 0)
|
||||||
|
|
||||||
|
def test_hicache_full_temp_lock_skips_evicted_anchor_and_mirrors_on_release(
|
||||||
|
self,
|
||||||
|
):
|
||||||
|
"""Acquire records the evicted anchor in skip_lock_node_ids (phase 1)
|
||||||
|
and locks device-on ancestors only (phase 2). After load_back
|
||||||
|
restores the anchor, a second acquire covers it; releasing the
|
||||||
|
first must mirror the skip so the anchor's lock_ref is not
|
||||||
|
decremented twice.
|
||||||
|
"""
|
||||||
|
if self._skip_unsupported_hicache_test():
|
||||||
|
return
|
||||||
|
ps = self.cfg.page_size
|
||||||
|
if 3 * ps > self.cfg.kv_size // 2:
|
||||||
|
self.skipTest("kv_size too small")
|
||||||
|
tree, allocator, req_to_token_pool = self._build_hicache_fixture()
|
||||||
|
chain = self._build_chain_pages(tree, allocator, req_to_token_pool, 3)
|
||||||
|
if len(chain) < 3:
|
||||||
|
self.skipTest("chain too short")
|
||||||
|
a, y, anchor = chain
|
||||||
|
self._simulate_backup_tree(tree)
|
||||||
|
|
||||||
|
cd_anchor = anchor.component_data[ComponentType.FULL]
|
||||||
|
cd_a = a.component_data[ComponentType.FULL]
|
||||||
|
cd_y = y.component_data[ComponentType.FULL]
|
||||||
|
anchor_value = cd_anchor.value
|
||||||
|
cd_anchor.value = None
|
||||||
|
|
||||||
|
self.assertEqual(cd_anchor.lock_ref, 0)
|
||||||
|
self.assertEqual(cd_y.lock_ref, 0)
|
||||||
|
self.assertEqual(cd_a.lock_ref, 0)
|
||||||
|
|
||||||
|
temp_lock = tree.inc_lock_ref(anchor)
|
||||||
|
self.assertEqual(cd_anchor.lock_ref, 0)
|
||||||
|
self.assertEqual(cd_y.lock_ref, 1)
|
||||||
|
self.assertEqual(cd_a.lock_ref, 1)
|
||||||
|
self.assertIn(ComponentType.FULL, temp_lock.skip_lock_node_ids)
|
||||||
|
self.assertIn(anchor.id, temp_lock.skip_lock_node_ids[ComponentType.FULL])
|
||||||
|
|
||||||
|
cd_anchor.value = anchor_value
|
||||||
|
|
||||||
|
second_lock = tree.inc_lock_ref(anchor)
|
||||||
|
self.assertEqual(cd_anchor.lock_ref, 1)
|
||||||
|
self.assertEqual(cd_y.lock_ref, 2)
|
||||||
|
self.assertEqual(cd_a.lock_ref, 2)
|
||||||
|
|
||||||
|
tree.dec_lock_ref(anchor, temp_lock.to_dec_params())
|
||||||
|
self.assertEqual(cd_anchor.lock_ref, 1)
|
||||||
|
self.assertEqual(cd_y.lock_ref, 1)
|
||||||
|
self.assertEqual(cd_a.lock_ref, 1)
|
||||||
|
|
||||||
|
tree.dec_lock_ref(anchor, second_lock.to_dec_params())
|
||||||
|
self.assertEqual(cd_anchor.lock_ref, 0)
|
||||||
|
self.assertEqual(cd_y.lock_ref, 0)
|
||||||
|
self.assertEqual(cd_a.lock_ref, 0)
|
||||||
|
|
||||||
def test_hicache_mamba_temp_lock_does_not_release_restored_tombstone(self):
|
def test_hicache_mamba_temp_lock_does_not_release_restored_tombstone(self):
|
||||||
"""A temporary scheduler lock that skipped a Mamba tombstone must not
|
"""A temporary scheduler lock that skipped a Mamba tombstone must not
|
||||||
release later load-back/request locks after the tombstone is restored.
|
release later load-back/request locks after the tombstone is restored.
|
||||||
|
|||||||
Reference in New Issue
Block a user