diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index edb8874d3..7a30de4da 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -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) diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index 32c7642ab..5c0fb7326 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -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) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index d82141228..5a56e8fbc 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -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") diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index 87a2488c5..96ba46651 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -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): """