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

Trace CP HiCache cache-hit state for GSM8K regressions

Browse Source 
The cache-hit accuracy drop is not isolated to one obvious kernel boundary, so the scheduler/radix/HiCache path needs request-correlated evidence.  Add default-off, rate-limited CP shared-KV debug logs for prefix matching, valid-tail flooring, pending backup attach/commit, load-back planning, and unfinished-request cache insertion.  The investigation notes capture current GSM8K pass1/pass2 evidence and rejected interpretations to avoid repeated log archaeology.

Constraint: Logs must be default-off and bounded because these paths are scheduler/control hot paths.

Rejected: Always-on INFO tracing | would distort the exact latency and cache-hit behavior under investigation.

Confidence: medium

Scope-risk: moderate

Directive: Remove or keep env-gated only after the GSM8K cache-hit root cause is closed; do not leave unbounded hot-path logs.

Tested: python -m py_compile on changed runtime files.

Not-tested: Local pytest blocked before collection by missing orjson dependency; fresh restarted GSM8K verification still pending.
(cherry picked from commit bfc6d1473dc2a5d72bc3a8d6fca1e2429537be0e)
This commit is contained in:
laoyao0822
2026-06-07 23:56:37 +08:00
parent ff4717c2fb
commit b324407def
4 changed files with 1093 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
python/sglang/srt/managers/schedule_policy.py
View File

@@ -198,7 +198,8 @@ class SchedulePolicy:
# NOTE: the prefix_indices must always be aligned with last_node
match_result = self.tree_cache.match_prefix(
MatchPrefixParams(
key=RadixKey(token_ids=prefix_ids, extra_key=extra_key)
key=RadixKey(token_ids=prefix_ids, extra_key=extra_key),
req=r,
)
)
(
@@ -857,6 +858,7 @@ class PrefillAdder:
last_host_node=req.last_host_node,
host_hit_length=req.host_hit_length,
mem_quota=self._get_load_back_mem_quota(real_input_tokens),
req=req,
)
)
req.prefix_indices = torch.cat([req.prefix_indices, new_indices])

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

