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

E1 lane-stats instrumentation + 2a reserve-at-admission (default-off)

Browse Source 
Two CP HiCache observability/admission additions, both behind default-off env flags (behavior-neutral until enabled).

E1 instrumentation (SGLANG_CP_HICACHE_LANE_STATS_S=<sec>, default 0):
- environ.py: flag. hiradix_cache.py: _cp_maybe_log_lane_stats() (rank-0 per-owner-lane L2 occupancy + imbalance + hot-prefix histogram) wired into check_hicache_events; per-cause drop counters in _warn_cp_hicache_fallback.
- Used for E1: verdict = lane imbalance is a non-issue (1.02-1.03 across full fill->evict->refill, zero host_reservation_failed).

2a reserve-at-admission (SGLANG_CP_HICACHE_RESERVE_ADMISSION, default off):
- hiradix_cache.py: cp_host_backup_admission_budget() (min-over-ranks available+evictable host tokens from the replicated radix snapshot) + _cp_host_evictable_tokens_by_rank.
- schedule_policy.py: PrefillAdder rem_l2_backup_tokens field + worst-lane footprint debit + budget_state gate (OTHER).
- scheduler.py: pass cp_size to PrefillAdder.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
孙远航
2026-06-20 12:03:14 +00:00
parent 7de882b622
commit 17cbfa31c9
4 changed files with 183 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -382,6 +382,13 @@ class Envs:
SGLANG_HICACHE_NIXL_BACKEND_STORAGE_DIR = EnvStr(None)
# Max fraction of cache (by token count) that can be pinned; 0 = disable pinning.
SGLANG_HICACHE_MAX_PINNED_RATIO = EnvFloat(0.0)
# E1 probe: emit per-owner-lane L2 (host) occupancy + imbalance + backup-drop
# counters every N seconds (0 = off). Read-only CP HiCache lane-imbalance probe.
SGLANG_CP_HICACHE_LANE_STATS_S = EnvFloat(0.0)
# Phase 2a: gate reserve-at-admission L2 backup-headroom budget. Default off
# during bring-up so E1 measures the baseline (reactive drops, not waits);
# flip on once validated (E3). Opt-in.
SGLANG_CP_HICACHE_RESERVE_ADMISSION = EnvBool(False)
# Mooncake KV Transfer
SGLANG_MOONCAKE_CUSTOM_MEM_POOL = EnvStr(None)

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

