Enhance runtime memory check in CI (#15192)

python/sglang/srt/managers/scheduler_runtime_checker_mixin.py

@@ -5,6 +5,7 @@ import signal
 import sys
 import threading
 import time
+import warnings
 from typing import TYPE_CHECKING
 
 import psutil
@@ -175,6 +176,13 @@ class SchedulerRuntimeCheckerMixin:
         if current_batch is None:
             return
 
+        spec_topk = self.server_args.speculative_eagle_topk or 1
+        if spec_topk > 1:
+            warnings.warn(
+                "Runtime memory check (busy) is not supported when speculation topk > 1."
+            )
+            return
+
         _, _, available_size, evictable_size = self._get_token_info()
         protected_size = self.tree_cache.protected_size()
 

python/sglang/srt/mem_cache/chunk_cache.py

@@ -51,14 +51,11 @@ class ChunkCache(BasePrefixCache):
         ]
         self.req_to_token_pool.free(req.req_pool_idx)
         self.token_to_kv_pool_allocator.free(kv_indices)
-        self.protected_size_ -= len(req.prefix_indices)
 
     def cache_unfinished_req(self, req: Req, chunked=False):
         kv_indices = self.req_to_token_pool.req_to_token[
             req.req_pool_idx, : len(req.fill_ids)
         ]
-        self.protected_size_ += len(kv_indices) - len(req.prefix_indices)
-
         # `req.prefix_indices` will be used in `PrefillAdder::add_chunked_req` later
         req.prefix_indices = kv_indices.to(dtype=torch.int64, copy=True)
 
@@ -72,7 +69,8 @@ class ChunkCache(BasePrefixCache):
         return 0
 
     def protected_size(self):
-        return self.protected_size_
+        # NOTE: no protected size in chunk cache. Chunk cache's eviction is the same with request's lifecycle.
+        return 0
 
     def pretty_print(self):
         return ""

python/sglang/srt/speculative/eagle_info.py

@@ -114,8 +114,6 @@ class EagleVerifyInput(SpecInput, EagleVerifyInputV2Mixin):
                 len(batch.input_ids),
             )
             end_offset = batch.seq_lens + self.draft_token_num
-            for req in batch.reqs:
-                req.kv_allocated_len += 1
         else:
             prefix_lens = batch.seq_lens
             prefix_lens_cpu = batch.seq_lens_cpu

python/sglang/srt/speculative/eagle_worker.py

@@ -382,8 +382,6 @@ class EAGLEWorker(TpModelWorker):
         # [       topk 0         ] [       topk 1         ]
         # [iter=0, iter=1, iter=2] [iter=0, iter=1, iter=2]
         if self.page_size == 1:
-            for req in batch.reqs:
-                req.kv_allocated_len += self.speculative_num_steps * self.topk
             # TODO: We only need self.speculative_num_steps - 1 * topk cache loc
             out_cache_loc, token_to_kv_pool_state_backup = alloc_token_slots(
                 batch.tree_cache,