@@ -65,6 +65,33 @@ if TYPE_CHECKING:
logger = logging.getLogger(__name__)
_CP_HICACHE_BS_GT1_DEBUG_COUNTS: Dict[str, int] = {}
def _cp_hicache_bs_gt1_debug(
key: str,
message: str,
*args,
limit: Optional[int] = None,
) -> None:
"""Env-gated, rate-limited CP HiCache debug logging.
Keep this local to mem_cache to avoid importing NSA attention helpers from
the radix cache path. It is intentionally default-off because these paths
run on scheduler/control hot paths.
"""
if not envs.SGLANG_CP_SHARED_KV_BS_GT1_DEBUG.get():
return
if limit is None:
limit = envs.SGLANG_CP_SHARED_KV_BS_GT1_DEBUG_LIMIT.get()
limit = int(limit)
count = _CP_HICACHE_BS_GT1_DEBUG_COUNTS.get(key, 0)
if limit > 0 and count >= limit:
return
_CP_HICACHE_BS_GT1_DEBUG_COUNTS[key] = count + 1
logger.info("[CP_SHARED_KV_BS_GT1_DEBUG][hicache] event=%s " + message, key, *args)
def _estimate_hicache_size_per_token(kv_cache) -> int:
dtype_size = kv_cache.store_dtype.itemsize
@@ -1992,6 +2019,20 @@ class HiRadixCache(RadixCache):
node.host_value = None
if pending.locked:
self.dec_node_lock_ref(node)
_cp_hicache_bs_gt1_debug(
"commit_pending_backup",
"node_id=%d key_len=%d host_len=%d padded_len=%s "
"owned_positions=%d host_indices=%d page_owner_pages=%d "
"remaining_pending=%d",
node.id,
len(node.key),
node.host_len,
getattr(pending.metadata, "padded_len", None),
pending.metadata.owned_positions.numel(),
pending.metadata.host_indices.numel(),
pending.metadata.page_owners.numel(),
len(self.pending_host_backups),
)
return node
def _rollback_pending_backup(self, node_id: int) -> TreeNode:
@@ -2156,6 +2197,21 @@ class HiRadixCache(RadixCache):
locked=True,
)
prepared.attached = True
_cp_hicache_bs_gt1_debug(
"attach_prepared_backup",
"node_id=%d key_len=%d value_len=%d logical_len=%d padded_len=%s "
"owned_positions=%d host_indices=%d page_owner_pages=%d "
"pending_backups=%d",
node.id,
len(node.key),
len(node.value),
prepared.logical_len,
getattr(prepared.metadata, "padded_len", None),
prepared.metadata.owned_positions.numel(),
prepared.metadata.host_indices.numel(),
prepared.metadata.page_owners.numel(),
len(self.pending_host_backups),
)
logger.info(
"[HiCache-write] attached prepared CP backup: node_id=%d logical_len=%d owned_positions=%d pending_backups=%d",
node.id,
@@ -2280,6 +2336,7 @@ class HiRadixCache(RadixCache):
request_len: int,
*,
floor_exact_key: bool = False,
debug_req_id: Optional[str] = None,
) -> int:
"""Floor exact CP valid-tail hits to the previous physical page.
@@ -2306,7 +2363,29 @@ class HiRadixCache(RadixCache):
or prefix_len % self.page_size == 0
):
return prefix_len
return floor_to_page_len(prefix_len, self.page_size)
floored_len = floor_to_page_len(prefix_len, self.page_size)
metadata = getattr(child, "cp_hicache", None)
_cp_hicache_bs_gt1_debug(
"match_floor_valid_tail",
"rid=%s node_id=%s prefix_len=%d floored_len=%d request_len=%d "
"child_key_len=%d page_size=%d evicted=%s host_len=%d "
"has_cp_hicache=%s padded_len=%s pending_write=%s floor_exact_key=%s",
debug_req_id,
getattr(child, "id", None),
prefix_len,
floored_len,
request_len,
len(child.key),
self.page_size,
child.evicted,
getattr(child, "host_len", 0),
metadata is not None,
getattr(metadata, "padded_len", None),
self._node_host_write_pending(child),
floor_exact_key,
limit=64,
)
return floored_len
def _cp_subtree_has_unprunable_state(self, node: TreeNode) -> bool:
stack = [node]
@@ -2558,6 +2637,25 @@ class HiRadixCache(RadixCache):
metadata=result.metadata,
logical_len=len(kv_indices),
)
_cp_hicache_bs_gt1_debug(
"prepare_write_backup",
"node_id=%d rid=%s logical_len=%d padded_len=%s "
"owned_positions=%d host_indices=%d draft_host_indices=%s "
"page_owner_pages=%d admission_checked=%s",
node_id,
getattr(req, "rid", "<unknown>"),
len(kv_indices),
getattr(result.metadata, "padded_len", None),
result.metadata.owned_positions.numel(),
result.metadata.host_indices.numel(),
(
None
if getattr(result.metadata, "draft_host_indices", None) is None
else result.metadata.draft_host_indices.numel()
),
result.metadata.page_owners.numel(),
admission_checked,
)
logger.info(
"[HiCache-write] prepared CP per-layer backup before forward: node_id=%d rid=%s logical_len=%d owned_positions=%d",
node_id,
@@ -3406,6 +3504,30 @@ class HiRadixCache(RadixCache):
self.dec_lock_ref(ancester_node)
raise
host_hit_len = load_back_plan.host_hit_len
node_debug_summary = [
(
getattr(x, "id", None),
len(x.key) if getattr(x, "key", None) is not None else 0,
getattr(x, "host_len", 0),
getattr(getattr(x, "cp_hicache", None), "padded_len", None),
)
for x in nodes_to_load
]
_cp_hicache_bs_gt1_debug(
"load_back_plan",
"node_id=%d nodes=%s host_hit_len=%d threshold=%d "
"required_by_owner=%s available_by_owner=%s deficit_by_owner=%s "
"free_room_deficit_by_owner=%s",
last_hit_node.id,
node_debug_summary,
host_hit_len,
self.load_back_threshold,
load_back_plan.required_by_owner,
load_back_plan.available_by_owner,
load_back_plan.deficit_by_owner,
load_back_plan.free_room_deficit_by_owner,
limit=64,
)
logger.info(
"[HiCache-load] load_back CP: node_id=%d nodes_to_load=%d "
"host_hit_len=%d threshold=%d required_by_owner=%s "
@@ -3638,16 +3760,46 @@ class HiRadixCache(RadixCache):
):
last_node = params.last_host_node
mem_quota = params.mem_quota
debug_req_id = getattr(getattr(params, "req", None), "rid", None)
if last_node.evicted:
if self._uses_cp_hicache:
_cp_hicache_bs_gt1_debug(
"init_load_back_start",
"rid=%s last_host_node=%s host_hit_length=%d mem_quota=%s "
"node_host_len=%d",
debug_req_id,
getattr(last_node, "id", None),
int(params.host_hit_length),
mem_quota,
getattr(last_node, "host_len", 0),
)
loading_values = self.load_back(last_node, mem_quota)
if loading_values is not None:
logger.info(
f"loading back {len(loading_values)} tokens for node {last_node.id}"
)
if self._uses_cp_hicache:
_cp_hicache_bs_gt1_debug(
"init_load_back_loaded",
"rid=%s loaded_tokens=%d loaded_node=%s host_hit_length=%d",
debug_req_id,
len(loading_values),
getattr(last_node, "id", None),
int(params.host_hit_length),
)
return loading_values, last_node
while last_node.evicted:
last_node = last_node.parent
if self._uses_cp_hicache:
_cp_hicache_bs_gt1_debug(
"init_load_back_unloaded",
"rid=%s fallback_device_node=%s host_hit_length=%d mem_quota=%s",
debug_req_id,
getattr(last_node, "id", None),
int(params.host_hit_length),
mem_quota,
)
return (
torch.empty((0,), dtype=torch.int64, device=self.device),
@@ -3870,10 +4022,13 @@ class HiRadixCache(RadixCache):
deferred_node = None
try:
debug_req = getattr(params, "req", None)
debug_req_id = getattr(debug_req, "rid", None)
value, last_node = self._match_prefix_helper(
self.root_node,
key,
floor_exact_key=getattr(params, "cp_floor_exact", True),
debug_req_id=debug_req_id,
)
except HiCachePendingBackupSplit as exc:
value = []
@@ -3906,6 +4061,28 @@ class HiRadixCache(RadixCache):
if not last_host_node.backuped:
last_host_node = self.root_node
if self._uses_cp_hicache:
debug_req = getattr(params, "req", None)
_cp_hicache_bs_gt1_debug(
"match_prefix_result",
"rid=%s key_len=%d device_tokens=%d host_hit_length=%d "
"last_device_node=%s last_device_evicted=%s last_device_host_len=%d "
"last_host_node=%s last_host_evicted=%s last_host_host_len=%d "
"deferred_node=%s cp_floor_exact=%s",
getattr(debug_req, "rid", None),
len(key),
len(value),
host_hit_length,
getattr(last_node, "id", None),
getattr(last_node, "evicted", None),
getattr(last_node, "host_len", 0),
getattr(last_host_node, "id", None),
getattr(last_host_node, "evicted", None),
getattr(last_host_node, "host_len", 0),
getattr(deferred_node, "id", None),
getattr(params, "cp_floor_exact", True),
)
return MatchResult(
device_indices=value,
last_device_node=last_node,
@@ -4003,7 +4180,12 @@ class HiRadixCache(RadixCache):
return matched_length
def _match_prefix_helper(
self, node: TreeNode, key: RadixKey, *, floor_exact_key: bool = True
self,
node: TreeNode,
key: RadixKey,
*,
floor_exact_key: bool = True,
debug_req_id: Optional[str] = None,
):
node.last_access_time = time.monotonic()
child_key = self.get_child_key_fn(key)
@@ -4021,6 +4203,7 @@ class HiRadixCache(RadixCache):
raw_prefix_len,
len(key),
floor_exact_key=floor_exact_key,
debug_req_id=debug_req_id,
)
stop_after_page_floor = prefix_len != raw_prefix_len
prune_stale_tail_after_split = stop_after_page_floor

90
python/sglang/srt/mem_cache/radix_cache.py
View File

@@ -31,6 +31,8 @@ from typing import TYPE_CHECKING, Any, Iterator, List, Optional, Tuple, Union
import torch
from sglang.srt.environ import envs
logger = logging.getLogger(__name__)
from sglang.srt.disaggregation.kv_events import (
@@ -66,6 +68,41 @@ from sglang.srt.mem_cache.hicache_storage import get_hash_str, hash_str_to_int64
if TYPE_CHECKING:
from sglang.srt.managers.schedule_batch import Req
_CP_SHARED_KV_RADIX_DEBUG_COUNTS: dict[str, int] = {}
def _cp_shared_kv_radix_debug(
key: str,
message: str,
*args,
limit: int | None = None,
) -> None:
if not envs.SGLANG_CP_SHARED_KV_BS_GT1_DEBUG.get():
return
if limit is None:
limit = int(envs.SGLANG_CP_SHARED_KV_BS_GT1_DEBUG_LIMIT.get())
count = _CP_SHARED_KV_RADIX_DEBUG_COUNTS.get(key, 0)
if limit > 0 and count >= limit:
return
_CP_SHARED_KV_RADIX_DEBUG_COUNTS[key] = count + 1
logger.info("[CP_SHARED_KV_BS_GT1_DEBUG][radix] event=%s " + message, key, *args)
def _cp_shared_kv_page_head_for_debug(
locs: torch.Tensor | None,
*,
page_size: int,
max_items: int = 8,
) -> list[int] | None:
if locs is None:
return None
if page_size <= 0 or locs.numel() == 0:
return []
pages = locs[::page_size]
if pages.numel() == 0:
return []
return (pages[:max_items].detach().to("cpu", non_blocking=False) // page_size).tolist()
class RadixKey:
def __init__(
@@ -650,6 +687,24 @@ class RadixCache(BasePrefixCache):
values = kv_indices[: len(keys)].to(dtype=torch.int64, copy=True)
radix_key = RadixKey(keys, req.extra_key, is_bigram=self.is_eagle)
prepared_cp_backup = getattr(req, "cp_hicache_prepared_backup", None)
if getattr(self, "_uses_cp_hicache", False):
_cp_shared_kv_radix_debug(
"unfinished_pre_insert",
"rid=%s chunked=%s key_len=%d kv_len=%d cache_protected_len=%d "
"req_pool_idx=%s value_pages=%s kv_pages=%s prefix_pages=%s",
getattr(req, "rid", None),
chunked,
len(keys),
int(kv_indices.numel()),
int(getattr(req, "cache_protected_len", 0) or 0),
getattr(req, "req_pool_idx", None),
_cp_shared_kv_page_head_for_debug(values, page_size=self.page_size),
_cp_shared_kv_page_head_for_debug(kv_indices, page_size=self.page_size),
_cp_shared_kv_page_head_for_debug(
getattr(req, "prefix_indices", None), page_size=self.page_size
),
limit=160,
)
# Radix Cache takes one ref in memory pool
result = self.insert(
@@ -706,11 +761,46 @@ class RadixCache(BasePrefixCache):
match_result.last_device_node,
)
assert len(new_indices) == len(keys), f"{len(new_indices)=}, {len(keys)=}"
if getattr(self, "_uses_cp_hicache", False):
_cp_shared_kv_radix_debug(
"unfinished_post_match",
"rid=%s chunked=%s key_len=%d old_cache_protected_len=%d "
"new_indices_len=%d new_node=%s new_pages=%s kv_pages_before=%s",
getattr(req, "rid", None),
chunked,
len(keys),
int(getattr(req, "cache_protected_len", 0) or 0),
int(new_indices.numel()),
getattr(new_last_node, "id", None),
_cp_shared_kv_page_head_for_debug(
new_indices, page_size=self.page_size
),
_cp_shared_kv_page_head_for_debug(kv_indices, page_size=self.page_size),
limit=160,
)
self.req_to_token_pool.write(
(req.req_pool_idx, slice(req.cache_protected_len, len(new_indices))),
new_indices[req.cache_protected_len :],
)
if getattr(self, "_uses_cp_hicache", False):
final_kv_indices = self.req_to_token_pool.req_to_token[
req.req_pool_idx, : len(token_ids)
]
_cp_shared_kv_radix_debug(
"unfinished_post_write",
"rid=%s chunked=%s key_len=%d cache_protected_len=%d "
"final_kv_len=%d final_pages=%s",
getattr(req, "rid", None),
chunked,
len(keys),
int(getattr(req, "cache_protected_len", 0) or 0),
int(final_kv_indices.numel()),
_cp_shared_kv_page_head_for_debug(
final_kv_indices, page_size=self.page_size
),
limit=160,
)
# The cache_protected_len is not always equal to len(req.prefix_indices)
# since for page_size > 1, the partial part is added to req.prefix_indices, but that part of kv indices is not added to the tree.