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:
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user