@@ -478,8 +478,10 @@ class PrefillAdder:
] = None,
prefill_delayer_single_pass: Optional[PrefillDelayerSinglePassExecutor] = None,
dllm_config: Optional[DllmConfig] = None,
cp_size: int = 1,
):
self.page_size = page_size
self.cp_size = max(1, int(cp_size or 1))
self.tree_cache = tree_cache
self.token_to_kv_pool_allocator = token_to_kv_pool_allocator
self.running_batch = running_batch
@@ -592,6 +594,35 @@ class PrefillAdder:
CPSharedKVPrefillBufferEstimate
] = None
# Reserve-at-admission (Phase 2a): scalar host (L2) backup headroom =
# min over CP ranks of per-rank (available + evictable) host tokens,
# snapshotted once per adder (once per tick). Debited per admitted backup
# by its worst-lane footprint; admission stops (backpressure) at <= 0.
# None unless under CP shared-KV HiCache. Replicated -> identical on every
# rank, so the admit/skip decision stays collective-free.
self.rem_l2_backup_tokens: Optional[int] = None
if (
self._is_cp_prefill_context()
and envs.SGLANG_CP_HICACHE_RESERVE_ADMISSION.get()
and hasattr(self.tree_cache, "cp_host_backup_admission_budget")
):
budget = self.tree_cache.cp_host_backup_admission_budget()
if budget is not None:
self.rem_l2_backup_tokens = int(budget)
def _cp_host_backup_footprint_tokens(self, extend_input_len_paged: int) -> int:
"""Worst-lane per-rank host-backup footprint for an extend of
``extend_input_len_paged`` (already page-aligned) tokens. The backup's
pages round-robin across ``cp_size`` owner lanes, so the busiest lane
receives at most ceil(num_pages / cp_size) pages. Debiting this against
the min-over-ranks budget is conservative (never under-counts the tightest
lane), which is the safe direction for an admission gate -- the exact
per-lane reservation still happens at prepare.
"""
num_pages = max(0, extend_input_len_paged) // self.page_size
per_rank_pages = (num_pages + self.cp_size - 1) // self.cp_size
return per_rank_pages * self.page_size
def _init_dllm_meta(self, dllm_config: DllmConfig):
self.dllm_block_size = dllm_config.block_size
max_running_reqs = dllm_config.max_running_requests
@@ -808,6 +839,13 @@ class PrefillAdder:
if self.rem_chunk_tokens is not None and self.rem_chunk_tokens <= 0:
return AddReqResult.OTHER
# Reserve-at-admission (Phase 2a): stop admitting when the host (L2) backup
# headroom is exhausted so a complete backup never fails reactively at
# prepare. OTHER (not NO_TOKEN): a deferred req re-enters the scan next tick
# (= backpressure/wait) without forcing batch_is_full.
if self.rem_l2_backup_tokens is not None and self.rem_l2_backup_tokens <= 0:
return AddReqResult.OTHER
return AddReqResult.CONTINUE
def _update_prefill_budget(
@@ -823,6 +861,12 @@ class PrefillAdder:
self.rem_input_tokens -= extend_input_len
if self._is_cp_prefill_context():
self.cp_shared_kv_prefill_total_extend_tokens += extend_input_len
if self.rem_l2_backup_tokens is not None:
# Debit this backup's worst-lane host footprint (extend_input_len is
# already page-aligned above) against the reserve-at-admission budget.
self.rem_l2_backup_tokens -= self._cp_host_backup_footprint_tokens(
extend_input_len
)
if self.cp_shared_kv_prefill_max_buffer_size is not None:
self.cp_shared_kv_prefill_estimate_prefix_lens.append(prefix_len)
self.cp_shared_kv_prefill_estimate_extend_lens.append(extend_input_len)

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

@@ -2554,6 +2554,7 @@ class Scheduler(
),
prefill_delayer_single_pass=prefill_delayer_single_pass,
dllm_config=self.dllm_config,
cp_size=self.attn_cp_size,
)
record_cp_timing_stage("create_adder")

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

@@ -935,6 +935,48 @@ class HiRadixCache(RadixCache):
snapshot.target_capacity, snapshot.target_used
)
def _cp_host_evictable_tokens_by_rank(self) -> Tuple[int, ...]:
"""Per-rank host (L2) tokens reclaimable by evicting all currently
plannable host leaves -- the host analog of tree_cache.evictable_size().
Computed from the REPLICATED radix tree (evictable_host_leaves + per-node
owner counts), so it is identical on every CP rank (collective-free).
"""
cp_size = self._cp_hicache_cp_size()
totals = self._cp_zero_counts(cp_size)
for node in getattr(self, "evictable_host_leaves", set()):
if not self._cp_host_leaf_is_plannable_victim(node):
continue
metadata = getattr(node, "cp_hicache", None)
if metadata is None:
continue
counts = self._cp_assert_metadata_counts(
metadata, context="host_evictable_tokens"
)
totals = self._cp_add_counts(totals, counts)
return totals
def cp_host_backup_admission_budget(self) -> Optional[int]:
"""Reserve-at-admission (Phase 2a) host headroom: min over CP ranks of
per-rank (available + evictable) host tokens.
A scalar, conservative budget -- the tightest lane's effective L2 capacity
-- mirroring the device-side ``rem_total_tokens = available_size() +
evictable_size()``. The scheduler debits each admitted backup's worst-lane
footprint and stops admitting (backpressure) when this hits zero, so a
complete backup never fails reactively at prepare. Returns None when not
under CP shared-KV HiCache. Derived purely from the replicated radix
snapshot -> identical on every rank (collective-free).
"""
if not self._uses_cp_hicache:
return None
snapshot = self._cp_host_capacity_snapshot()
available = self._cp_host_available_tokens_by_rank(snapshot)
evictable = self._cp_host_evictable_tokens_by_rank()
cp_size = min(len(available), len(evictable))
if cp_size == 0:
return None
return min(int(available[r]) + int(evictable[r]) for r in range(cp_size))
def _cp_host_evict_key(self, node: TreeNode):
# 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
@@ -2363,6 +2405,13 @@ class HiRadixCache(RadixCache):
def _warn_cp_hicache_fallback(
self, fallback_name: str, reason: str, *, rate_limit_s: float = 0.0, **kwargs
):
# E1 probe: count every fallback/drop by name BEFORE rate-limiting, so the
# true drop rate is captured even when the warning log is suppressed.
fallback_counts = getattr(self, "_cp_fallback_counts", None)
if fallback_counts is None:
fallback_counts = {}
self._cp_fallback_counts = fallback_counts
fallback_counts[fallback_name] = fallback_counts.get(fallback_name, 0) + 1
details = " ".join(f"{key}={value}" for key, value in kwargs.items())
if details:
details = " " + details
@@ -2396,6 +2445,87 @@ class HiRadixCache(RadixCache):
details,
)
def _cp_maybe_log_lane_stats(self) -> None:
"""E1 lane-imbalance probe (read-only, measurement only).
Every ``SGLANG_CP_HICACHE_LANE_STATS_S`` seconds, log the per-owner-lane L2
(host) occupancy vector, the cross-lane imbalance (max/min ratio + spread),
a hot-prefix length histogram (does the hot working set skew owners?), and
the backup-drop counters accumulated since the last emit. The host capacity
snapshot is derived from the REPLICATED radix tree, so every CP rank
computes the identical 8-lane vector; emit only on cp_rank 0 to avoid a
cp_size-fold of duplicate lines. Tests the load-bearing Phase-2 premise:
is per-rank L2 saturation (lane imbalance) real, or just global pressure?
"""
interval = float(envs.SGLANG_CP_HICACHE_LANE_STATS_S.get() or 0.0)
if interval <= 0.0 or not self._uses_cp_hicache:
return
if self._cp_hicache_cp_rank() != 0:
return
now = time.monotonic()
last = getattr(self, "_cp_lane_stats_last_emit", 0.0)
if last and (now - last) < interval:
return
self._cp_lane_stats_last_emit = now
snapshot = self._cp_host_capacity_snapshot()
capacity = snapshot.target_capacity
used = snapshot.target_used
cp_size = len(capacity)
occ = [
(used[i] / capacity[i]) if capacity[i] > 0 else 0.0
for i in range(cp_size)
]
max_occ = max(occ) if occ else 0.0
min_occ = min(occ) if occ else 0.0
spread = max_occ - min_occ
imbalance = (max_occ / min_occ) if min_occ > 1e-9 else float("inf")
# Hot-prefix length histogram: bucket reused (hit_count>=2) backed-up nodes
# at the owner round-robin period (cp_size * page_size tokens). A prefix
# shorter than one period cannot span all lanes -> it skews ownership; a
# working set dominated by such short hot prefixes is what makes per-rank
# pools imbalance. Long prefixes round-robin evenly and do not skew.
skew_token_threshold = cp_size * self.page_size
hot_short = 0
hot_long = 0
for node in self._cp_walk_radix_nodes():
if int(getattr(node, "host_len", 0) or 0) <= 0:
continue
if int(getattr(node, "hit_count", 0) or 0) < 2:
continue
if int(getattr(node, "host_len", 0)) < skew_token_threshold:
hot_short += 1
else:
hot_long += 1
fallback_counts = getattr(self, "_cp_fallback_counts", {})
last_counts = getattr(self, "_cp_fallback_counts_last", {})
drops_since_last = {
name: fallback_counts[name] - last_counts.get(name, 0)
for name in fallback_counts
if fallback_counts[name] - last_counts.get(name, 0) > 0
}
self._cp_fallback_counts_last = dict(fallback_counts)
logger.info(
"[CP_HICACHE_LANE_STATS] occ=%s cap_tokens=%s used_tokens=%s max=%.3f "
"min=%.3f spread=%.3f imbalance=%s hot_prefix_short=%d hot_prefix_long=%d "
"skew_token_threshold=%d drops_since_last=%s drops_total=%s",
[round(x, 3) for x in occ],
list(capacity),
list(used),
max_occ,
min_occ,
spread,
("inf" if imbalance == float("inf") else round(imbalance, 2)),
hot_short,
hot_long,
skew_token_threshold,
drops_since_last,
dict(fallback_counts),
)
def _probe_existing_radix_prefix_len_no_split(self, key: RadixKey) -> int:
"""Return the already-present radix prefix length without mutating the tree.
@@ -3923,6 +4053,7 @@ class HiRadixCache(RadixCache):
self.storage_metrics_collector.log_storage_metrics(
self.cache_controller.storage_backend.get_stats()
)
self._cp_maybe_log_lane_stats()
def drain_storage_control_queues(self):
"""