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

[unified-memory] Stop eviction when shared allocation capacity is sufficient (#33091)

Browse Source 
Co-authored-by: seokwoosong <seokwoosong@users.noreply.github.com>
This commit is contained in:
SuperSong
2026-08-26 02:06:32 -07:00
committed by GitHub
co-authored by seokwoosong
parent ffc431cd4c
commit 2511743bd7
22 changed files with 885 additions and 123 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/disaggregation/decode.py
+6 -2
View File
@@ -427,7 +427,9 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
required = ceil_align(swa_tail_len, page_size)
available = self.token_to_kv_pool_allocator.swa_available_size()
if available < required:
self.tree_cache.evict(EvictParams(swa_num_tokens=required - available))
self.tree_cache.evict_for_alloc(
EvictParams(swa_num_tokens=required - available)
)
available = self.token_to_kv_pool_allocator.swa_available_size()
if available < required:
@@ -1789,7 +1791,9 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
and self._radix_full_available() < required_alloc_tokens
):
num_to_evict = required_alloc_tokens - self._radix_full_available()
result = self.tree_cache.evict(EvictParams(num_tokens=num_to_evict))
result = self.tree_cache.evict_for_alloc(
EvictParams(num_tokens=num_to_evict)
)
if self._radix_full_available() < required_alloc_tokens:
logger.warning(
f"Eviction insufficient: needed {required_alloc_tokens} tokens, "
python/sglang/srt/hardware_backend/mlx/kv_cache/auxiliary_state.py
+29 -2
View File
@@ -10,7 +10,7 @@ storing model-agnostic native cache snapshots.
from __future__ import annotations
from dataclasses import dataclass
from typing import Any, Iterable, Optional
from typing import Any, Iterable, Iterator, Optional
import mlx.core as mx
import torch
@@ -95,6 +95,7 @@ class MlxAuxiliaryStatePool:
self.mamba_cache = None
self.mem_usage = 0
self._snapshots: dict[int, dict[int, _CacheSnapshot]] = {}
self._alloc_iter: Optional[Iterator[torch.Tensor]] = None
self.clear()
def _tensor(self, indices: Any) -> torch.Tensor:
@@ -108,7 +109,31 @@ class MlxAuxiliaryStatePool:
def available_size(self) -> int:
return int(self.free_slots.numel())
def schedulable_available_size(self) -> int:
return self.available_size()
def alloc_group_begin(self, num_reqs: int) -> None:
self._alloc_iter = None
if num_reqs > 0:
slots = self._do_alloc(num_reqs)
if slots is not None:
self._alloc_iter = iter(slots.split(1))
def alloc_group_end(self) -> None:
if self._alloc_iter is not None:
remaining = list(self._alloc_iter)
if remaining:
self.free(torch.cat(remaining))
self._alloc_iter = None
def alloc(self, need_size: int) -> Optional[torch.Tensor]:
if self._alloc_iter is not None and need_size == 1:
slot = next(self._alloc_iter, None)
if slot is not None:
return slot
return self._do_alloc(need_size)
def _do_alloc(self, need_size: int) -> Optional[torch.Tensor]:
if need_size > self.available_size():
return None
slots = self.free_slots[:need_size].clone()
@@ -128,6 +153,7 @@ class MlxAuxiliaryStatePool:
self.free_slots = torch.cat([self.free_slots, indices])
def clear(self) -> None:
self._alloc_iter = None
self.free_slots = torch.arange(
1, self.size + 1, dtype=torch.int64, device=self.device
)
@@ -227,6 +253,7 @@ class MlxAuxiliaryStateReqToTokenPool(ReqToTokenPool):
size=auxiliary_state_size,
device=device,
)
self.mamba_allocator = self.mamba_pool
# The unified radix base MAMBA component still reads ``mamba_pool``.
# Keep the MLX-owned name beside it so local code can avoid model-
# specific terminology.
@@ -352,7 +379,7 @@ class MlxAuxiliaryStateComponent(MambaComponent):
source_value
)
if forked_value is None:
self.cache.evict(EvictParams(num_tokens=0, mamba_num=1))
self.cache.evict_for_alloc(EvictParams(num_tokens=0, mamba_num=1))
forked_value = (
self.cache.req_to_token_pool.auxiliary_state_pool.fork_from(
source_value
python/sglang/srt/mem_cache/allocation.py
+3 -1
View File
@@ -257,7 +257,9 @@ def alloc_req_slots(
if mamba_available_size < mamba_state_needed:
if tree_cache is not None and tree_cache.supports_mamba():
mamba_num = max(0, mamba_state_needed - mamba_available_size)
tree_cache.evict(EvictParams(num_tokens=0, mamba_num=mamba_num))
tree_cache.evict_for_alloc(
EvictParams(num_tokens=0, mamba_num=mamba_num)
)
req_pool_indices = req_to_token_pool.alloc(reqs)
if req_pool_indices is None:
raise RuntimeError(
python/sglang/srt/mem_cache/base_prefix_cache.py
+10
View File
@@ -324,6 +324,16 @@ class BasePrefixCache(ABC, PrefixCacheTrait):
def evict(self, params: EvictParams) -> EvictResult:
pass
def evict_for_alloc(self, params: EvictParams) -> EvictResult:
"""Evict cache entries to cover allocator shortfalls.
The default implementation preserves the component-count semantics of
:meth:`evict`. Multi-component caches backed by shared memory can
override this entry point to stop once collateral frees make the
requested allocation feasible.
"""
return self.evict(params)
@abstractmethod
def inc_lock_ref(self, node: Any) -> IncLockRefResult:
pass
python/sglang/srt/mem_cache/buffer_mode/pipeline.py
+6 -2
View File
@@ -832,8 +832,12 @@ class BufferModePipeline:
avail = cache.token_to_kv_pool_allocator.available_size()
if avail < f.num_tokens:
needed = f.num_tokens - avail
evicted = cache.evict(EvictParams(num_tokens=needed))
if evicted.num_tokens_evicted < needed:
cache.evict_for_alloc(EvictParams(num_tokens=needed))
if cache.supports_swa():
avail = cache.token_to_kv_pool_allocator.full_available_size()
else:
avail = cache.token_to_kv_pool_allocator.available_size()
if avail < f.num_tokens:
# Genuinely no room (locked pages): recompute.
return _drop()
python/sglang/srt/mem_cache/common.py
+4 -2
View File
@@ -130,14 +130,16 @@ 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(
tree_cache.evict_for_alloc(
EvictParams(num_tokens=full_num_tokens, swa_num_tokens=swa_num_tokens)
)
else:
# Standard allocator: evict only the shortfall (mirrors the SWA arm)
available_size = allocator.available_size()
if available_size < num_tokens:
tree_cache.evict(EvictParams(num_tokens=num_tokens - available_size))
tree_cache.evict_for_alloc(
EvictParams(num_tokens=num_tokens - available_size)
)
def retraction_backup(
python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py
+25 -13
View File
@@ -48,6 +48,26 @@ def _get_allocator_type(server_args: ServerArgs) -> str:
return get_allocator_type(server_args)
def _evict_swa_for_device_alloc(cache: UnifiedRadixCache, required_size: int) -> None:
from sglang.srt.mem_cache.base_prefix_cache import EvictParams
available_size = cache.token_to_kv_pool_allocator.swa_available_size()
shortfall = max(0, required_size - available_size)
if shortfall > 0:
cache.evict_for_alloc(EvictParams(swa_num_tokens=shortfall))
def _evict_mamba_for_device_alloc(cache: UnifiedRadixCache, required_size: int) -> None:
from sglang.srt.mem_cache.base_prefix_cache import EvictParams
available_size = (
cache.req_to_token_pool.mamba_allocator.schedulable_available_size()
)
shortfall = max(0, required_size - available_size)
if shortfall > 0:
cache.evict_for_alloc(EvictParams(mamba_num=shortfall))
def _make_layer_mapper(
layer_mapping: dict[int, int],
transfer_layer_num: int,
@@ -1210,8 +1230,6 @@ class _DeepSeekV4Strategy(StackStrategy):
model_name=None,
enable_storage_metrics=False,
):
from sglang.srt.mem_cache.base_prefix_cache import EvictParams
host_pool_group, cache_controller = build_deepseek_v4_hicache_stack(
params=params,
server_args=server_args,
@@ -1219,7 +1237,7 @@ class _DeepSeekV4Strategy(StackStrategy):
load_cache_event=load_cache_event,
storage_backend=storage_backend,
host_swa_evict_fn=lambda n: cache.evict_host(n, ComponentType.SWA),
device_swa_evict_fn=lambda n: cache.evict(EvictParams(swa_num_tokens=n)),
device_swa_evict_fn=lambda n: _evict_swa_for_device_alloc(cache, n),
prefetch_threshold=prefetch_threshold,
model_name=model_name,
storage_backend_extra_config=storage_backend_extra_config,
@@ -1286,8 +1304,6 @@ class _MambaStrategy(StackStrategy):
model_name=None,
enable_storage_metrics=False,
):
from sglang.srt.mem_cache.base_prefix_cache import EvictParams
full_layer_mapping = dict(kvcache.full_attention_layer_id_mapping)
mamba_layer_mapping = dict(params.req_to_token_pool.mamba_map)
host_pool_group, cache_controller = build_hybrid_mamba_stack(
@@ -1301,7 +1317,7 @@ class _MambaStrategy(StackStrategy):
storage_backend=storage_backend,
use_mla=kvcache.use_mla,
host_mamba_evict_fn=lambda n: cache.evict_host(n, ComponentType.MAMBA),
device_mamba_evict_fn=lambda n: cache.evict(EvictParams(mamba_num=n)),
device_mamba_evict_fn=lambda n: _evict_mamba_for_device_alloc(cache, n),
prefetch_threshold=prefetch_threshold,
model_name=model_name,
storage_backend_extra_config=storage_backend_extra_config,
@@ -1355,8 +1371,6 @@ class _SwaStrategy(StackStrategy):
model_name=None,
enable_storage_metrics=False,
):
from sglang.srt.mem_cache.base_prefix_cache import EvictParams
full_layer_mapping, swa_layer_mapping = _swa_layer_mappings(kvcache)
host_pool_group, cache_controller = build_hybrid_swa_stack(
params=params,
@@ -1369,7 +1383,7 @@ class _SwaStrategy(StackStrategy):
storage_backend=storage_backend,
use_mla=False,
host_swa_evict_fn=lambda n: cache.evict_host(n, ComponentType.SWA),
device_swa_evict_fn=lambda n: cache.evict(EvictParams(swa_num_tokens=n)),
device_swa_evict_fn=lambda n: _evict_swa_for_device_alloc(cache, n),
prefetch_threshold=prefetch_threshold,
model_name=model_name,
storage_backend_extra_config=storage_backend_extra_config,
@@ -1417,8 +1431,6 @@ class _MambaSwaStrategy(StackStrategy):
model_name=None,
enable_storage_metrics=False,
):
from sglang.srt.mem_cache.base_prefix_cache import EvictParams
full_layer_mapping, swa_layer_mapping = _swa_layer_mappings(kvcache)
mamba_layer_mapping = dict(params.req_to_token_pool.mamba_map)
host_pool_group, cache_controller = build_hybrid_mamba_swa_stack(
@@ -1438,9 +1450,9 @@ class _MambaSwaStrategy(StackStrategy):
pp_group=params.pp_cache_group,
storage_backend=storage_backend,
host_swa_evict_fn=lambda n: cache.evict_host(n, ComponentType.SWA),
device_swa_evict_fn=lambda n: cache.evict(EvictParams(swa_num_tokens=n)),
device_swa_evict_fn=lambda n: _evict_swa_for_device_alloc(cache, n),
host_mamba_evict_fn=lambda n: cache.evict_host(n, ComponentType.MAMBA),
device_mamba_evict_fn=lambda n: cache.evict(EvictParams(mamba_num=n)),
device_mamba_evict_fn=lambda n: _evict_mamba_for_device_alloc(cache, n),
prefetch_threshold=prefetch_threshold,
model_name=model_name,
storage_backend_extra_config=storage_backend_extra_config,
python/sglang/srt/mem_cache/pool_host/group.py
+2
View File
@@ -16,6 +16,8 @@ class PoolEntry:
device_pool: Any
layer_mapper: Callable[[int], int | None]
is_primary_index_anchor: bool = False
# Reclaim callbacks receive the absolute allocation size n. The host
# callback evicts n slots; the device callback makes alloc(n) feasible.
host_evict_fn: Callable[[int], Any] | None = None
device_evict_fn: Callable[[int], Any] | None = None
device_alloc_fn: Callable[[int], Any] | None = None
python/sglang/srt/mem_cache/unified_cache/components/mamba_component.py
+39 -33
View File
@@ -203,7 +203,7 @@ class MambaComponent(TreeComponent):
# stops at this request's window boundary instead of walking to
# root and over-decrementing locks held by other requests.
lock_result = self.cache.inc_lock_ref(result.best_match_node)
self.cache.evict(EvictParams(num_tokens=0, mamba_num=1))
self.cache.evict_for_alloc(EvictParams(num_tokens=0, mamba_num=1))
dst_index = self.cache.req_to_token_pool.mamba_allocator.alloc(1)
self.cache.dec_lock_ref(
result.best_match_node, lock_result.to_dec_params()
@@ -374,10 +374,14 @@ class MambaComponent(TreeComponent):
device_frees: dict[ComponentType, list[torch.Tensor]],
host_frees: dict[ComponentType, list[torch.Tensor]],
) -> Optional[NodeId]:
"""Return the next device-leaf node for the driver to evict, or None.
Internal nodes are tombstoned inline (no IO). If the previous node's
eviction removed the cursor, the walk resumes from the partition
sentinel with session refs on, else it restarts at the LRU tail."""
"""Advance one device-eviction step and return a leaf, if selected.
An internal tombstone is one complete step so the caller can apply its
pending frees and recheck allocator capacity before the next mutation.
If the previous node's eviction removed the cursor, the walk resumes
from the partition sentinel with session refs on, else it restarts at
the LRU tail.
"""
ct = self.component_type
lru = self.tree_core.lru_lists[ct]
enabled = self.tree_core.enable_session_radix_cache
@@ -387,34 +391,36 @@ class MambaComponent(TreeComponent):
self._evict_device_cursor = (
lru.cursor_next() if enabled else lru.get_lru_no_lock()
)
while (
tracker[ct] < self._evict_device_request_cnt
and self._evict_device_cursor is not None
and lru.in_list(self._evict_device_cursor)
if (
tracker[ct] >= self._evict_device_request_cnt
or self._evict_device_cursor is None
or not lru.in_list(self._evict_device_cursor)
):
x = self._evict_device_cursor
assert x.component_data[ct].value is not None
if x in self.tree_core.evictable_device_leaves and (
not enabled or self._can_evict_leaf_atomically(x)
):
self._evict_device_cursor = (
lru.cursor_next() if enabled else lru.get_prev_no_lock(x)
)
return x.id
if not enabled:
x_next = lru.get_prev_no_lock(x)
self.tree_core._evict_component_and_detach_lru(
x,
self,
target=EvictLayer.DEVICE,
tracker=tracker,
device_frees=device_frees,
host_frees=host_frees,
return None
x = self._evict_device_cursor
assert x.component_data[ct].value is not None
if x in self.tree_core.evictable_device_leaves and (
not enabled or self._can_evict_leaf_atomically(x)
):
self._evict_device_cursor = (
lru.cursor_next() if enabled else lru.get_prev_no_lock(x)
)
self.tree_core._cascade_evict(
x, self, tracker, device_frees=device_frees, host_frees=host_frees
)
self._evict_device_cursor = lru.cursor_next() if enabled else x_next
return x.id
if not enabled:
x_next = lru.get_prev_no_lock(x)
self.tree_core._evict_component_and_detach_lru(
x,
self,
target=EvictLayer.DEVICE,
tracker=tracker,
device_frees=device_frees,
host_frees=host_frees,
)
self.tree_core._cascade_evict(
x, self, tracker, device_frees=device_frees, host_frees=host_frees
)
self._evict_device_cursor = lru.cursor_next() if enabled else x_next
return None
def _evict_device_end(self) -> None:
@@ -487,7 +493,7 @@ class MambaComponent(TreeComponent):
"""Allocate one mamba pool slot, evicting if necessary."""
slot = self.cache.req_to_token_pool.mamba_allocator.alloc(1)
if slot is None:
self.cache.evict(EvictParams(num_tokens=0, mamba_num=1))
self.cache.evict_for_alloc(EvictParams(num_tokens=0, mamba_num=1))
slot = self.cache.req_to_token_pool.mamba_allocator.alloc(1)
assert slot is not None, "Can not alloc mamba cache"
return slot
@@ -660,7 +666,7 @@ class MambaComponent(TreeComponent):
return PrepareLoadBackResult()
dst = self.cache.req_to_token_pool.mamba_allocator.alloc(1)
if dst is None:
self.cache.evict(EvictParams(num_tokens=0, mamba_num=1))
self.cache.evict_for_alloc(EvictParams(num_tokens=0, mamba_num=1))
dst = self.cache.req_to_token_pool.mamba_allocator.alloc(1)
assert dst is not None, "Cannot alloc mamba for load_back"
req.mamba_pool_idx = dst[0]
python/sglang/srt/mem_cache/unified_cache/components/swa_component.py
+36 -30
View File
@@ -535,10 +535,14 @@ class SWAComponent(TreeComponent):
device_frees: dict[ComponentType, list[torch.Tensor]],
host_frees: dict[ComponentType, list[torch.Tensor]],
) -> Optional[NodeId]:
"""Return the next device-leaf node for the driver to evict, or None.
Internal nodes are tombstoned inline (no IO). If the previous node's
eviction removed the cursor, the walk resumes from the partition
sentinel with session refs on, else it restarts at the LRU tail."""
"""Advance one device-eviction step and return a leaf, if selected.
An internal tombstone is one complete step so the caller can apply its
pending frees and recheck allocator capacity before the next mutation.
If the previous node's eviction removed the cursor, the walk resumes
from the partition sentinel with session refs on, else it restarts at
the LRU tail.
"""
ct = self.component_type
lru = self.tree_core.lru_lists[ct]
enabled = self.tree_core.enable_session_radix_cache
@@ -548,34 +552,36 @@ class SWAComponent(TreeComponent):
self._evict_device_cursor = (
lru.cursor_next() if enabled else lru.get_lru_no_lock()
)
while (
tracker[ct] < self._evict_device_request_cnt
and self._evict_device_cursor is not None
and lru.in_list(self._evict_device_cursor)
if (
tracker[ct] >= self._evict_device_request_cnt
or self._evict_device_cursor is None
or not lru.in_list(self._evict_device_cursor)
):
x = self._evict_device_cursor
assert x.component_data[ct].value is not None
if x in self.tree_core.evictable_device_leaves and (
not enabled or self._can_evict_leaf_atomically(x)
):
self._evict_device_cursor = (
lru.cursor_next() if enabled else lru.get_prev_no_lock(x)
)
return x.id
if not enabled:
x_next = lru.get_prev_no_lock(x)
self.tree_core._evict_component_and_detach_lru(
x,
self,
target=EvictLayer.DEVICE,
tracker=tracker,
device_frees=device_frees,
host_frees=host_frees,
return None
x = self._evict_device_cursor
assert x.component_data[ct].value is not None
if x in self.tree_core.evictable_device_leaves and (
not enabled or self._can_evict_leaf_atomically(x)
):
self._evict_device_cursor = (
lru.cursor_next() if enabled else lru.get_prev_no_lock(x)
)
self.tree_core._cascade_evict(
x, self, tracker, device_frees=device_frees, host_frees=host_frees
)
self._evict_device_cursor = lru.cursor_next() if enabled else x_next
return x.id
if not enabled:
x_next = lru.get_prev_no_lock(x)
self.tree_core._evict_component_and_detach_lru(
x,
self,
target=EvictLayer.DEVICE,
tracker=tracker,
device_frees=device_frees,
host_frees=host_frees,
)
self.tree_core._cascade_evict(
x, self, tracker, device_frees=device_frees, host_frees=host_frees
)
self._evict_device_cursor = lru.cursor_next() if enabled else x_next
return None
def _evict_device_end(self) -> None:
python/sglang/srt/mem_cache/unified_cache/components/tree_component.py
+6 -3
View File
@@ -507,8 +507,11 @@ class TreeComponent(ABC):
device_frees: dict[ComponentType, list[torch.Tensor]],
host_frees: dict[ComponentType, list[torch.Tensor]],
) -> Optional[NodeId]:
"""Return the next device-leaf node for the driver to evict, or None.
Internal nodes are tombstoned inline (no IO)."""
"""Advance one eviction step and return a device leaf, if selected.
Implementations must return after one allocator-relevant internal
mutation so the caller can drain pending frees before continuing.
"""
assert (
self.is_evict_device_ongoing
), f"{self.component_type} device eviction not started"
@@ -534,7 +537,7 @@ class TreeComponent(ABC):
device_frees: dict[ComponentType, list[torch.Tensor]],
host_frees: dict[ComponentType, list[torch.Tensor]],
) -> Optional[NodeId]:
"""Advance the walk; return the next device leaf or None."""
"""Advance the walk by at most one allocator-relevant mutation."""
...
@abstractmethod
python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py
+2 -1
View File
@@ -1224,7 +1224,7 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
def evict_device_next_node(
self, component_type: ComponentType, tracker: dict[ComponentType, int]
) -> EvictDeviceNextNodeResult:
"""Return the next device leaf to evict for a component, or None when done."""
"""Advance one component eviction step and report whether it progressed."""
result = EvictDeviceNextNodeResult()
# The walk reads running totals for its doneness check; the result
# carries only this step's delta.
@@ -1236,6 +1236,7 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
delta = n - tracker.get(ct, 0)
if delta:
result.tracker[ct] = delta
result.made_progress = result.node_id is not None or bool(result.tracker)
return result
def evict_device_end(self, component_type: ComponentType) -> None:
python/sglang/srt/mem_cache/unified_cache/unified_tree_core_interface.py
+13 -2
View File
@@ -41,7 +41,15 @@ class BaseEvictionResult(msgspec.Struct):
class EvictDeviceNextNodeResult(BaseEvictionResult):
"""One device-walk step.
``node_id`` selects a leaf for the Controller to evict. ``made_progress``
also covers an internal tombstone that returned no leaf, distinguishing it
from true walk exhaustion.
"""
node_id: Optional[NodeId] = None
made_progress: bool = False
class EvictDeviceLeafResult(BaseEvictionResult):
@@ -230,8 +238,11 @@ class UnifiedTreeCoreInterface(ABC):
def evict_device_next_node(
self, component_type: ComponentType, tracker: dict[ComponentType, int]
) -> EvictDeviceNextNodeResult:
"""The next evictable node (None node_id when the walk is exhausted);
tracker is the caller's running totals, read for the doneness check."""
"""Advance one eviction step.
A missing ``node_id`` is exhausted only when ``made_progress`` is also
false. ``tracker`` is the caller's running totals, read for doneness.
"""
...
@abstractmethod
python/sglang/srt/mem_cache/unified_radix_cache.py
+83 -17
View File
@@ -531,18 +531,68 @@ class UnifiedRadixCache(BasePrefixCache):
self._apply_cache_actions(self.tree_core.end_insert())
def evict(self, params: EvictParams) -> EvictResult:
return self._evict(params)
def evict_for_alloc(self, params: EvictParams) -> EvictResult:
"""Evict until the requested component allocations become feasible.
``params`` contains allocator shortfalls, not absolute eviction quotas.
A component eviction can cascade to its peers; with a shared memory pool,
those collateral frees can satisfy the original allocation before the
triggering component's requested count is reached.
"""
if self.disable:
return EvictResult()
start_time = time.perf_counter()
tracker = {ct: 0 for ct in self.tree_components}
request_by_type = {
request_by_type = self._evict_request_by_type(params)
available_size_targets = {
ct: self._component_available_size(ct) + request_cnt
for ct, request_cnt in request_by_type.items()
if request_cnt > 0
}
return self._evict(params, available_size_targets)
@staticmethod
def _evict_request_by_type(params: EvictParams) -> dict[ComponentType, int]:
return {
ComponentType.FULL: params.num_tokens,
ComponentType.SWA: params.swa_num_tokens,
ComponentType.MAMBA: params.mamba_num,
ComponentType.C128: 0,
}
self._evict_components(request_by_type, tracker)
def _component_available_size(self, component_type: ComponentType) -> int:
"""Return capacity usable by the component's next allocation.
Shared allocators expose schedulable capacity, which includes peer holes
that an urgent allocator flush can reclaim without further eviction.
"""
if component_type == ComponentType.FULL:
if self.supports_swa():
return self.token_to_kv_pool_allocator.full_available_size()
return self.token_to_kv_pool_allocator.available_size()
if component_type == ComponentType.SWA:
return self.token_to_kv_pool_allocator.swa_available_size()
if component_type == ComponentType.MAMBA:
return self.req_to_token_pool.mamba_allocator.schedulable_available_size()
raise ValueError(f"Unsupported cache component: {component_type}")
def _evict(
self,
params: EvictParams,
available_size_targets: Optional[dict[ComponentType, int]] = None,
) -> EvictResult:
if self.disable:
return EvictResult()
start_time = time.perf_counter()
tracker = {ct: 0 for ct in self.tree_components}
request_by_type = self._evict_request_by_type(params)
self._evict_components(
request_by_type,
tracker,
available_size_targets=available_size_targets,
)
if (
self.cache_controller is not None
@@ -581,12 +631,12 @@ class UnifiedRadixCache(BasePrefixCache):
def _evict_device_next_node(
self, component_type: ComponentType, tracker: dict[ComponentType, int]
) -> Optional[NodeId]:
) -> tuple[Optional[NodeId], bool]:
"""Advance the eviction walk one node, consuming its step result."""
result = self.tree_core.evict_device_next_node(component_type, tracker)
self._free_values(result.device_frees, result.host_frees)
self._accumulate_tracker(tracker, result.tracker)
return result.node_id
return result.node_id, result.made_progress
def _evict_device_leaf(
self, node_id: NodeId, tracker: dict[ComponentType, int]
@@ -617,20 +667,39 @@ class UnifiedRadixCache(BasePrefixCache):
self,
request_by_type: dict[ComponentType, int],
tracker: dict[ComponentType, int],
available_size_targets: Optional[dict[ComponentType, int]] = None,
) -> None:
# Buffer mode: eviction always wins over queued backup intents — a
# destroyed victim's intent is stale-swept and the content rewrites
# after its recompute.
def target_reached(component_type: ComponentType) -> bool:
if available_size_targets is None:
return False
target = available_size_targets.get(component_type)
# Do not compact on every eviction step. Shared allocators include
# drainable peer holes here and flush the peer once in alloc().
return (
target is not None
and self._component_available_size(component_type) >= target
)
for ct in self.tree_components:
request_cnt = request_by_type[ct]
# Skip eviction walk if request is already met
if tracker[ct] >= request_cnt:
# A preceding component may have cascade-evicted this component or,
# on a shared pool, released enough bytes to satisfy its allocation.
if tracker[ct] >= request_cnt or target_reached(ct):
continue
self.tree_core.evict_device_start(ct, request_cnt)
try:
while (
node_id := self._evict_device_next_node(ct, tracker)
) is not None:
while not target_reached(ct):
node_id, made_progress = self._evict_device_next_node(ct, tracker)
if node_id is None:
if made_progress:
# Internal tombstone frees are now allocator-visible;
# recheck the allocation target before walking again.
continue
break
backup_kv = self._evict_device_leaf(node_id, tracker)
if backup_kv is not None:
# Deferred demote: run the D->H backup, demote only on success.
@@ -1395,14 +1464,11 @@ class UnifiedRadixCache(BasePrefixCache):
self.dec_host_lock_ref(node_id, host_anchor_params)
return False
if self.supports_swa():
avail = self.token_to_kv_pool_allocator.full_available_size()
else:
avail = self.token_to_kv_pool_allocator.available_size()
avail = self._component_available_size(ComponentType.FULL)
if avail < kv_tokens:
needed = kv_tokens - avail
result = self.evict(EvictParams(num_tokens=needed))
if result.num_tokens_evicted < needed:
self.evict_for_alloc(EvictParams(num_tokens=needed))
if self._component_available_size(ComponentType.FULL) < kv_tokens:
self.dec_lock_ref(node_id, ancestor_lock_params)
self.dec_host_lock_ref(node_id, host_anchor_params)
return False
python/sglang/srt/session/streaming_session.py
+3
View File
@@ -417,6 +417,9 @@ class StreamingSession(BasePrefixCache):
def evict(self, params: EvictParams) -> EvictResult:
return self.inner.evict(params)
def evict_for_alloc(self, params: EvictParams) -> EvictResult:
return self.inner.evict_for_alloc(params)
def inc_lock_ref(self, node: Any) -> IncLockRefResult:
result = self.try_inc_lock_ref(node)
if result is not None: