diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 28708d39a..161dc36ce 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -271,6 +271,9 @@ class Envs: SGLANG_CP_SHARED_KV_MLA_PREFETCH_MIN_EXTEND_TOKENS = EnvInt(-1) SGLANG_CP_DRAFT_SHARED_KV = EnvBool(False) SGLANG_CP_DRAFT_SHARED_KV_DEBUG = EnvBool(False) + # CP HiCache round-trip KV-corruption tracing. 0=off, 1=structural lifecycle + # (rid_map/backup/split/evict/reload), 2=+compose/free/ack timing. + SGLANG_CP_HICACHE_KV_TRACE = EnvInt(0) SGLANG_EAGLE_ACCEPT_DEBUG = EnvBool(False) SGLANG_EAGLE_ACCEPT_DEBUG_INTERVAL = EnvInt(128) SGLANG_DISABLE_TAI_BIGRAM = EnvBool(False) diff --git a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py index c368b3e75..44f8eb9ae 100644 --- a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py +++ b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py @@ -9,6 +9,9 @@ from typing import Any import torch from sglang.srt.environ import envs +from sglang.srt.mem_cache.cp_hicache_trace import cptrace as _cptrace +from sglang.srt.mem_cache.cp_hicache_trace import rng as _cprng +from sglang.srt.mem_cache.cp_hicache_trace import trace_enabled as _cptrace_enabled from sglang.srt.layers.attention.nsa.utils import ( cp_shared_kv_bs_gt1_timing_start, get_cp_shared_kv_local_out_cache_loc, @@ -2071,6 +2074,22 @@ def fill_current_index_page_slots( safe_req = torch.clamp(req_long, min=0, max=max(batch_rows - 1, 0)) dense_pages = page_inverse[safe_req, safe_pages.to(torch.long)].to(torch.long) valid_rows = valid_pages & (dense_pages > 0) + if _cptrace_enabled(2): + # H3: in-range logical locs that fail to map to a real dense page + # (dense<=0) right after a reload = stale/wrong page_inverse. + _in_range = int(valid_pages.sum().item()) + _mapped = int(valid_rows.sum().item()) + _cptrace( + 2, + "remap", + total=int(current_locs.numel()), + in_range=_in_range, + mapped=_mapped, + unmapped=_in_range - _mapped, + batch_rows=int(batch_rows), + capacity=int(capacity), + dense=_cprng(dense_pages), + ) if not torch.any(valid_rows): return dense_page_buffer @@ -4426,10 +4445,28 @@ def get_cp_shared_kv_flattened_request_row_ids( cached = getattr(forward_batch, "cp_flattened_row_ids", None) cached_key = getattr(forward_batch, "cp_flattened_row_ids_key", None) if cached is not None and cached_key == key and cached.device == device: + if _cptrace_enabled(2): + _cptrace( + 2, + "compose", + fb=id(forward_batch), + key=key, + hit=True, + row_ids=_cprng(cached), + ) return cached row_ids = build_flattened_request_row_ids(seq_lens_cpu, device=device) forward_batch.cp_flattened_row_ids = row_ids forward_batch.cp_flattened_row_ids_key = key + if _cptrace_enabled(2): + _cptrace( + 2, + "compose", + fb=id(forward_batch), + key=key, + hit=False, + row_ids=_cprng(row_ids), + ) return row_ids diff --git a/python/sglang/srt/managers/cache_controller.py b/python/sglang/srt/managers/cache_controller.py index d23dc6378..87e5edbc6 100644 --- a/python/sglang/srt/managers/cache_controller.py +++ b/python/sglang/srt/managers/cache_controller.py @@ -23,6 +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, rng, trace_enabled from sglang.srt.mem_cache.hicache_storage import HiCacheStorageConfig if TYPE_CHECKING: @@ -1141,6 +1142,22 @@ class HiCacheController: self.draft_mem_pool_host.free(draft_host_indices) raise + if trace_enabled(1): + cptrace( + 1, + "backup_reserve", + node_id=node_id, + logical_len=logical_len, + padded_len=padded_len, + owned=owned_positions.numel(), + host_idx=rng(host_indices), + phys=rng(physical_device_indices), + page_owners=page_owners.numel() + if page_owners is not None + else 0, + draft=self.has_draft_hicache, + draft_host_idx=rng(draft_host_indices), + ) return HiCacheWriteReservation( metadata=CpHiCacheNodeMetadata( logical_len=logical_len, @@ -1401,6 +1418,20 @@ class HiCacheController: self.io_backend, ) grouped_tensors.extend([target_host_indices, target_device_indices]) + if trace_enabled(1) and layer_id == 0: + cptrace( + 1, + "backup_d2h", + node_ids=[s.reservation.node_id for s in target_transfer_states], + layer=layer_id, + host_idx=rng(target_host_indices), + phys=rng(target_device_indices), + len_match=( + target_host_indices.numel() + == target_device_indices.numel() + ), + draft=False, + ) if draft_transfer_states: draft_host_indices = self._concat_layer_write_tensors( @@ -1420,6 +1451,20 @@ class HiCacheController: self.io_backend, ) grouped_tensors.extend([draft_host_indices, draft_device_indices]) + if trace_enabled(1) and layer_id == 0: + cptrace( + 1, + "backup_d2h", + node_ids=[s.reservation.node_id for s in draft_transfer_states], + layer=layer_id, + host_idx=rng(draft_host_indices), + phys=rng(draft_device_indices), + len_match=( + draft_host_indices.numel() + == draft_device_indices.numel() + ), + draft=True, + ) for tensor in grouped_tensors: self._record_tensor_on_stream(tensor, self.write_stream) @@ -1675,6 +1720,19 @@ class HiCacheController: host_chunks.append(meta.host_indices) if draft_host_indices is not None: draft_host_chunks.append(draft_host_indices) + if trace_enabled(1): + cptrace( + 1, + "reload_node", + node_id=getattr(node, "id", -1), + valid_len=valid_len, + padded_len=padded_len, + owned=owned_positions.numel(), + host_idx=rng(meta.host_indices), + sel_logical=rng(selected_logical_locs), + phys=rng(physical_chunks[-1]), + draft_host_idx=rng(draft_host_indices), + ) record_stage("build_chunks") visible_device_indices = ( diff --git a/python/sglang/srt/mem_cache/allocator.py b/python/sglang/srt/mem_cache/allocator.py index bf2cceff8..1cdd7cd2b 100644 --- a/python/sglang/srt/mem_cache/allocator.py +++ b/python/sglang/srt/mem_cache/allocator.py @@ -27,6 +27,7 @@ import torch import triton from sglang.srt.environ import envs +from sglang.srt.mem_cache.cp_hicache_trace import cptrace, rng, trace_enabled import triton.language as tl from sglang.srt.utils import get_bool_env_var, get_num_new_pages, next_power_of_2 @@ -847,6 +848,14 @@ class CPSharedPagedTokenToKVPoolAllocator(PagedTokenToKVPoolAllocator): if free_index.numel() == 0: return + if trace_enabled(2): + cptrace( + 2, + "dev_free", + tokens=rng(free_index), + pages=rng(torch.unique(free_index // self.page_size)), + in_free_group=not self.is_not_in_free_group, + ) if self.is_not_in_free_group: free_page_indices = torch.unique(free_index // self.page_size) self._append_pages_to_owner_buckets( diff --git a/python/sglang/srt/mem_cache/cp_hicache_trace.py b/python/sglang/srt/mem_cache/cp_hicache_trace.py new file mode 100644 index 000000000..21338b07b --- /dev/null +++ b/python/sglang/srt/mem_cache/cp_hicache_trace.py @@ -0,0 +1,71 @@ +"""Env-gated structured tracing for CP HiCache KV round-trip corruption debugging. + +All output is keyed by ``node_id`` (which threads prepare->backup->evict->reload) +plus a ``rid_map`` line that ties the client request id to its node_id. The +companion analyzer joins these lines by node_id and checks four invariants: + H1 split-misattribution, H2 stale/wrong host ref on reload, + H3 stale compose cache, H4 page-recycle-before-backup race. + +Gate: SGLANG_CP_HICACHE_KV_TRACE (0=off, 1=structural lifecycle, 2=+compose/free/ack). +Off by default and cheap to leave in place: when the level is below the call's +level, nothing is computed or logged. +""" +import logging + +from sglang.srt.environ import envs + +logger = logging.getLogger(__name__) + + +def trace_level() -> int: + try: + return int(envs.SGLANG_CP_HICACHE_KV_TRACE.get()) + except Exception: + return 0 + + +def trace_enabled(level: int = 1) -> bool: + return trace_level() >= level + + +def rng(t) -> str: + """Compact summary of a 1-D index tensor/list: (cnt,min,max,first,last). + + Debug-only; .item() syncs are acceptable here. Empty -> (0).""" + if t is None: + return "(none)" + try: + n = t.numel() if hasattr(t, "numel") else len(t) + if n == 0: + return "(0)" + if hasattr(t, "numel"): + tf = t.flatten() + return ( + f"(cnt={n},min={int(tf.min().item())},max={int(tf.max().item())}," + f"first={int(tf[0].item())},last={int(tf[-1].item())})" + ) + return f"(cnt={n},min={min(t)},max={max(t)},first={t[0]},last={t[-1]})" + except Exception as e: # never let tracing crash the hot path + return f"(err:{e})" + + +def _fmt(v) -> str: + """Render a value with NO internal spaces so the log is space-splittable + into key=value pairs by the analyzer.""" + if isinstance(v, (list, tuple)): + return "[" + ",".join(_fmt(x) for x in v) + "]" + return str(v).replace(" ", "") + + +def cptrace(level: int, tag: str, **fields) -> None: + """Emit one ``[CPTRACE ] k=v ...`` line if the gate >= level. + + Every value is rendered space-free so ``line.split()`` recovers the + key=value pairs cleanly in the offline analyzer.""" + if trace_level() < level: + return + try: + parts = " ".join(f"{k}={_fmt(v)}" for k, v in fields.items()) + logger.info("[CPTRACE %s] %s", tag, parts) + except Exception: + pass diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index 48aa66219..464b6f13c 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -15,6 +15,7 @@ from typing import TYPE_CHECKING, Dict, List, Optional, Set, Tuple import torch from sglang.srt.environ import envs +from sglang.srt.mem_cache.cp_hicache_trace import cptrace, rng, trace_enabled from sglang.srt.managers.cache_controller import ( HiCacheController, HiCacheWriteFailure, @@ -2685,6 +2686,20 @@ class HiRadixCache(RadixCache): len(kv_indices), result.metadata.owned_positions.numel(), ) + if trace_enabled(1): + md = result.metadata + cptrace( + 1, + "rid_map", + rid=getattr(req, "rid", ""), + node_id=node_id, + logical_len=len(kv_indices), + owned=md.owned_positions.numel(), + host_idx=rng(getattr(md, "host_indices", None)), + page_owners=getattr(md, "page_owners", None).numel() + if getattr(md, "page_owners", None) is not None + else 0, + ) def prepare_write_backup_for_req(self, req) -> None: if self.disable or not self._uses_cp_hicache: @@ -3000,6 +3015,13 @@ class HiRadixCache(RadixCache): ) finish_count = int(queue_size.item()) + if trace_enabled(2) and finish_count > 0: + finished_ids = [ + aid + for entry in self.cache_controller.ack_write_queue[:finish_count] + for aid in entry[2] + ] + cptrace(2, "write_ack", finished=finish_count, node_ids=finished_ids) logger.debug( "[HiCache-write] writing_check: ongoing=%d ack_queue=%d local_finished=%d sync_finished=%d tp_size=%d", len(self.ongoing_write_through), @@ -3328,6 +3350,8 @@ class HiRadixCache(RadixCache): def _evict_backuped(self, node: TreeNode): # GPU -> CPU demotion: no BlockRemoved since block is still reachable via load_back + _trace_wp = self._node_host_write_pending(node) if trace_enabled(1) else None + _trace_dev = rng(node.value) if trace_enabled(1) else None device_resident_len = self._node_device_resident_len(node) freed_len = self.cache_controller.evict_device(node.value) assert freed_len > 0 @@ -3340,6 +3364,17 @@ class HiRadixCache(RadixCache): node.lock_ref, self._node_backuped(node), ) + if trace_enabled(1): + cptrace( + 1, + "evict", + node_id=node.id, + freed_len=freed_len, + device_resident=device_resident_len, + backed=self._node_backuped(node), + write_pending=_trace_wp, + dev=_trace_dev, + ) node.value = None self._update_leaf_status(node) self._update_host_leaf_status(node) @@ -3658,6 +3693,15 @@ class HiRadixCache(RadixCache): for loaded_node in nodes_to_load: host_len = self._node_host_len(loaded_node) loaded_node.value = device_indices[offset : offset + host_len].clone() + if trace_enabled(1): + cptrace( + 1, + "reload_assign", + node_id=loaded_node.id, + host_len=host_len, + offset=offset, + value=rng(loaded_node.value), + ) offset += host_len metadata = getattr(loaded_node, "cp_hicache", None) physical_loaded_len += int( @@ -4247,11 +4291,43 @@ class HiRadixCache(RadixCache): child.value = child.value[split_len:].clone() if self._uses_cp_hicache: if self._node_backuped(child): + if trace_enabled(1): + _orig = child.cp_hicache + _orig_owned = _orig.owned_positions.numel() + _orig_host = rng(getattr(_orig, "host_indices", None)) + _orig_po = ( + _orig.page_owners.numel() + if getattr(_orig, "page_owners", None) is not None + else 0 + ) new_node.cp_hicache, child.cp_hicache = child.cp_hicache.split( split_len ) new_node.host_len = split_len child.host_len = child.host_len - split_len + if trace_enabled(1): + pf, sf = new_node.cp_hicache, child.cp_hicache + cptrace( + 1, + "split", + orig_id=child.id, + prefix_id=new_node.id, + suffix_id=child.id, + split_len=split_len, + orig_owned=_orig_owned, + orig_host=_orig_host, + orig_po=_orig_po, + pre_owned=pf.owned_positions.numel(), + pre_host=rng(getattr(pf, "host_indices", None)), + pre_po=pf.page_owners.numel() + if getattr(pf, "page_owners", None) is not None + else 0, + suf_owned=sf.owned_positions.numel(), + suf_host=rng(getattr(sf, "host_indices", None)), + suf_po=sf.page_owners.numel() + if getattr(sf, "page_owners", None) is not None + else 0, + ) elif child.backuped: new_node.host_value = child.host_value[:split_len].clone() child.host_value = child.host_value[split_len:].clone()