diff --git a/python/sglang/srt/managers/cache_controller.py b/python/sglang/srt/managers/cache_controller.py index c3b5df40b..1b9cb50c9 100644 --- a/python/sglang/srt/managers/cache_controller.py +++ b/python/sglang/srt/managers/cache_controller.py @@ -23,7 +23,7 @@ from typing import TYPE_CHECKING, Dict, List, NamedTuple, Optional, Set import torch from sglang.srt.environ import envs -from sglang.srt.mem_cache.cp_hicache_trace import cptrace, khash, rng, trace_enabled +from sglang.srt.mem_cache.cp_hicache_trace import cptrace, khash, knz, rng, trace_enabled from sglang.srt.mem_cache.hicache_storage import HiCacheStorageConfig if TYPE_CHECKING: @@ -1497,12 +1497,15 @@ class HiCacheController: if trace_enabled(2): for state in target_transfer_states: kv = self.mem_pool_device.get_key_buffer(layer_id) + rows = kv[state.physical_device_indices] cptrace( 2, "backup_kv_hash", node_id=state.reservation.node_id, layer=layer_id, - kvhash=khash(kv[state.physical_device_indices]), + host=rng(state.host_indices), + kvhash=khash(rows), + nz=knz(rows), nrows=int(state.physical_device_indices.numel()), ) @@ -1975,12 +1978,15 @@ class HiCacheController: # op's nodes (single node in the c=1 rehit pass). if trace_enabled(2): kv = self.mem_pool_device.get_key_buffer(i) + rows = kv[device_indices] cptrace( 2, "reload_kv_hash", node_ids=node_ids, layer=i, - kvhash=khash(kv[device_indices]), + host=rng(host_indices), + kvhash=khash(rows), + nz=knz(rows), nrows=int(device_indices.numel()), ) producer_event.complete(i) diff --git a/python/sglang/srt/mem_cache/cp_hicache_trace.py b/python/sglang/srt/mem_cache/cp_hicache_trace.py index e24af4e92..e588c0a1e 100644 --- a/python/sglang/srt/mem_cache/cp_hicache_trace.py +++ b/python/sglang/srt/mem_cache/cp_hicache_trace.py @@ -79,6 +79,22 @@ def khash(t) -> int: return -1 +def knz(t) -> int: + """Count of nonzero bytes in a KV slice. Real KV is never all-zero, so a + reload (or backup) slice with nz==0 = uninitialized/zero KV -- a direct, + self-contained corruption flag (no cross-stage matching needed). Catches the + observed '0|0|0...' garbage where the store never filled the backed-up pages.""" + try: + import torch + + if t is None or t.numel() == 0: + return 0 + b = t.detach().contiguous().view(torch.uint8).reshape(-1) + return int((b != 0).sum().item()) + except Exception: + return -1 + + def cptrace(level: int, tag: str, **fields) -> None: """Emit one ``[CPTRACE ] k=v ...`` line if the gate >= level.