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,

test/registered/spec/eagle/test_eagle_infer_b.py

@@ -316,6 +316,12 @@ class TestEAGLEServerPageSize(TestEAGLEServerBasic):
         "--attention-backend=flashinfer",
     ]
 
+    @classmethod
+    def setUpClass(cls):
+        # Runtime check only supported for topk=1, and can help to find a leak.
+        with envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY.override(1):
+            super().setUpClass()
+
 
 class TestEAGLEServerPageSizeTopk(TestEAGLEServerBasic):
     # default topk=8 and tokens=64

test/registered/spec/eagle/test_eagle_infer_beta.py

@@ -57,7 +57,9 @@ class TestEagleServerBase(CustomTestCase, MatchedStopMixin):
             *[str(i) for i in range(1, cls.max_running_requests + 1)],
         ]
         launch_args.extend(cls.other_launch_args)
-        with envs.SGLANG_ENABLE_SPEC_V2.override(True):
+        with envs.SGLANG_ENABLE_SPEC_V2.override(
+            True
+        ), envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY.override(1):
             cls.process = popen_launch_server(
                 cls.model,
                 cls.base_url,

test/srt/test_radix_attention.py

@@ -84,4 +84,5 @@ class TestRadixCacheNonOverlapLPM(TestRadixCacheFCFS):
 
 if __name__ == "__main__":
     envs.SGLANG_TEST_RETRACT.set(True)
+    envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY.set(1)
     unittest.main()

test/srt/test_retract_decode.py

@@ -27,7 +27,9 @@ class TestRetractDecode(CustomTestCase):
         cls.model = DEFAULT_MODEL_NAME_FOR_TEST
         cls.base_url = DEFAULT_URL_FOR_TEST
         launch_args = ["--chunked-prefill-size", "128"] + cls.other_args
-        with envs.SGLANG_TEST_RETRACT.override(True):
+        with envs.SGLANG_TEST_RETRACT.override(
+            True
+        ), envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY.override(1):
             cls.process = popen_launch_server(
                 cls.model,
                 cls.base_url,