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

CP HiCache: replicated-clock SLRU eviction, collective-safe (Phase 1)

Browse Source 
Make CP shared-KV HiCache eviction collective-free and scan-resistant by
removing every per-rank wall-clock input from the eviction/admission decisions
(any such input desyncs the must-be-replicated victim set / batch across the 8
CP ranks and can deadlock the collective-coupled writeback/load-back).

- Replicated logical clock: TreeNode.next_access_time() (a process-global
  monotonic counter, mirrors mamba/swa's get_last_access_time) replaces
  time.monotonic() as the source of last_access_time at every radix bump site
  (radix_cache __init__/match/insert; hiradix CP match/insert/_insert_helper_host;
  reset() zeroes it). creation_time/pin_expiry intentionally stay wall-clock.
  Because match/insert events are replicated (reqs broadcast from rank 0 over the
  replicated tree), last_access_time is now identical on every rank. Unique ints
  also give a strict total order (no LRU tie ambiguity). No duration arithmetic
  reads last_access_time, so non-CP LRU is unchanged; mamba/swa use their own
  TreeNode and are untouched.
- CpReplicatedSLRUStrategy = (is_protected[hit>=2], last_access_time, node.id):
  scan-resistant (cold one-shots stay probationary, evicted before reused/
  protected prefixes), recency-within-segment ages out cold (not LFU), node.id
  is the deterministic total-order tiebreak (heaps never compare TreeNodes).
  Overridden for CP so it reaches every get_priority eviction site; the host
  write-admission key _cp_host_evict_key is switched from FIFO-by-id to it.
- Co-fixes for the same per-rank-wall-clock class: pin_prefix is forbidden under
  CP (its pin_expiry victim-eligibility check is per-rank wall-clock; SLRU
  scan-resistance covers the benefit); the affinity head_age_s admission input is
  disabled (=0.0, relying on the replicated head_defer_count bound); and a
  server_args fail-fast guards CP shared-KV against SGLANG_REQ_WAITING_TIMEOUT>0
  (its waiting-queue pruning is per-rank wall-clock).

Tests: test_evict_policy.py adds TestCpReplicatedSLRUStrategy (scan-resistance,
recency, total-order, full ordering) -- 30 passed in the dev-cu13 container;
test_radix_cache_unit.py adds logical-clock determinism/total-order tests.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
孙远航
2026-06-19 18:59:01 +00:00
parent 8d73919d02
commit 8f1e85a992
7 changed files with 220 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
python/sglang/srt/managers/scheduler.py
View File

@@ -2665,12 +2665,16 @@ class Scheduler(
batch_warm_led=bool(affinity_batch_warm_led),
is_head=is_head,
head_defer_count=req.affinity_defer_count,
head_age_s=(
time.perf_counter()
- req.time_stats.wait_queue_entry_time
if req.time_stats.wait_queue_entry_time > 0.0
else 0.0
),
# head_age_s DISABLED (=0.0) for CP collective-safety:
# time.perf_counter() is per-rank, so at the MAX_AGE_S boundary
# ranks could admit DIFFERENT batches → permanent node.id /
# access-clock / eviction desync (breaks the replicated
# event-stream premise). The replicated head_defer_count
# (MAX_DEFER passes) already bounds cold-head starvation, and
# the SKIP_COLD path only triggers under load (frequent passes).
# Restore a time-like bound only via a replicated logical age
# (ticks-since-enqueue). MAX_AGE_S is currently unused.
head_age_s=0.0,
window_used=affinity_window_used,
)
if decision is AffinityDecision.SKIP_COLD:

32
python/sglang/srt/mem_cache/evict_policy.py
View File

@@ -63,3 +63,35 @@ class SLRUStrategy(EvictionStrategy):
is_protected = 1 if node.hit_count >= self.protected_threshold else 0
return (is_protected, node.last_access_time)
class CpReplicatedSLRUStrategy(EvictionStrategy):
"""SLRU for CP shared-KV HiCache, made collective-safe by a strict total order.
Value = (is_protected, last_access_time, node.id), smallest evicted first:
- is_protected (hit_count >= threshold): probationary nodes (cold one-shots,
e.g. a long-request injector) are evicted before ANY protected node →
scan-resistance for the shared working set.
- last_access_time: LRU within a segment, so a now-cold protected prefix
ages out by recency (NOT by frequency — avoids LFU pinning a
stale-high-hit-count prefix).
- node.id: final deterministic tiebreak → a strict TOTAL order (the
per-call logical clock can repeat a value across nodes touched in one
match; node.id disambiguates).
Every field is rank-replicated: last_access_time is the logical clock
(TreeNode.next_access_time()), node.id/hit_count advance identically on all
CP ranks. So all ranks select identical victims with no collective. (Plain
LRU/LFU/SLRU here tiebreak on per-rank wall-clock last_access_time and would
diverge victims → collective deadlock.)
"""
def __init__(self, protected_threshold: int = 2):
# protected_threshold T=2: a one-shot (hit_count==1) stays probationary;
# a genuinely-reused prefix (hit>=2) is protected. T is the scan-resistance
# tuning knob (design item #3, to be measured).
self.protected_threshold = protected_threshold
def get_priority(self, node: "TreeNode") -> Tuple[int, float, int]:
is_protected = 1 if node.hit_count >= self.protected_threshold else 0
return (is_protected, node.last_access_time, node.id)

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

@@ -671,6 +671,20 @@ class HiRadixCache(RadixCache):
super().__init__(params=params)
if self._uses_cp_hicache:
# CP shared-KV eviction must be COLLECTIVE-FREE: every rank must pick
# identical victims from rank-replicated state, else partial backups /
# all-reduce-shape deadlock. The base LRU/SLRU strategies tiebreak on
# per-rank wall-clock last_access_time and diverge. Override with the
# replicated-clock SLRU value (is_protected, last_access_time[logical],
# node.id) — a strict total order over replicated state. Reaches every
# eviction key site that uses eviction_strategy.get_priority (device
# owner-lane score, host physical-slot heap, generic evict); the host
# write-admission key _cp_host_evict_key is switched to it explicitly.
from sglang.srt.mem_cache.evict_policy import CpReplicatedSLRUStrategy
self.eviction_strategy = CpReplicatedSLRUStrategy()
def _cp_hicache_all_reduce(self, tensor, *, op, group, tag: str) -> None:
start = time.perf_counter()
torch.distributed.all_reduce(tensor, op=op, group=group)
@@ -922,10 +936,11 @@ class HiRadixCache(RadixCache):
)
def _cp_host_evict_key(self, node: TreeNode):
return (
int(getattr(node, "priority", 0)),
int(getattr(node, "id", 0)),
)
# Replicated-clock SLRU total order (is_protected, last_access_time[logical],
# node.id) — replaces the old FIFO-by-node.id key. Same strategy the device
# owner-lane / physical-slot eviction paths use, so all host/device victim
# selection is consistent and collective-free.
return self.eviction_strategy.get_priority(node)
def _cp_host_leaf_is_plannable_victim(self, node: TreeNode) -> bool:
if node == getattr(self, "root_node", None):
@@ -1249,6 +1264,10 @@ class HiRadixCache(RadixCache):
)
if contribution <= 0 and unlock_contribution <= 0:
continue
# get_priority returns the CpReplicatedSLRUStrategy tuple
# (is_protected, last_access_time[logical], node.id) under CP — a
# rank-replicated value; the trailing node.id keeps the full score a
# strict total order even if the strategy is ever swapped.
score = (
-int(contribution),
-int(unlock_contribution),
@@ -1996,6 +2015,7 @@ class HiRadixCache(RadixCache):
def reset(self):
TreeNode.counter = 0
TreeNode.last_access_counter = 0
self.cache_controller.reset()
self.token_to_kv_pool_host.clear()
self.cache_controller.clear_draft_host_pool()
@@ -3149,6 +3169,14 @@ class HiRadixCache(RadixCache):
self, token_ids: List[int], ttl_seconds: int = 300
) -> Tuple[int, Optional[str]]:
"""Pin nodes along a prefix path. Returns (nodes_pinned, reject_reason)."""
if self._uses_cp_hicache:
# Forbidden under CP shared-KV: pin_expiry is a per-rank wall-clock TTL
# and the eviction victim-eligibility check (time.monotonic() <=
# pin_expiry) would then diverge across CP ranks → non-replicated
# victim set → collective-free eviction breaks. SLRU scan-resistance
# already protects the hot working set; re-enable only with a
# replicated logical TTL.
return (0, "pin_prefix is not supported under CP shared-KV HiCache")
if self.disable or not token_ids:
return (0, None)
@@ -4176,7 +4204,11 @@ class HiRadixCache(RadixCache):
def _insert_helper_host(
self, node: TreeNode, key: RadixKey, host_value, hash_value
):
node.last_access_time = time.monotonic()
# NOTE: dead under CP today (storage forbidden); converted for whole-repo
# correctness so a future CP-aware L3 prefetch insert stays rank-replicated.
# (The prefetch-completion insert is MIN-synchronized, so the bump order is
# replicated; if it ever ran at a non-replicated point this must be excluded.)
node.last_access_time = TreeNode.next_access_time()
if len(key) == 0:
return 0
@@ -4185,7 +4217,7 @@ class HiRadixCache(RadixCache):
matched_length = 0
while len(key) > 0 and child_key in node.children.keys():
node = node.children[child_key]
node.last_access_time = time.monotonic()
node.last_access_time = TreeNode.next_access_time()
# Refresh pin TTL on host insert hit
if self._is_pinned(node):
node.pin_expiry = time.monotonic() + node.pin_ttl
@@ -4219,13 +4251,13 @@ class HiRadixCache(RadixCache):
def _match_prefix_helper(
self, node: TreeNode, key: RadixKey, *, floor_exact_key: bool = True
):
node.last_access_time = time.monotonic()
node.last_access_time = TreeNode.next_access_time()
child_key = self.get_child_key_fn(key)
value = []
while len(key) > 0 and child_key in node.children.keys():
child = node.children[child_key]
child.last_access_time = time.monotonic()
child.last_access_time = TreeNode.next_access_time()
# Refresh pin TTL on cache hit
if self._is_pinned(child):
child.pin_expiry = time.monotonic() + child.pin_ttl
@@ -4351,7 +4383,7 @@ class HiRadixCache(RadixCache):
while len(key) > 0 and child_key in node.children.keys():
parent_node = node
node = node.children[child_key]
node.last_access_time = time.monotonic()
node.last_access_time = TreeNode.next_access_time()
node.priority = max(node.priority, priority)
raw_prefix_len = self.key_match_fn(node.key, key)
prefix_len = self._cp_floor_exact_valid_tail_extension_len(

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

@@ -186,6 +186,23 @@ def floor_to_page_len(length: int, page_size: int) -> int:
class TreeNode:
counter = 0
# Replicated logical access clock (mirrors mamba_radix_cache's
# get_last_access_time pattern). A process-global monotonically-increasing
# counter used as the source of ``last_access_time``. Under CP shared-KV
# every rank advances it through the IDENTICAL sequence of match/insert
# events (requests are broadcast from rank 0 and the radix tree is
# replicated), so ``last_access_time`` is rank-replicated and eviction
# victim order is collective-free. The old per-rank ``time.monotonic()``
# diverged victim order across ranks and could deadlock the collective-
# coupled writeback/load-back. Unique integers also give a strict total
# order (no tie ambiguity). Reset alongside ``counter`` in HiRadixCache.reset.
last_access_counter = 0
@staticmethod
def next_access_time() -> int:
v = TreeNode.last_access_counter
TreeNode.last_access_counter += 1
return v
def __init__(self, id: Optional[int] = None, priority: int = 0):
self.children = {}
@@ -197,7 +214,11 @@ class TreeNode:
0.0 # absolute expiry time (time.monotonic()), 0 = not pinned
)
self.pin_ttl: int = 0 # original TTL in seconds, for refresh-on-hit
self.last_access_time = time.monotonic()
self.last_access_time = TreeNode.next_access_time()
# creation_time stays wall-clock: it is only read by FIFO/FILO eviction
# (not the CP/default LRU/SLRU path). FIFO/FILO are not CP-safe today and
# are not the default; if ever enabled under CP, route this to the
# logical clock too.
self.creation_time = time.monotonic()
self.hit_count = 0
@@ -926,7 +947,7 @@ class RadixCache(BasePrefixCache):
##### Internal Helper Functions #####
def _match_prefix_helper(self, node: TreeNode, key: RadixKey):
access_time = time.monotonic()
access_time = TreeNode.next_access_time()
node.last_access_time = access_time
child_key = self.get_child_key_fn(key)
@@ -994,7 +1015,7 @@ class RadixCache(BasePrefixCache):
# Convert None priority to 0
if priority is None:
priority = 0
access_time = time.monotonic()
access_time = TreeNode.next_access_time()
node.last_access_time = access_time
# Update priority along the path (take max to propagate higher priority)
node.priority = max(node.priority, priority)

13
python/sglang/srt/server_args.py
View File

@@ -994,6 +994,19 @@ class ServerArgs:
"hicache_storage_backend in the host-only CP HiCache stage. "
"Disable hicache_storage_backend or disable CP shared KV."
)
assert not (
self.enable_nsa_prefill_cp_shared_kv
and envs.SGLANG_REQ_WAITING_TIMEOUT.get() > 0
), (
"enable_nsa_prefill_cp_shared_kv is incompatible with "
"SGLANG_REQ_WAITING_TIMEOUT > 0: _abort_on_waiting_timeout prunes the "
"waiting queue using per-rank time.perf_counter(), so CP ranks could "
"drop different requests at the deadline boundary and desync the "
"(must-be-replicated) waiting queue -> radix node ids / eviction order "
"-> collective deadlock. Unset SGLANG_REQ_WAITING_TIMEOUT or disable CP "
"shared KV. (Same per-rank-wall-clock class as the disabled affinity "
"head_age_s and the prefetch-gating sync requirement.)"
)
assert not (
self.enable_nsa_prefill_cp_shared_kv
and self.disaggregation_mode == "prefill"