chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Add env-gated CP HiCache KV round-trip corruption tracing

Browse Source 
Heavily-forked CP HiCache produces KV corruption only after an L1->L2->L1
round-trip (eviction + reload), scaling with cached volume. Static reading and
upstream-diff are exhausted, so trace the data flow and check four invariants
directly, gated by SGLANG_CP_HICACHE_KV_TRACE (0=off, 1=structural lifecycle,
2=+compose/free/ack), off by default and cheap when off.

Trace points (keyed by node_id; rid_map ties client rid->node_id):
- backup_reserve / backup_d2h: host<->phys slots written           (H2/H4)
- split: prefix+suffix partition vs parent                         (H1)
- evict: write_pending at device free                              (H4)
- reload_node: host slots read on reload                           (H2)
- reload_assign: reloaded device locs -> forward
- dev_free / write_ack: device free vs backup-complete ordering    (H4)
- compose / remap: per-forward row-id cache reuse + unmapped count  (H3)

New mem_cache/cp_hicache_trace.py (cptrace/rng helper; values rendered
space-free for offline parsing). The companion analyzer joins these by
node_id; the decisive, request-independent check is H2: any node whose reload
host-slot fingerprint differs from its backup fingerprint is the corruption
source.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
孙远航
2026-06-17 01:14:31 +00:00
parent 34763d92b1
commit 38ef664b4c
6 changed files with 254 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
python/sglang/srt/environ.py
View File

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

37
python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py
View File

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

58
python/sglang/srt/managers/cache_controller.py
View File

@@ -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 = (

9
python/sglang/srt/mem_cache/allocator.py
View File

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

71
python/sglang/srt/mem_cache/cp_hicache_trace.py Normal file
View File

@@ -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 <tag>] 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

76
python/sglang/srt/mem_cache/hiradix_cache.py
View File

@@ -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", "<unknown>"),
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()