chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Further fix memory pool leak error (#9298)

Browse Source
This commit is contained in:
fzyzcjy
2025-08-18 15:58:06 +08:00
committed by GitHub
parent 968e181826
commit 4c0bb411e5
2 changed files with 7 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
python/sglang/srt/mem_cache/allocator.py
View File

@@ -434,15 +434,12 @@ class PagedTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
device: str,
kvcache: KVCache,
need_sort: bool,
max_num_extend_tokens: int,
):
super().__init__(size, page_size, dtype, device, kvcache, need_sort)
self.num_pages = size // page_size
self.max_num_extend_tokens_next_power_of_2 = next_power_of_2(
max_num_extend_tokens
)
self.debug_mode = get_bool_env_var("SGLANG_DEBUG_MEMORY_POOL")
self.ret_values = torch.empty((), dtype=torch.int64, device=self.device)
self.seen_max_num_extend_tokens_next_power_of_2 = 1
self.clear()
def alloc(self, need_size: int):
@@ -480,17 +477,17 @@ class PagedTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
(last_loc + 1) % self.page_size == prefix_lens % self.page_size
)
self.seen_max_num_extend_tokens_next_power_of_2 = max(
self.seen_max_num_extend_tokens_next_power_of_2,
next_power_of_2(extend_num_tokens),
)
bs = len(prefix_lens)
if self.need_sort and extend_num_tokens // self.page_size + bs + 1 > len(
self.free_pages
):
self.merge_and_sort_free()
assert self.max_num_extend_tokens_next_power_of_2 >= extend_num_tokens, (
f"{self.max_num_extend_tokens_next_power_of_2=} >= {extend_num_tokens=} does not hold. "
f"If this happens in PD, consider letting chunked_prefill_size in D be as large as in P"
)
out_indices = torch.empty(
(extend_num_tokens,), dtype=torch.int64, device=self.device
)
@@ -503,7 +500,7 @@ class PagedTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
self.ret_values,
next_power_of_2(bs),
self.page_size,
self.max_num_extend_tokens_next_power_of_2,
self.seen_max_num_extend_tokens_next_power_of_2,
)
if self.debug_mode: