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

fix: preserve HiCache load-back allocator headroom

Browse Source
This commit is contained in:
谢骞
2026-05-08 00:18:33 +08:00
parent 1f074f434e
commit 95bacb8862
10 changed files with 172 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
python/sglang/srt/managers/schedule_policy.py
View File

@@ -499,6 +499,15 @@ class PrefillAdder:
def ceil_paged_tokens(self, tokens: int) -> int:
return -(-tokens // self.page_size) * self.page_size
def _get_available_device_tokens_for_load_back(self) -> int:
if self.is_hybrid_swa:
return self.token_to_kv_pool_allocator.full_available_size()
return self.token_to_kv_pool_allocator.available_size()
def _get_load_back_mem_quota(self, real_input_tokens: int) -> int:
reserve_tokens = real_input_tokens + self.page_size
return max(self._get_available_device_tokens_for_load_back() - reserve_tokens, 0)
def budget_state(self):
if self.rem_total_tokens <= 0 or self.cur_rem_tokens <= 0:
return AddReqResult.NO_TOKEN
@@ -777,6 +786,7 @@ class PrefillAdder:
InitLoadBackParams(
last_host_node=req.last_host_node,
host_hit_length=req.host_hit_length,
mem_quota=self._get_load_back_mem_quota(real_input_tokens),
)
)
req.prefix_indices = torch.cat([req.prefix_indices, new_indices])

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

