From 820790b97fba106f1d69855240787a32f5b1cdb9 Mon Sep 17 00:00:00 2001 From: leavelet Date: Wed, 17 Jun 2026 19:59:01 +0000 Subject: [PATCH] Remove CP HiCache debug instrumentation The CP shared-KV cache-hit corruption bug is fixed and verified end-to-end, so drop the per-layer/per-request trace probes added during the investigation. Each probe was a Python call plus an UNCACHED os.getenv (EnvField.get() re-reads os.environ on every call) in the hot path even when disabled -- ~1000 per forward from the 11 per-layer deepseek probes alone -- plus a latent risk that an accidentally-set SGLANG_NSA_DUMP_DIR would dump tensors in-forward and tank throughput. Delete cp_hicache_trace.py and every reference to it: - deepseek_v2: 11 per-layer fwd_hash probes + the final dump-flush block - nsa_backend: the per-compose NSA tensor-dump block - cp_shared_kv_runtime: 5 imports + the if _cptrace_enabled(...) compose blocks - cache_controller / hiradix_cache / allocator / memory_pool_host: the cptrace/khash/knz/rng round-trip + lifecycle hashes - environ: the SGLANG_CP_HICACHE_KV_TRACE and SGLANG_NSA_DUMP_DIR flags Kept: SGLANG_DEBUG_CP_SHARED_KV / SGLANG_CP_TRANSFER_LOG transfer logging and the x-request-id passthrough (independent of the trace module). Verified: imports clean, 235 mem_cache + reasoning unit tests pass. Co-Authored-By: Claude Opus 4.8 (1M context) --- python/sglang/srt/environ.py | 6 - .../attention/nsa/cp_shared_kv_runtime.py | 96 ------ .../srt/layers/attention/nsa_backend.py | 17 - .../sglang/srt/managers/cache_controller.py | 137 -------- python/sglang/srt/mem_cache/allocator.py | 23 -- .../sglang/srt/mem_cache/cp_hicache_trace.py | 295 ------------------ python/sglang/srt/mem_cache/hiradix_cache.py | 96 ------ .../sglang/srt/mem_cache/memory_pool_host.py | 27 -- python/sglang/srt/models/deepseek_v2.py | 24 -- 9 files changed, 721 deletions(-) delete mode 100644 python/sglang/srt/mem_cache/cp_hicache_trace.py diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index d0f3c82a1..12c42a3ca 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -271,16 +271,10 @@ 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) # Logging-only: emit the CP shared-KV sender/worker transfer-partition dumps # (main-KV pages/positions vs NSA-state pages/positions) WITHOUT the side effects # of SGLANG_DEBUG_CP_SHARED_KV (which disables tai materialize -> fail-fast). SGLANG_CP_TRANSFER_LOG = EnvBool(False) - # Dir to torch.save per-stage attention tensors for L1-hit vs L2-reload diff. - # Empty=off. Only requests whose rid starts with "dump-" are dumped. - SGLANG_NSA_DUMP_DIR = EnvStr("") 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 626d2d41e..ffc707934 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,11 +9,6 @@ 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 khash as _khash -from sglang.srt.mem_cache.cp_hicache_trace import knz as _knz -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, @@ -2076,22 +2071,6 @@ 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 @@ -2242,21 +2221,6 @@ def build_cp_shared_kv_ipc_page_descriptors( invalid_value = torch.full_like(owner_ranks, -1) owner_ranks = torch.where(invalid, invalid_value, owner_ranks) src_page_indices = torch.where(invalid, invalid_value, src_page_indices) - if _cptrace_enabled(3): - try: - _valid = ~invalid - _cptrace( - 3, - "ipc_desc", - cprank=layout.cp_rank, - slots=int(logical_pages.numel()), - valid=int(_valid.sum().item()), - owners=_cprng(owner_ranks[_valid]), - src=_cprng(src_page_indices[_valid]), - lpg=_cprng(logical_pages[_valid]), - ) - except Exception: - pass return owner_ranks.contiguous(), src_page_indices.contiguous() @@ -4462,28 +4426,10 @@ 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 @@ -5662,48 +5608,6 @@ def materialize_prefix_and_reuse_current_kv_page_slots( materialized_by_ipc, kv_cache.dtype, ) - if _cptrace_enabled(3): - try: - _live = mixed_locs.reshape(-1) - _live = _live[_live >= 0] - _gath = ( - mixed_kv_cache.index_select(0, _live.to(torch.long)) - if _live.numel() - else mixed_kv_cache[:0] - ) - _cptrace( - 3, - "gather_out", - cprank=layout.cp_rank, - layer=layer_id, - total_slots=int(total_slots), - live_rows=int(_live.numel()), - prefix_pages=sum(int(e) - int(s) for s, e in prefix_spans), - current_pages=sum(int(e) - int(s) for s, e in merged_current_spans), - h=_khash(_gath), - nz=_knz(_gath), - ) - _slp = slot_remap.slot_logical_pages.reshape(-1) - for _kind, _spans in (("prefix", prefix_spans), ("current", merged_current_spans)): - for _s, _e in _spans: - _sl = slot_range_to_token_slice(page_size, int(_s), int(_e)) - _lp = _slp[int(_s) : int(_e)].to(torch.long) - _cptrace( - 3, - "span_owner", - cprank=layout.cp_rank, - layer=layer_id, - kind=_kind, - s=int(_s), - e=int(_e), - owners=_cprng(layout.owner_for_logical_pages(_lp)), - phys=_cprng(layout.logical_pages_to_physical(_lp)), - lpg=_cprng(_lp), - h=_khash(mixed_kv_cache[_sl]), - nz=_knz(mixed_kv_cache[_sl]), - ) - except Exception: - pass return mixed_kv_cache, mixed_locs diff --git a/python/sglang/srt/layers/attention/nsa_backend.py b/python/sglang/srt/layers/attention/nsa_backend.py index 93addc554..4faf58c8f 100644 --- a/python/sglang/srt/layers/attention/nsa_backend.py +++ b/python/sglang/srt/layers/attention/nsa_backend.py @@ -2696,23 +2696,6 @@ class NativeSparseAttnBackend( sm_scale=layer.scaling, v_head_dim=layer.v_head_dim, ) - # EXACT L1-hit vs L2-reload dump (env+rid gated): the query, the - # composed/dequantized KV the kernel reads, the selection, and the - # output. Offline relerr per tensor localizes reload corruption past - # the fp-nondeterminism that defeats binary hashes. - try: - from sglang.srt.mem_cache.cp_hicache_trace import ( - dump_tensors as _cp_dump, - ) - - _cp_dump( - forward_batch, - layer.layer_id, - "nsa", - {"q_all": q_all}, - ) - except Exception: - pass elif nsa_impl == "flashmla_kv": if ( self.nsa_kv_cache_store_fp8 diff --git a/python/sglang/srt/managers/cache_controller.py b/python/sglang/srt/managers/cache_controller.py index 6fdb0ac32..2ae5d4a53 100644 --- a/python/sglang/srt/managers/cache_controller.py +++ b/python/sglang/srt/managers/cache_controller.py @@ -23,7 +23,6 @@ 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, knz, rng, trace_enabled from sglang.srt.mem_cache.hicache_storage import HiCacheStorageConfig if TYPE_CHECKING: @@ -1142,22 +1141,6 @@ 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, @@ -1418,21 +1401,6 @@ 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( [state.draft_host_indices for state in draft_transfer_states] @@ -1451,21 +1419,6 @@ 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) @@ -1494,21 +1447,6 @@ class HiCacheController: # correct post-store value -- NOT the racy write_stream copy's view. If a # per-layer store/copy ordering race exists, the host copy will diverge # from THIS hash, and the reload hash (post-H2D) will not match it. - 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, - host=rng(state.host_indices), - kvhash=khash(rows), - nz=knz(rows), - nrows=int(state.physical_device_indices.numel()), - ) - for state in final_states: self._append_layer_write_ack(state) @@ -1705,26 +1643,6 @@ class HiCacheController: node_device_indices = device_indices[offset : offset + padded_len] offset += padded_len visible_chunks.append(node_device_indices[:valid_len]) - if trace_enabled(1): - # CROSS-RANK invariant: the visible logical-loc sequence for this - # node MUST be byte-identical on every CP rank — the cross-rank - # gather reads token j from its owner rank at physical(loc_j), so - # if rank A and rank R disagree on loc_j, A reads R's slot for a - # DIFFERENT token (right bytes, wrong place). Emitted for EVERY - # node (incl. zero-owned, before the continue below) so the - # analyzer can diff this hash across ranks by node_id. Divergence - # here is produced by wall-clock-LRU device eviction giving ranks - # different free buckets in alloc_pages_with_owners. - _vis = node_device_indices[:valid_len] - cptrace( - 1, - "visible_locs_hash", - node_id=getattr(node, "id", -1), - cprank=self.cp_shared_kv_layout.cp_rank, - n=int(_vis.numel()), - locs_hash=khash(_vis), - locs=rng(_vis), - ) if self.has_draft_hicache: draft_host_indices = getattr( meta, "draft_host_indices", None @@ -1760,48 +1678,6 @@ 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), - ) - # MAKE-OR-BREAK reload invariant (the code above SKIPS this for - # perf): every page this rank reloads into must be owned by THIS - # rank, i.e. owner_for_logical_pages((loc)//page_size) == cp_rank. - # A nonzero bad_owner = alloc_pages_with_owners/padding produced a - # per-page owner permutation -> this rank loads its host bytes into - # pages owned by another rank -> attention reads garbage (the - # aggregate byte-hash can't see this). - try: - _lay = self.cp_shared_kv_layout - _ps = int(getattr(_lay, "page_size", getattr(self, "page_size", 0))) - if _ps > 0: - _pages = selected_logical_locs // _ps - _own = _lay.owner_for_logical_pages(_pages) - _bad = int((_own != _lay.cp_rank).sum().item()) - cptrace( - 1, - "owner_check", - node_id=getattr(node, "id", -1), - cprank=_lay.cp_rank, - owned=int(owned_positions.numel()), - bad_owner=_bad, - valid_len=valid_len, - padded_len=padded_len, - tail_pad=padded_len - valid_len, - pages=rng(_pages), - owners=rng(_own), - ) - except Exception: - pass record_stage("build_chunks") visible_device_indices = ( @@ -2025,19 +1901,6 @@ class HiCacheController: # Compare to backup_kv_hash(node,layer): inequality = # round-trip delivered wrong KV. node_ids is the merged # 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, - host=rng(host_indices), - kvhash=khash(rows), - nz=knz(rows), - nrows=int(device_indices.numel()), - ) producer_event.complete(i) elif op is None and i < self.layer_num: producer_event.complete(i) diff --git a/python/sglang/srt/mem_cache/allocator.py b/python/sglang/srt/mem_cache/allocator.py index b1fd6ea6e..bf2cceff8 100644 --- a/python/sglang/srt/mem_cache/allocator.py +++ b/python/sglang/srt/mem_cache/allocator.py @@ -27,7 +27,6 @@ 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 @@ -848,14 +847,6 @@ 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( @@ -1093,20 +1084,6 @@ class CPSharedPagedTokenToKVPoolAllocator(PagedTokenToKVPoolAllocator): page_size, dtype=torch.int64, device=self.device ).unsqueeze(0) out_indices = (base + offsets).reshape(-1) - if trace_enabled(1) and any(c > 0 for c in selected_release_counts): - # Tapped the DEFERRED-FREE (release) bucket because a lane's free - # bucket was exhausted — only under eviction pressure. A release page - # may still be read by an in-flight kernel (deferred free exists to - # prevent exactly this), so handing it to an H2D reload here is a - # use-after-free / aliasing risk the byte-round-trip hash can't see. - cptrace( - 1, - "release_draw", - cprank=self.cp_rank, - pages=len(page_compute_owners), - free_by_owner=selected_free_counts, - release_by_owner=selected_release_counts, - ) self._consume_owner_bucket_prefix( release=False, counts_by_owner=selected_free_counts ) diff --git a/python/sglang/srt/mem_cache/cp_hicache_trace.py b/python/sglang/srt/mem_cache/cp_hicache_trace.py deleted file mode 100644 index 085931115..000000000 --- a/python/sglang/srt/mem_cache/cp_hicache_trace.py +++ /dev/null @@ -1,295 +0,0 @@ -"""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 khash(t) -> int: - """Full-tensor, position-weighted int64 fingerprint of a KV slice. - - Sensitive to BOTH value and position (catches permutation/scatter, unlike a - plain sum), covers the whole tensor (not a byte sample), and reduces to one - int. Debug-only: the .item() sync is acceptable when tracing is on. Used to - compare the KV bytes a node's pages hold at backup (correct, post-store, on - the default stream) vs what comes back at reload (post-H2D); inequality = - the round-trip delivered wrong KV (store-vs-copy race or transfer corruption).""" - try: - import torch - - if t is None or t.numel() == 0: - return 0 - b = t.detach().contiguous().view(torch.uint8).reshape(-1).to(torch.int64) - n = b.numel() - idx = torch.arange(1, n + 1, device=b.device, dtype=torch.int64) - return int((((b * idx).sum() ^ (b.sum() * 1000003) ^ n) & 0x7FFFFFFFFFFFFFFF).item()) - except Exception: - 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 _content_key(forward_batch) -> int: - """Stable per-request content fingerprint so the SAME request forwarded from - an L1-hit (known-good) and an L2-reload (suspect) can be JOINED across the - log WITHOUT rid (the Rust PD gateway strips the client rid -> server mints a - UUID). Derived from the extend input-ids + total seq length, which are - identical for the same content regardless of where the prefix KV came from. - Cached on the ForwardBatch (one compute, shared by every layer/stage). - Only meaningful for a SINGLE-request (bs=1) forward; for bs>1 it mixes - requests and simply won't join (harmless).""" - ck = getattr(forward_batch, "_cp_content_key", None) - if ck is not None: - return ck - ck = -1 - try: - import torch - - ids = getattr(forward_batch, "input_ids", None) - sl = getattr(forward_batch, "seq_lens", None) - h = khash(ids) if isinstance(ids, torch.Tensor) else 0 - slsum = int(sl.sum().item()) if isinstance(sl, torch.Tensor) else 0 - nreq = len(getattr(forward_batch, "rids", []) or []) - ck = ((h ^ (slsum * 1000003) ^ (nreq * 998244353)) & 0x7FFFFFFFFFFFFFFF) - except Exception: - ck = -1 - try: - forward_batch._cp_content_key = ck - except Exception: - pass - return ck - - -def fwd_hash(forward_batch, layer_id, stage, t, *, level: int = 3) -> None: - """Hash a per-layer forward tensor (attn-in/out, topk_indices, MoE-in) to - localize where a reload forward diverges from a fresh one. Level 3 (so the - level-2 KV run is unaffected). Guards: only the EAGER EXTEND path (the reload - case) -- never under cuda-graph decode (.item() sync would corrupt capture); - handles topk_indices=None and tuple/quant hidden_states.""" - # LEAN dump (gated internally by SGLANG_NSA_DUMP_DIR + rid 'dump-'): only the - # residual-stream trajectory (attn_in/attn_out per layer) — enough to localize - # the first divergent layer. (The full-stage/kv dump was too heavy and hung the - # server.) Stage 2 drills into the localized layer. - if stage in ("attn_in", "attn_out"): - try: - dump_tensors(forward_batch, layer_id, stage, {"v": t}) - except Exception: - pass - if trace_level() < level: - return - try: - import torch - - fm = getattr(forward_batch, "forward_mode", None) - if fm is not None and hasattr(fm, "is_extend") and not fm.is_extend(): - return - # Single-request (bs<=1) forwards only: the content-key join (L1-hit vs - # L2-reload of the SAME content) is meaningful only at bs==1 (bs>1 mixes - # requests -> ck is a blend), and this also skips the c=24 flood-evict - # forwards so the level-3 log stays small and focused on the comparison. - _rids0 = getattr(forward_batch, "rids", None) - if _rids0 is not None and len(_rids0) > 1: - return - if isinstance(t, torch.Tensor): - h, nz, rows = khash(t), knz(t), int(t.shape[0]) if t.dim() else 0 - elif t is None: - h, nz, rows = 0, 0, 0 - else: - h, nz, rows = -2, -2, -2 # tuple/quant: not a plain tensor - rids = getattr(forward_batch, "rids", None) - lay = getattr(forward_batch, "cp_shared_kv_layout", None) - cptrace( - level, - "fwd_hash", - rid=(rids[0] if rids else "?"), - nreq=(len(rids) if rids else 0), - ck=_content_key(forward_batch), - layer=layer_id, - stage=stage, - h=h, - nz=nz, - rows=rows, - cprank=(getattr(lay, "cp_rank", -1) if lay is not None else -1), - ) - except Exception: - pass - - -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 - - -def _dump_enabled_rid(forward_batch): - """Return (dump_dir, rid, cprank, mode) if this forward should dump, else None. - Two rid-controlled modes (gated by SGLANG_NSA_DUMP_DIR + EXTEND forward): - 'dumpfh-' -> mode 'fh' : final_hidden ONLY (tiny), works for FRESH prefills too - -> the cache-BUST-vs-cache-HIT discriminator. - 'dump-' -> mode 'full': lean per-layer trajectory (attn_in/out, q_all), and - only CACHE-HIT extends (skip huge fresh prefills).""" - try: - d = envs.SGLANG_NSA_DUMP_DIR.get() - except Exception: - d = "" - if not d: - return None - fm = getattr(forward_batch, "forward_mode", None) - if fm is not None and hasattr(fm, "is_extend") and not fm.is_extend(): - return None - rids = getattr(forward_batch, "rids", None) - rid = str(rids[0]) if rids else "" - if rid.startswith("dumpfh-"): - mode = "fh" - elif rid.startswith("dump-"): - mode = "full" - # 'full' only on cache-hit extends; fresh full prefills are huge + not the repro. - epl = getattr(forward_batch, "extend_prefix_lens_cpu", None) - try: - if epl is not None and int(epl.sum().item()) <= 0: - return None - except Exception: - pass - else: - return None - lay = getattr(forward_batch, "cp_shared_kv_layout", None) - cprank = int(getattr(lay, "cp_rank", -1)) if lay is not None else -1 - return d, rid, cprank, mode - - -def dump_tensors(forward_batch, layer_id, tag, tensors: dict, *, big_layers=None) -> None: - """ACCUMULATE per-(layer,stage) tensors into a per-forward buffer for a COMPLETE - L1-hit vs L2-reload diff. Written once per (rid,rank) by dump_flush() at model end - (token order, comparable; relerr offline beats fp-nondeterminism). 'big' tensors - (e.g. kv_cache, passed with big_layers=) are kept ONLY at those layers to - bound size; everything else at every layer. Gated via _dump_enabled_rid.""" - info = _dump_enabled_rid(forward_batch) - if info is None or info[3] != "full": - return # per-layer trajectory only in 'full' mode (cache-hit); 'fh' skips it - try: - import torch - - buf = getattr(forward_batch, "_cp_dump_buf", None) - if buf is None: - buf = {} - forward_batch._cp_dump_buf = buf - lbuf = buf.setdefault(int(layer_id), {}) - for k, v in tensors.items(): - if big_layers is not None and int(layer_id) not in big_layers: - continue - if isinstance(v, torch.Tensor): - if v.numel() > 200_000_000: - lbuf[f"{tag}/{k}_skipped_numel"] = int(v.numel()) - continue - lbuf[f"{tag}/{k}"] = v.detach().to("cpu") - except Exception: - pass - - -def dump_flush(forward_batch, *, final_hidden=None, positions=None) -> None: - """Write the accumulated per-layer dump buffer to ONE file per (rid,rank) at model - end, plus final hidden / positions / alignment metadata, then clear the buffer.""" - info = _dump_enabled_rid(forward_batch) - if info is None: - return - d, rid, cprank, mode = info - try: - import os - import torch - - buf = getattr(forward_batch, "_cp_dump_buf", None) - forward_batch._cp_dump_buf = None - payload = {"rid": rid, "cprank": cprank, "mode": mode} - if mode == "full": - payload["layers"] = buf or {} - if isinstance(final_hidden, torch.Tensor): - # 'fh' mode: keep only the LAST token's hidden (the first-token source) - # so a fresh full-prefill dump stays tiny; 'full' keeps the extend rows. - fh = final_hidden.detach() - if mode == "fh" and fh.dim() >= 1 and fh.shape[0] > 1: - fh = fh[-1:] - payload["final_hidden"] = fh.to("cpu") - if isinstance(positions, torch.Tensor): - payload["positions"] = positions.detach().to("cpu") - for name in ("seq_lens_cpu", "extend_prefix_lens_cpu", "extend_seq_lens_cpu"): - v = getattr(forward_batch, name, None) - if isinstance(v, torch.Tensor): - payload[name] = v.detach().to("cpu") - os.makedirs(d, exist_ok=True) - torch.save(payload, os.path.join(d, f"{rid}_cp{cprank}.pt")) - except Exception: - pass diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index c8be24dd2..dfd9a90df 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -15,7 +15,6 @@ 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, @@ -1310,25 +1309,6 @@ class HiRadixCache(RadixCache): ) -> CpLoadBackPlan: evict_start_time = time.perf_counter() eviction_plan = self._plan_cp_load_back_owner_lane_evictions(plan) - if trace_enabled(1): - # ROOT-CAUSE signal: device load-back eviction victim selection is - # keyed on wall-clock last_access_time (per-rank monotonic clock, NOT - # replicated across CP ranks — unlike host eviction which uses the - # deterministic (priority, node.id)). So ranks can pick DIFFERENT - # victims -> free DIFFERENT pages -> diverge alloc_pages_with_owners - # -> req_to_token disagrees across ranks. Log the victim id set per - # rank for this load op; the analyzer diffs it across ranks. A - # cross-rank mismatch confirms the divergence trigger. - cptrace( - 1, - "victim_set", - node_id=node_id, - cprank=self._cp_hicache_cp_rank(), - nvictims=len(eviction_plan.victims), - victims=sorted( - int(getattr(v, "id", -1)) for v in eviction_plan.victims - ), - ) plan_elapsed_ms = (time.perf_counter() - evict_start_time) * 1000.0 if plan_elapsed_ms >= 1000.0: logger.warning( @@ -2705,21 +2685,6 @@ 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: return @@ -3034,13 +2999,6 @@ 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), @@ -3369,8 +3327,6 @@ 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 @@ -3383,17 +3339,6 @@ 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) @@ -3712,15 +3657,6 @@ 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( @@ -4310,43 +4246,11 @@ 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() diff --git a/python/sglang/srt/mem_cache/memory_pool_host.py b/python/sglang/srt/mem_cache/memory_pool_host.py index 8d0844b1e..fd5cafba0 100644 --- a/python/sglang/srt/mem_cache/memory_pool_host.py +++ b/python/sglang/srt/mem_cache/memory_pool_host.py @@ -12,7 +12,6 @@ from typing import Optional, Tuple import psutil import torch -from sglang.srt.mem_cache.cp_hicache_trace import cptrace, khash, knz, rng, trace_enabled from sglang.jit_kernel.hicache import ( can_use_hicache_jit_kernel, @@ -2552,18 +2551,6 @@ class NSATokenToKVPoolHost(MLATokenToKVPoolHost): host_page_indices, device_page_indices = self._get_indexer_page_indices( host_indices, device_indices ) - if trace_enabled(2): - _dev = device_pool.index_k_with_scale_buffer[device_layer_slot] - _rows = _dev[device_page_indices] - cptrace( - 2, - "indexer_backup_hash", - layer=layer_id, - host=rng(host_indices), - kvhash=khash(_rows), - nz=knz(_rows), - npages=int(device_page_indices.numel()), - ) use_kernel = io_backend == "kernel" and self.indexer_page_stride_size % 8 == 0 if use_kernel: if self.layout == "layer_first": @@ -2638,20 +2625,6 @@ class NSATokenToKVPoolHost(MLATokenToKVPoolHost): self._load_indexer_to_device_per_layer( device_pool, host_indices, device_indices, layer_id, io_backend ) - if trace_enabled(2) and self._is_device_index_layer_active(device_pool, layer_id): - _dls = self._device_index_layer_slot(device_pool, layer_id) - _, _dpi = self._get_indexer_page_indices(host_indices, device_indices) - _rows = device_pool.index_k_with_scale_buffer[_dls][_dpi] - cptrace( - 2, - "indexer_reload_hash", - layer=layer_id, - host=rng(host_indices), - kvhash=khash(_rows), - nz=knz(_rows), - npages=int(_dpi.numel()), - ) - def backup_from_device_all_layer( self, device_pool, host_indices, device_indices, io_backend ): diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 2ddc4863c..62d5321a0 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -73,7 +73,6 @@ from sglang.srt.layers.communicator import ( enable_moe_dense_fully_dp, get_attn_tp_context, ) -from sglang.srt.mem_cache.cp_hicache_trace import fwd_hash as _cp_fwd_hash from sglang.srt.layers.communicator_nsa_cp import NSACPLayerCommunicator from sglang.srt.layers.dp_attention import ( get_attention_cp_rank, @@ -775,7 +774,6 @@ class DeepseekV2MoE(nn.Module): hidden_states, ) - _cp_fwd_hash(forward_batch, getattr(self, "layer_id", -1), "moe_postsel", hidden_states) shared_output = None sbo_enabled_flag = self._fuse_shared_experts_inside_sbo and not self.is_nextn sbo_overlap_dispatch_flag = ( @@ -788,7 +786,6 @@ class DeepseekV2MoE(nn.Module): if hidden_states.shape[0] > 0: # router_logits: (num_tokens, n_experts) router_logits = self.gate(hidden_states, forward_batch=forward_batch) - _cp_fwd_hash(forward_batch, getattr(self, "layer_id", -1), "router_logits", router_logits) if not sbo_enabled_flag: if self.alt_stream is not None: self.alt_stream.wait_stream(torch.cuda.current_stream()) @@ -806,8 +803,6 @@ class DeepseekV2MoE(nn.Module): layer_id=self.layer_id, ), ) - _cp_fwd_hash(forward_batch, getattr(self, "layer_id", -1), "topk_ids", getattr(topk_output, "topk_ids", None)) - _cp_fwd_hash(forward_batch, getattr(self, "layer_id", -1), "topk_w", getattr(topk_output, "topk_weights", None)) else: topk_output = self.topk.empty_topk_output(hidden_states.device) @@ -960,7 +955,6 @@ class DeepseekV2MoE(nn.Module): hidden_states=hidden_states, topk_output=topk_output, ) - _cp_fwd_hash(forward_batch, getattr(self, "layer_id", -1), "experts_out", final_hidden_states) if ( hidden_states.shape[0] > 0 @@ -984,14 +978,12 @@ class DeepseekV2MoE(nn.Module): ): final_hidden_states *= self.routed_scaling_factor - _cp_fwd_hash(forward_batch, getattr(self, "layer_id", -1), "moe_out_compact", final_hidden_states) if local_compute_hidden_states is not None: final_hidden_states = restore_cp_local_valid_rows_for_moe( forward_batch, final_hidden_states, local_compute_hidden_states, ) - _cp_fwd_hash(forward_batch, getattr(self, "layer_id", -1), "moe_out_restored", final_hidden_states) return final_hidden_states @@ -1738,7 +1730,6 @@ class DeepseekV2DecoderLayer(nn.Module): forward_batch, quant_format, ) - _cp_fwd_hash(forward_batch, getattr(self, "layer_id", -1), "attn_in", hidden_states) previous_cp_shared_kv_num_model_layers = getattr( forward_batch, "cp_shared_kv_num_model_layers", None @@ -1770,13 +1761,10 @@ class DeepseekV2DecoderLayer(nn.Module): hidden_states, topk_indices = hidden_states else: topk_indices = None - _cp_fwd_hash(forward_batch, getattr(self, "layer_id", -1), "attn_out", hidden_states) - _cp_fwd_hash(forward_batch, getattr(self, "layer_id", -1), "topk", topk_indices) hidden_states, residual = self.layer_communicator.prepare_mlp( hidden_states, residual, forward_batch ) - _cp_fwd_hash(forward_batch, getattr(self, "layer_id", -1), "moe_in", hidden_states) should_allreduce_fusion = ( self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer( @@ -2162,18 +2150,6 @@ class DeepseekV2Model(nn.Module): torch.cuda.current_stream(), ) - # Dump flush at model end (gated by SGLANG_NSA_DUMP_DIR + rid 'dump-'/'dumpfh-'). - # Placed AFTER the CP gather so final_hidden is the model's actual output (the - # first-token source); for 'fh' mode this is the cache-bust-vs-cache-hit probe. - try: - from sglang.srt.mem_cache.cp_hicache_trace import dump_flush as _cp_dump_flush - - _cp_dump_flush( - forward_batch, final_hidden=hidden_states, positions=positions - ) - except Exception: - pass - if len(aux_hidden_states) == 0: return hidden_states return hidden_states, aux_hidden_states