feat: add structured debug logging for CP HiCache write/evict/load paths
Add logger.info statements with [CacheCtrl-write], [MemCache-evict], [MemCache-alloc], [HiCache-write], [HiCache-evict], [HiCache-load] tags across cache_controller.py, common.py, and hiradix_cache.py to improve observability of cache operations.
This commit is contained in:
@@ -686,16 +686,32 @@ class HiCacheController:
|
||||
"""
|
||||
Back up KV caches from device memory to host memory.
|
||||
"""
|
||||
logger.info(
|
||||
"[CacheCtrl-write] write: node_id=%d len=%d cp=%s",
|
||||
node_id,
|
||||
len(device_indices),
|
||||
self.uses_cp_hicache,
|
||||
)
|
||||
if self.uses_cp_hicache:
|
||||
return self._write_cp(device_indices, priority, node_id)
|
||||
|
||||
host_indices = self.mem_pool_host.alloc(len(device_indices))
|
||||
if host_indices is None:
|
||||
logger.info(
|
||||
"[CacheCtrl-write] write non-CP FAILED (host full): node_id=%d len=%d",
|
||||
node_id,
|
||||
len(device_indices),
|
||||
)
|
||||
return None
|
||||
self.write_queue.append(
|
||||
CacheOperation(host_indices, device_indices, node_id, priority)
|
||||
)
|
||||
self.start_writing()
|
||||
logger.info(
|
||||
"[CacheCtrl-write] write non-CP submitted: node_id=%d len=%d",
|
||||
node_id,
|
||||
len(device_indices),
|
||||
)
|
||||
return host_indices
|
||||
|
||||
def start_writing(self) -> None:
|
||||
@@ -760,6 +776,11 @@ class HiCacheController:
|
||||
logical_len = len(device_indices)
|
||||
if owned_positions.numel() == 0:
|
||||
self._append_completed_write_ack(node_id)
|
||||
logger.info(
|
||||
"[CacheCtrl-write] _write_cp zero-owned rank: node_id=%d logical_len=%d",
|
||||
node_id,
|
||||
logical_len,
|
||||
)
|
||||
return HiCacheWriteResult(
|
||||
metadata=CpHiCacheNodeMetadata(
|
||||
logical_len=logical_len,
|
||||
@@ -772,6 +793,12 @@ class HiCacheController:
|
||||
physical_device_indices = layout.logical_locs_to_physical(owned_logical_indices)
|
||||
host_indices = self.mem_pool_host.alloc(len(physical_device_indices))
|
||||
if host_indices is None:
|
||||
logger.info(
|
||||
"[CacheCtrl-write] _write_cp FAILED (host full): node_id=%d logical_len=%d owned=%d",
|
||||
node_id,
|
||||
logical_len,
|
||||
owned_positions.numel(),
|
||||
)
|
||||
return HiCacheWriteFailure(required_host_slots=len(physical_device_indices))
|
||||
try:
|
||||
self._validate_cp_hicache_page_indices(host_indices, physical_device_indices)
|
||||
@@ -783,6 +810,13 @@ class HiCacheController:
|
||||
CacheOperation(host_indices, physical_device_indices, node_id, priority)
|
||||
)
|
||||
self.start_writing()
|
||||
logger.info(
|
||||
"[CacheCtrl-write] _write_cp submitted: node_id=%d logical_len=%d owned=%d physical=%d",
|
||||
node_id,
|
||||
logical_len,
|
||||
owned_positions.numel(),
|
||||
len(physical_device_indices),
|
||||
)
|
||||
return HiCacheWriteResult(
|
||||
metadata=CpHiCacheNodeMetadata(
|
||||
logical_len=logical_len,
|
||||
@@ -1104,7 +1138,7 @@ class HiCacheController:
|
||||
# not to prefetch if not enough benefits
|
||||
self.prefetch_revoke_queue.put(operation.request_id)
|
||||
self.append_host_mem_release(operation.host_indices)
|
||||
logger.debug(
|
||||
logger.info(
|
||||
f"Revoking prefetch for request {operation.request_id} due to insufficient hits ({storage_hit_count})."
|
||||
)
|
||||
else:
|
||||
@@ -1116,7 +1150,7 @@ class HiCacheController:
|
||||
operation.host_indices[storage_hit_count:]
|
||||
)
|
||||
operation.host_indices = operation.host_indices[:storage_hit_count]
|
||||
logger.debug(
|
||||
logger.info(
|
||||
f"Prefetching {len(operation.hash_value)} pages for request {operation.request_id}."
|
||||
)
|
||||
self.prefetch_buffer.put(operation)
|
||||
|
||||
@@ -283,7 +283,14 @@ def evict_from_tree_cache(
|
||||
return EvictResult()
|
||||
else:
|
||||
# Standard allocator
|
||||
if allocator.available_size() < num_tokens:
|
||||
available = allocator.available_size()
|
||||
if available < num_tokens:
|
||||
logger.info(
|
||||
"[MemCache-evict] evict_from_tree_cache: available=%d < num_tokens=%d deficit=%d, triggering eviction",
|
||||
available,
|
||||
num_tokens,
|
||||
num_tokens - available,
|
||||
)
|
||||
return tree_cache.evict(EvictParams(num_tokens=num_tokens))
|
||||
return EvictResult()
|
||||
|
||||
@@ -302,7 +309,7 @@ def _evict_for_compute_owner_lanes(
|
||||
return
|
||||
|
||||
max_attempts = max(2, min(8, int(getattr(allocator, "cp_size", 1))))
|
||||
for _ in range(max_attempts):
|
||||
for attempt in range(max_attempts):
|
||||
_required, _available, deficits = compute_owner_lane_stats(page_compute_owners)
|
||||
deficit_pages = sum(deficits)
|
||||
if deficit_pages <= 0:
|
||||
@@ -315,6 +322,10 @@ def _evict_for_compute_owner_lanes(
|
||||
if isinstance(evictable_size, tuple):
|
||||
evictable_size = evictable_size[0]
|
||||
if evictable_size <= 0:
|
||||
logger.info(
|
||||
"[MemCache-evict] _evict_for_compute_owner_lanes: evictable_size=%d <= 0, giving up",
|
||||
evictable_size,
|
||||
)
|
||||
return
|
||||
|
||||
evict_tokens = max(
|
||||
@@ -322,9 +333,23 @@ def _evict_for_compute_owner_lanes(
|
||||
deficit_pages * allocator.page_size * int(getattr(allocator, "cp_size", 1)),
|
||||
)
|
||||
before_available = allocator.available_size()
|
||||
logger.info(
|
||||
"[MemCache-evict] _evict_for_compute_owner_lanes attempt=%d: deficit_pages=%d evict_tokens=%d before_available=%d evictable_size=%d",
|
||||
attempt,
|
||||
deficit_pages,
|
||||
evict_tokens,
|
||||
before_available,
|
||||
evictable_size,
|
||||
)
|
||||
evict_result = tree_cache.evict(EvictParams(num_tokens=evict_tokens))
|
||||
after_available = allocator.available_size()
|
||||
evicted_tokens = getattr(evict_result, "num_tokens_evicted", 0)
|
||||
logger.info(
|
||||
"[MemCache-evict] _evict_for_compute_owner_lanes attempt=%d result: evicted=%d after_available=%d",
|
||||
attempt,
|
||||
evicted_tokens,
|
||||
after_available,
|
||||
)
|
||||
if after_available <= before_available and evicted_tokens <= 0:
|
||||
return
|
||||
|
||||
@@ -344,6 +369,17 @@ def alloc_paged_token_slots_extend(
|
||||
num_tokens = extend_num_tokens + len(seq_lens_cpu) * allocator.page_size
|
||||
evict_result = evict_from_tree_cache(tree_cache, num_tokens)
|
||||
|
||||
logger.info(
|
||||
"[MemCache-alloc] alloc_paged_token_slots_extend: extend_num_tokens=%d batch_size=%d num_tokens=%d page_size=%d "
|
||||
"available_size=%d evicted=%d",
|
||||
extend_num_tokens,
|
||||
len(seq_lens_cpu),
|
||||
num_tokens,
|
||||
allocator.page_size,
|
||||
allocator.available_size(),
|
||||
getattr(evict_result, "num_tokens_evicted", 0),
|
||||
)
|
||||
|
||||
alloc_extend_compute_owner = getattr(
|
||||
allocator, "alloc_extend_compute_owner", None
|
||||
)
|
||||
|
||||
@@ -727,6 +727,13 @@ class HiRadixCache(RadixCache):
|
||||
return node.host_value
|
||||
|
||||
def write_backup(self, node: TreeNode, write_back=False):
|
||||
logger.info(
|
||||
"[HiCache-write] write_backup triggered: node_id=%d len=%d cp=%s write_back=%s",
|
||||
node.id,
|
||||
len(node.value) if node.value is not None else 0,
|
||||
self._uses_cp_hicache,
|
||||
write_back,
|
||||
)
|
||||
if self._uses_cp_hicache:
|
||||
result = self.cache_controller.write(
|
||||
device_indices=node.value,
|
||||
@@ -741,12 +748,22 @@ class HiRadixCache(RadixCache):
|
||||
synchronize_across_ranks=getattr(self, "tp_world_size", 1) > 1,
|
||||
)
|
||||
if metadata is None:
|
||||
logger.info(
|
||||
"[HiCache-write] write_backup CP retry after host eviction: node_id=%d needed_slots=%d",
|
||||
node.id,
|
||||
required_host_slots,
|
||||
)
|
||||
result = self.cache_controller.write(
|
||||
device_indices=node.value,
|
||||
node_id=node.id,
|
||||
)
|
||||
metadata = getattr(result, "metadata", None)
|
||||
if metadata is None:
|
||||
logger.info(
|
||||
"[HiCache-write] write_backup CP FAILED (host full): node_id=%d len=%d",
|
||||
node.id,
|
||||
len(node.value) if node.value is not None else 0,
|
||||
)
|
||||
return 0
|
||||
|
||||
node.host_len = len(node.value)
|
||||
@@ -755,6 +772,13 @@ class HiRadixCache(RadixCache):
|
||||
self.ongoing_write_through[node.id] = node
|
||||
if not write_back:
|
||||
self.inc_node_lock_ref(node)
|
||||
logger.info(
|
||||
"[HiCache-write] write_backup CP SUCCESS: node_id=%d host_len=%d owned_positions=%d ongoing_writes=%d",
|
||||
node.id,
|
||||
node.host_len,
|
||||
node.cp_hicache.owned_positions.numel(),
|
||||
len(self.ongoing_write_through),
|
||||
)
|
||||
return len(node.value)
|
||||
|
||||
host_indices = self.cache_controller.write(
|
||||
@@ -762,6 +786,11 @@ class HiRadixCache(RadixCache):
|
||||
node_id=node.id,
|
||||
)
|
||||
if host_indices is None:
|
||||
logger.info(
|
||||
"[HiCache-write] write_backup non-CP retry after host eviction: node_id=%d len=%d",
|
||||
node.id,
|
||||
len(node.value) if node.value is not None else 0,
|
||||
)
|
||||
self.evict_host(len(node.value))
|
||||
host_indices = self.cache_controller.write(
|
||||
device_indices=node.value,
|
||||
@@ -778,7 +807,18 @@ class HiRadixCache(RadixCache):
|
||||
# parents with non-evicted children), so path locking is
|
||||
# unnecessary and would over-reduce evictable_size_.
|
||||
self.inc_node_lock_ref(node)
|
||||
logger.info(
|
||||
"[HiCache-write] write_backup non-CP SUCCESS: node_id=%d len=%d ongoing_writes=%d",
|
||||
node.id,
|
||||
len(host_indices),
|
||||
len(self.ongoing_write_through),
|
||||
)
|
||||
else:
|
||||
logger.info(
|
||||
"[HiCache-write] write_backup non-CP FAILED (host full): node_id=%d len=%d",
|
||||
node.id,
|
||||
len(node.value) if node.value is not None else 0,
|
||||
)
|
||||
return 0
|
||||
|
||||
return len(host_indices)
|
||||
@@ -804,6 +844,13 @@ class HiRadixCache(RadixCache):
|
||||
|
||||
if not self._node_backuped(node):
|
||||
if node.hit_count >= self.write_through_threshold:
|
||||
logger.info(
|
||||
"[HiCache-write] _inc_hit_count triggering write_through: node_id=%d hit_count=%d threshold=%d len=%d",
|
||||
node.id,
|
||||
node.hit_count,
|
||||
self.write_through_threshold,
|
||||
len(node.value) if node.value is not None else 0,
|
||||
)
|
||||
# write to host if the node is not backuped
|
||||
self.write_backup(node)
|
||||
|
||||
@@ -825,6 +872,7 @@ class HiRadixCache(RadixCache):
|
||||
if len(self.ongoing_write_through) == 0:
|
||||
return
|
||||
|
||||
ack_queue_len = len(self.cache_controller.ack_write_queue)
|
||||
finish_count = 0
|
||||
for _, finish_event, ack_list in self.cache_controller.ack_write_queue:
|
||||
if not finish_event.query():
|
||||
@@ -840,15 +888,32 @@ class HiRadixCache(RadixCache):
|
||||
)
|
||||
|
||||
finish_count = int(queue_size.item())
|
||||
logger.info(
|
||||
"[HiCache-write] writing_check: ongoing=%d ack_queue=%d local_finished=%d sync_finished=%d tp_size=%d",
|
||||
len(self.ongoing_write_through),
|
||||
ack_queue_len,
|
||||
finish_count if self.tp_world_size <= 1 else queue_size.item(),
|
||||
finish_count,
|
||||
getattr(self, "tp_world_size", 1),
|
||||
)
|
||||
released_nodes = []
|
||||
while finish_count > 0:
|
||||
_, finish_event, ack_list = self.cache_controller.ack_write_queue.pop(0)
|
||||
finish_event.synchronize()
|
||||
for ack_id in ack_list:
|
||||
backuped_node = self.ongoing_write_through.pop(ack_id)
|
||||
released_nodes.append(ack_id)
|
||||
self.dec_node_lock_ref(backuped_node)
|
||||
if self.enable_storage:
|
||||
self.write_backup_storage(backuped_node)
|
||||
finish_count -= 1
|
||||
if released_nodes:
|
||||
logger.info(
|
||||
"[HiCache-write] writing_check released %d write locks: node_ids=%s remaining_ongoing=%d",
|
||||
len(released_nodes),
|
||||
released_nodes,
|
||||
len(self.ongoing_write_through),
|
||||
)
|
||||
|
||||
def loading_check(self):
|
||||
finish_count = 0
|
||||
@@ -1011,12 +1076,21 @@ class HiRadixCache(RadixCache):
|
||||
]
|
||||
heapq.heapify(eviction_heap)
|
||||
|
||||
logger.info(
|
||||
"[HiCache-evict] evict START: num_tokens=%d heap_size=%d evictable_size=%d available_size=%d",
|
||||
num_tokens,
|
||||
len(eviction_heap),
|
||||
self.evictable_size_,
|
||||
self.token_to_kv_pool_allocator.available_size(),
|
||||
)
|
||||
num_evicted = 0
|
||||
num_locked_skipped = 0
|
||||
write_back_nodes = []
|
||||
while num_evicted < num_tokens and len(eviction_heap):
|
||||
_priority, x = heapq.heappop(eviction_heap)
|
||||
|
||||
if x.lock_ref > 0:
|
||||
num_locked_skipped += 1
|
||||
continue
|
||||
|
||||
if self._is_pinned(x):
|
||||
@@ -1066,6 +1140,14 @@ class HiRadixCache(RadixCache):
|
||||
self._evict_backuped(node)
|
||||
|
||||
self.update_eviction_metrics(num_evicted, start_time)
|
||||
logger.info(
|
||||
"[HiCache-evict] evict END: num_tokens=%d num_evicted=%d num_locked_skipped=%d evictable_size_after=%d available_size_after=%d",
|
||||
num_tokens,
|
||||
num_evicted,
|
||||
num_locked_skipped,
|
||||
self.evictable_size_,
|
||||
self.token_to_kv_pool_allocator.available_size(),
|
||||
)
|
||||
return EvictResult(num_tokens_evicted=num_evicted)
|
||||
|
||||
def _evict_backuped(self, node: TreeNode):
|
||||
@@ -1073,6 +1155,13 @@ class HiRadixCache(RadixCache):
|
||||
num_evicted = self.cache_controller.evict_device(node.value)
|
||||
assert num_evicted > 0
|
||||
self.evictable_size_ -= num_evicted
|
||||
logger.info(
|
||||
"[HiCache-evict] _evict_backuped: node_id=%d num_evicted=%d lock_ref=%d backed=%s",
|
||||
node.id,
|
||||
num_evicted,
|
||||
node.lock_ref,
|
||||
self._node_backuped(node),
|
||||
)
|
||||
node.value = None
|
||||
self._update_leaf_status(node)
|
||||
self._update_host_leaf_status(node)
|
||||
@@ -1082,9 +1171,14 @@ class HiRadixCache(RadixCache):
|
||||
|
||||
def _evict_regular(self, node: TreeNode):
|
||||
# evict a node not initiated write to host -- emit BlockRemoved
|
||||
num_evicted = len(node.value)
|
||||
logger.info(
|
||||
"[HiCache-evict] _evict_regular: node_id=%d num_evicted=%d",
|
||||
node.id,
|
||||
num_evicted,
|
||||
)
|
||||
self._record_remove_event(node)
|
||||
self.cache_controller.mem_pool_device_allocator.free(node.value)
|
||||
num_evicted = len(node.value)
|
||||
self._delete_leaf(node)
|
||||
return num_evicted
|
||||
|
||||
@@ -1110,6 +1204,12 @@ class HiRadixCache(RadixCache):
|
||||
return 0
|
||||
|
||||
leaves = list(self.evictable_host_leaves)
|
||||
logger.info(
|
||||
"[HiCache-evict] _evict_host_for_physical_slots: required_slots=%d sync=%s leaves=%d",
|
||||
required_host_slots,
|
||||
synchronize_across_ranks,
|
||||
len(leaves),
|
||||
)
|
||||
eviction_heap = [
|
||||
(self.eviction_strategy.get_priority(node), node) for node in leaves
|
||||
]
|
||||
@@ -1232,17 +1332,34 @@ class HiRadixCache(RadixCache):
|
||||
|
||||
# load it all or not at all
|
||||
host_hit_len = sum(self._node_host_len(n) for n in nodes_to_load)
|
||||
logger.info(
|
||||
"[HiCache-load] load_back CP: node_id=%d nodes_to_load=%d host_hit_len=%d threshold=%d",
|
||||
last_hit_node.id,
|
||||
len(nodes_to_load),
|
||||
host_hit_len,
|
||||
self.load_back_threshold,
|
||||
)
|
||||
if host_hit_len < self.load_back_threshold or (
|
||||
host_hit_len > mem_quota + delta if mem_quota is not None else False
|
||||
):
|
||||
# skip loading back if the total size is too small or exceeding the memory quota
|
||||
self.dec_lock_ref(ancester_node)
|
||||
logger.info(
|
||||
"[HiCache-load] load_back CP SKIP: host_hit_len=%d below threshold=%d or over quota",
|
||||
host_hit_len,
|
||||
self.load_back_threshold,
|
||||
)
|
||||
return None
|
||||
|
||||
device_indices = self.cache_controller.load_cp(
|
||||
nodes_to_load, node_id=last_hit_node.id
|
||||
)
|
||||
if device_indices is None:
|
||||
logger.info(
|
||||
"[HiCache-load] load_back CP retry with eviction: node_id=%d tokens_needed=%d",
|
||||
last_hit_node.id,
|
||||
host_hit_len,
|
||||
)
|
||||
self.evict(EvictParams(num_tokens=host_hit_len))
|
||||
device_indices = self.cache_controller.load_cp(
|
||||
nodes_to_load, node_id=last_hit_node.id
|
||||
@@ -1276,6 +1393,11 @@ class HiRadixCache(RadixCache):
|
||||
len(device_indices)
|
||||
)
|
||||
|
||||
logger.info(
|
||||
"[HiCache-load] load_back CP SUCCESS: node_id=%d loaded_tokens=%d",
|
||||
last_hit_node.id,
|
||||
len(device_indices),
|
||||
)
|
||||
return device_indices
|
||||
|
||||
start_time = time.perf_counter()
|
||||
@@ -1296,17 +1418,34 @@ class HiRadixCache(RadixCache):
|
||||
|
||||
# load it all or not at all
|
||||
host_indices = torch.cat([n.host_value for n in nodes_to_load])
|
||||
logger.info(
|
||||
"[HiCache-load] load_back non-CP: node_id=%d nodes_to_load=%d host_hit_len=%d threshold=%d",
|
||||
last_hit_node.id,
|
||||
len(nodes_to_load),
|
||||
len(host_indices),
|
||||
self.load_back_threshold,
|
||||
)
|
||||
if len(host_indices) < self.load_back_threshold or (
|
||||
len(host_indices) > mem_quota + delta if mem_quota is not None else False
|
||||
):
|
||||
# skip loading back if the total size is too small or exceeding the memory quota
|
||||
self.dec_lock_ref(ancester_node)
|
||||
logger.info(
|
||||
"[HiCache-load] load_back non-CP SKIP: host_hit_len=%d below threshold=%d or over quota",
|
||||
len(host_indices),
|
||||
self.load_back_threshold,
|
||||
)
|
||||
return None
|
||||
|
||||
device_indices = self.cache_controller.load(
|
||||
host_indices=host_indices, node_id=last_hit_node.id
|
||||
)
|
||||
if device_indices is None:
|
||||
logger.info(
|
||||
"[HiCache-load] load_back non-CP retry with eviction: node_id=%d tokens_needed=%d",
|
||||
last_hit_node.id,
|
||||
len(host_indices),
|
||||
)
|
||||
self.evict(EvictParams(num_tokens=len(host_indices)))
|
||||
device_indices = self.cache_controller.load(
|
||||
host_indices=host_indices, node_id=last_hit_node.id
|
||||
@@ -1337,6 +1476,12 @@ class HiRadixCache(RadixCache):
|
||||
)
|
||||
self.metrics_collector.increment_load_back_num_tokens(len(device_indices))
|
||||
|
||||
logger.info(
|
||||
"[HiCache-load] load_back non-CP SUCCESS: node_id=%d loaded_tokens=%d",
|
||||
last_hit_node.id,
|
||||
len(device_indices),
|
||||
)
|
||||
|
||||
return device_indices
|
||||
|
||||
def init_load_back(
|
||||
@@ -1348,7 +1493,7 @@ class HiRadixCache(RadixCache):
|
||||
if last_node.evicted:
|
||||
loading_values = self.load_back(last_node, mem_quota)
|
||||
if loading_values is not None:
|
||||
logger.debug(
|
||||
logger.info(
|
||||
f"loading back {len(loading_values)} tokens for node {last_node.id}"
|
||||
)
|
||||
return loading_values, last_node
|
||||
@@ -1369,7 +1514,15 @@ class HiRadixCache(RadixCache):
|
||||
return self.cache_controller.start_loading()
|
||||
|
||||
def flush_write_through_acks(self) -> None:
|
||||
ongoing_before = len(self.ongoing_write_through)
|
||||
self.writing_check()
|
||||
ongoing_after = len(self.ongoing_write_through)
|
||||
if ongoing_before > 0:
|
||||
logger.info(
|
||||
"[HiCache-write] flush_write_through_acks: before=%d after=%d",
|
||||
ongoing_before,
|
||||
ongoing_after,
|
||||
)
|
||||
|
||||
def check_hicache_events(self):
|
||||
self.writing_check()
|
||||
@@ -1478,7 +1631,7 @@ class HiRadixCache(RadixCache):
|
||||
completed_tokens, hash_value = self.cache_controller.terminate_prefetch(
|
||||
operation
|
||||
)
|
||||
logger.debug(f"Prefetch {req_id} completed with {completed_tokens} tokens")
|
||||
logger.info(f"Prefetch {req_id} completed with {completed_tokens} tokens")
|
||||
|
||||
min_completed_tokens = completed_tokens
|
||||
if self.tp_world_size > 1:
|
||||
|
||||
Reference in New Issue
Block a user