@@ -69,6 +69,38 @@ class BaseTokenToKVPoolAllocator(abc.ABC):
def available_size(self):
return (len(self.free_pages) + len(self.release_pages)) * self.page_size
def immediate_available_pages(self) -> int:
if self.free_pages is None:
return 0
return len(self.free_pages)
def deferred_available_pages(self) -> int:
if self.release_pages is None:
return 0
return len(self.release_pages)
def immediate_available_size(self) -> int:
if self.free_pages is None:
return self.available_size()
return self.immediate_available_pages() * self.page_size
def deferred_available_size(self) -> int:
if self.release_pages is None:
return 0
return self.deferred_available_pages() * self.page_size
def allocator_state_str(self) -> str:
if self.free_pages is None:
return f"allocator_available_size={self.available_size()}"
return (
"allocator_available_size="
f"{self.available_size()} "
f"(allocator_free_size={self.immediate_available_size()} "
f"[free_pages={self.immediate_available_pages()}] + "
f"allocator_release_size={self.deferred_available_size()} "
f"[release_pages={self.deferred_available_pages()}])"
)
def get_kvcache(self):
return self._kvcache
@@ -434,10 +466,15 @@ class PagedTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
)
bs = len(prefix_lens)
if self.need_sort and extend_num_tokens // self.page_size + bs + 1 > len(
self.free_pages
):
num_new_pages = get_num_new_pages(
seq_lens=seq_lens_cpu,
page_size=self.page_size,
prefix_lens=prefix_lens_cpu,
)
if self.need_sort and num_new_pages > len(self.free_pages):
self.merge_and_sort_free()
if num_new_pages > len(self.free_pages):
return None
out_indices = torch.empty(
(extend_num_tokens,), dtype=torch.int64, device=self.device
@@ -456,14 +493,6 @@ class PagedTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
if self.debug_mode:
assert len(torch.unique(out_indices)) == len(out_indices)
num_new_pages = get_num_new_pages(
seq_lens=seq_lens_cpu,
page_size=self.page_size,
prefix_lens=prefix_lens_cpu,
)
if num_new_pages > len(self.free_pages):
return None
self.free_pages = self.free_pages[num_new_pages:]
return out_indices
@@ -479,8 +508,15 @@ class PagedTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
)
bs = len(seq_lens)
if self.need_sort and bs > len(self.free_pages):
num_new_pages = get_num_new_pages(
seq_lens=seq_lens_cpu,
page_size=self.page_size,
decode=True,
)
if self.need_sort and num_new_pages > len(self.free_pages):
self.merge_and_sort_free()
if num_new_pages > len(self.free_pages):
return None
out_indices = torch.empty((bs,), dtype=torch.int64, device=self.device)
alloc_decode_kernel[(bs,)](
@@ -495,14 +531,6 @@ class PagedTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
if self.debug_mode:
assert len(torch.unique(out_indices)) == len(out_indices)
num_new_pages = get_num_new_pages(
seq_lens=seq_lens_cpu,
page_size=self.page_size,
decode=True,
)
if num_new_pages > len(self.free_pages):
return None
self.free_pages = self.free_pages[num_new_pages:]
return out_indices

12
python/sglang/srt/mem_cache/base_prefix_cache.py
View File

@@ -268,6 +268,14 @@ class BasePrefixCache(ABC, PrefixCacheTrait):
return not self.is_chunk_cache()
def available_and_evictable_str(self) -> str:
available_size = self.token_to_kv_pool_allocator.available_size()
allocator = self.token_to_kv_pool_allocator
allocator_state_str = getattr(allocator, "allocator_state_str", None)
if allocator_state_str is None:
allocator_state = f"allocator_available_size={allocator.available_size()}"
else:
allocator_state = allocator_state_str()
evictable_size = self.evictable_size()
return f"Available tokens: {available_size + evictable_size} ({available_size=} + {evictable_size=})\n"
return (
f"Available tokens: {allocator.available_size() + evictable_size} "
f"({allocator_state} + {evictable_size=})\n"
)

37
python/sglang/srt/mem_cache/common.py
View File

@@ -7,7 +7,11 @@ import torch
import triton
import triton.language as tl
from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache, EvictParams
from sglang.srt.mem_cache.base_prefix_cache import (
BasePrefixCache,
EvictParams,
EvictResult,
)
from sglang.srt.mem_cache.cp_shared_kv_compute_owner import (
build_in_seq_page_compute_owners,
get_in_seq_page_compute_owner_unavailable_reason,
@@ -220,7 +224,7 @@ def alloc_token_slots(
backup_state: bool = False,
):
allocator = tree_cache.token_to_kv_pool_allocator
evict_from_tree_cache(tree_cache, num_tokens)
evict_result = evict_from_tree_cache(tree_cache, num_tokens)
state = None
if backup_state:
@@ -233,6 +237,7 @@ def alloc_token_slots(
f"Out of memory. Try to lower your batch size.\n"
f"Try to allocate {num_tokens} tokens.\n"
f"{available_and_evictable_str(tree_cache)}"
f"{_evict_result_str(evict_result)}\n"
)
logger.error(error_msg)
if tree_cache is not None:
@@ -242,12 +247,25 @@ def alloc_token_slots(
return (out_cache_loc, state) if backup_state else out_cache_loc
def evict_from_tree_cache(tree_cache: BasePrefixCache | None, num_tokens: int):
def _evict_result_str(evict_result: EvictResult | None) -> str:
if evict_result is None:
return "evict_result=None"
return (
"evict_result=("
f"num_tokens_evicted={evict_result.num_tokens_evicted}, "
f"swa_num_tokens_evicted={evict_result.swa_num_tokens_evicted}, "
f"mamba_num_evicted={evict_result.mamba_num_evicted})"
)
def evict_from_tree_cache(
tree_cache: BasePrefixCache | None, num_tokens: int
) -> EvictResult:
if tree_cache is None:
return
return EvictResult()
if tree_cache.is_chunk_cache():
return
return EvictResult()
allocator = tree_cache.token_to_kv_pool_allocator
@@ -259,13 +277,15 @@ def evict_from_tree_cache(tree_cache: BasePrefixCache | None, num_tokens: int):
if full_available_size < num_tokens or swa_available_size < num_tokens:
full_num_tokens = max(0, num_tokens - full_available_size)
swa_num_tokens = max(0, num_tokens - swa_available_size)
tree_cache.evict(
return tree_cache.evict(
EvictParams(num_tokens=full_num_tokens, swa_num_tokens=swa_num_tokens)
)
return EvictResult()
else:
# Standard allocator
if allocator.available_size() < num_tokens:
tree_cache.evict(EvictParams(num_tokens=num_tokens))
return tree_cache.evict(EvictParams(num_tokens=num_tokens))
return EvictResult()
def _evict_for_compute_owner_lanes(
@@ -322,7 +342,7 @@ def alloc_paged_token_slots_extend(
# Over estimate the number of tokens: assume each request needs a new page.
allocator = tree_cache.token_to_kv_pool_allocator
num_tokens = extend_num_tokens + len(seq_lens_cpu) * allocator.page_size
evict_from_tree_cache(tree_cache, num_tokens)
evict_result = evict_from_tree_cache(tree_cache, num_tokens)
alloc_extend_compute_owner = getattr(
allocator, "alloc_extend_compute_owner", None
@@ -448,6 +468,7 @@ def alloc_paged_token_slots_extend(
f"Prefill out of memory. Try to lower your batch size.\n"
f"Try to allocate {extend_num_tokens} tokens.\n"
f"{available_and_evictable_str(tree_cache)}"
f"{_evict_result_str(evict_result)}\n"
)
logger.error(error_msg)
if tree_cache is not None:

3
python/sglang/srt/mem_cache/mamba_radix_cache.py
View File

@@ -20,7 +20,6 @@ The radix tree data structure for managing the hybrid (full and Mamba) KV cache.
"""
import heapq
from collections import defaultdict
from functools import lru_cache, partial
from typing import TYPE_CHECKING, List, Optional, Tuple
@@ -69,7 +68,7 @@ class TreeNode:
last_access_time_counter_float = float64(1.0)
def __init__(self, id: Optional[int] = None):
self.children = defaultdict(TreeNode)
self.children = {}
self.parent: TreeNode = None
self.key: RadixKey = None
self.value: Optional[torch.Tensor] = None

3
python/sglang/srt/mem_cache/radix_cache.py
View File

@@ -26,7 +26,6 @@ import heapq
import logging
import sys
import time
from collections import defaultdict
from functools import lru_cache, partial
from typing import TYPE_CHECKING, Any, Iterator, List, Optional, Tuple, Union
@@ -123,7 +122,7 @@ class TreeNode:
counter = 0
def __init__(self, id: Optional[int] = None, priority: int = 0):
self.children = defaultdict(TreeNode)
self.children = {}
self.parent: TreeNode = None
self.key: RadixKey = None
self.value: Optional[torch.Tensor] = None

3
python/sglang/srt/mem_cache/swa_radix_cache.py
View File

@@ -21,7 +21,6 @@ The radix tree data structure for managing the hybrid (full and SWA) KV cache.
import heapq
import time
from collections import defaultdict
from functools import partial
from typing import TYPE_CHECKING, List, Optional, Tuple
@@ -67,7 +66,7 @@ class TreeNode:
last_access_time_counter_float = float64(1.0)
def __init__(self, id: Optional[int] = None):
self.children = defaultdict(TreeNode)
self.children = {}
self.parent: TreeNode = None
self.key: RadixKey = None
self.value: Optional[torch.Tensor] = None