fix(sglang): account snapshot-reserved slots in radix mem leak check

Phase 2 prepare_receive allocates kv_pool slots that aren't visible
to radix / session bookkeeping until finalize_ingest. Without this
fix, the scheduler's idle self_check fires:

  ValueError: token_to_kv_pool_allocator memory leak detected!
    available=288391, evictable=5, protected=0, session_held=0
    (expected sum == 288460)

_check_radix_cache_memory now subtracts
  sum(len(rec.slot_indices) for rec in ctrl._ingest_records.values())
from the expected total before flagging a leak. Snapshot_reserved is
also printed in the leak message for diagnostics.

Smoke confirmed (scripts/smoke_snapshot_sglang_integration.py):
  [smoke] prepare_receive on P → 200: ok=true (96 layer bufs)
  [smoke] dump on D → 200: ok=false, reason=session-not-resident
  [smoke] finalize on P → 200: ok=true, inserted_prefix_len=0
  [smoke] OVERALL: PASS

End-to-end KV-correctness (snapshot ingest yields cache hit on next
prefill) still requires the agentic+router stack — covered in the E4
sweep, not this smoke.
This commit is contained in:
Claude Code Agent
2026-05-13 08:26:16 +08:00
parent b9b0cf0fac
commit a369722efe
2 changed files with 39 additions and 57 deletions

View File

@@ -184,10 +184,25 @@ class SchedulerRuntimeCheckerMixin:
_, _, available_size, evictable_size = self._get_token_info()
protected_size = self.tree_cache.protected_size()
session_held = self._session_held_tokens()
# Snapshot link prepare_receive reserves slots that aren't yet visible
# to radix / session bookkeeping until finalize_ingest. Count them so
# the leak check doesn't fire while a snapshot ingest is in-flight.
snapshot_reserved = 0
ctrl = getattr(self, "snapshot_link_controller", None)
if ctrl is not None:
try:
snapshot_reserved = sum(
len(rec.slot_indices) for rec in ctrl._ingest_records.values()
)
except Exception:
snapshot_reserved = 0
memory_leak = (available_size + evictable_size) != (
self.max_total_num_tokens - protected_size - session_held
self.max_total_num_tokens - protected_size - session_held - snapshot_reserved
)
token_msg = (
f"{self.max_total_num_tokens=}, {available_size=}, {evictable_size=}, "
f"{protected_size=}, {session_held=}, {snapshot_reserved=}\n"
)
token_msg = f"{self.max_total_num_tokens=}, {available_size=}, {evictable_size=}, {protected_size=}, {session_held=}\n"
return memory_leak, token_msg
def _get_batch_uncached_size(self: Scheduler, batch: ScheduleBatch) -> int: