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:
2026-05-11 00:00:04 +08:00
parent 3a5928ef53
commit 29a395fe95
3 changed files with 230 additions and 7 deletions

View File

@@ -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)

View File

@@ -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
)

View File

@@ -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: