From ecb401ed42cb73a604a050f302a5134b03e83186 Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Tue, 16 Dec 2025 21:40:38 +0800 Subject: [PATCH] Enhance runtime memory check in CI (#15192) --- .../srt/managers/scheduler_runtime_checker_mixin.py | 8 ++++++++ python/sglang/srt/mem_cache/chunk_cache.py | 6 ++---- python/sglang/srt/speculative/eagle_info.py | 2 -- python/sglang/srt/speculative/eagle_worker.py | 2 -- test/registered/spec/eagle/test_eagle_infer_b.py | 6 ++++++ test/registered/spec/eagle/test_eagle_infer_beta.py | 4 +++- test/srt/test_radix_attention.py | 1 + test/srt/test_retract_decode.py | 4 +++- 8 files changed, 23 insertions(+), 10 deletions(-) diff --git a/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py b/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py index b6fa4bde3..25e3d3429 100644 --- a/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py +++ b/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() diff --git a/python/sglang/srt/mem_cache/chunk_cache.py b/python/sglang/srt/mem_cache/chunk_cache.py index dcc899e46..a543a62d1 100644 --- a/python/sglang/srt/mem_cache/chunk_cache.py +++ b/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 "" diff --git a/python/sglang/srt/speculative/eagle_info.py b/python/sglang/srt/speculative/eagle_info.py index 75fc1070f..1bf3816e9 100644 --- a/python/sglang/srt/speculative/eagle_info.py +++ b/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 diff --git a/python/sglang/srt/speculative/eagle_worker.py b/python/sglang/srt/speculative/eagle_worker.py index 239e8136b..b84c8ed72 100644 --- a/python/sglang/srt/speculative/eagle_worker.py +++ b/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, diff --git a/test/registered/spec/eagle/test_eagle_infer_b.py b/test/registered/spec/eagle/test_eagle_infer_b.py index 60a02195a..10c4f6774 100644 --- a/test/registered/spec/eagle/test_eagle_infer_b.py +++ b/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 diff --git a/test/registered/spec/eagle/test_eagle_infer_beta.py b/test/registered/spec/eagle/test_eagle_infer_beta.py index f79c0de6a..b11ee31af 100644 --- a/test/registered/spec/eagle/test_eagle_infer_beta.py +++ b/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, diff --git a/test/srt/test_radix_attention.py b/test/srt/test_radix_attention.py index c8c3f539e..2f501e882 100644 --- a/test/srt/test_radix_attention.py +++ b/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() diff --git a/test/srt/test_retract_decode.py b/test/srt/test_retract_decode.py index da688b400..0b0f17bfb 100644 --- a/test/srt/test_retract_decode.py +++ b/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,