Fix SWA eviction tombstoning the last leaf (#29860)
This commit is contained in:
@@ -881,7 +881,6 @@ class Envs:
|
||||
SGLANG_OPT_SWA_RADIX_CACHE_COMPACT = EnvBool(False)
|
||||
SGLANG_OPT_SWA_SPLIT_LEAF_ON_INSERT = EnvBool(False)
|
||||
SGLANG_OPT_SWA_RELEASE_LEAF_LOCK_AFTER_WINDOW = EnvBool(False)
|
||||
SGLANG_OPT_SWA_EVICT_DROP_PAGE_MARGIN = EnvBool(False)
|
||||
|
||||
# Unified radix cache
|
||||
SGLANG_OPT_UNIFIED_CACHE_FREE_OUT_OF_WINDOW_SLOTS = EnvBool(False)
|
||||
@@ -1037,6 +1036,7 @@ def _convert_SGL_to_SGLANG():
|
||||
"SGLANG_ENABLE_TP_MEMORY_INBALANCE_CHECK",
|
||||
)
|
||||
_print_deprecated_env("SGLANG_PER_TOKEN_GROUP_QUANT_8BIT_V2")
|
||||
_print_deprecated_env("SGLANG_OPT_SWA_EVICT_DROP_PAGE_MARGIN")
|
||||
_print_deprecated_env("SGLANG_ENABLE_THINKING", "SGLANG_DEFAULT_THINKING")
|
||||
_print_deprecated_env("SGLANG_REASONING_EFFORT", "SGLANG_DSV4_REASONING_EFFORT")
|
||||
_print_deprecated_env(
|
||||
|
||||
@@ -2930,7 +2930,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
|
||||
page_size=self.tree_cache.page_size,
|
||||
req_to_token_pool=self.req_to_token_pool,
|
||||
token_to_kv_pool_allocator=self.token_to_kv_pool_allocator,
|
||||
drop_page_margin=self.tree_cache.is_chunk_cache(),
|
||||
is_chunk_cache=self.tree_cache.is_chunk_cache(),
|
||||
)
|
||||
|
||||
def __str__(self):
|
||||
|
||||
@@ -73,10 +73,8 @@ def free_swa_out_of_window_slots(
|
||||
page_size: int,
|
||||
req_to_token_pool: ReqToTokenPool,
|
||||
token_to_kv_pool_allocator: BaseTokenToKVPoolAllocator,
|
||||
drop_page_margin: bool = False,
|
||||
is_chunk_cache: bool = False,
|
||||
) -> None:
|
||||
from sglang.srt.environ import envs
|
||||
|
||||
# For swa radix cache, we need to evict the tokens that are not in the tree cache and also not in the sliding window
|
||||
assert (
|
||||
req.cache_protected_len % page_size == 0
|
||||
@@ -86,13 +84,16 @@ def free_swa_out_of_window_slots(
|
||||
evict_floor = -(-evict_floor // page_size) * page_size
|
||||
req.swa_evicted_seqlen = max(req.swa_evicted_seqlen, evict_floor)
|
||||
|
||||
# Subtract an extra page_size so the eviction frontier never reaches the
|
||||
# radix tree insert boundary, keeping >=1 page of non-evicted SWA KV for the
|
||||
# tree to store as a non-tombstone node (else leaf nodes tombstone -> SWA leak).
|
||||
if drop_page_margin or envs.SGLANG_OPT_SWA_EVICT_DROP_PAGE_MARGIN.get():
|
||||
if is_chunk_cache:
|
||||
# Chunk cache builds no radix tree, so no tombstone-leaf concern; evict
|
||||
# up to the window boundary (the trailing floor keeps it page-aligned).
|
||||
evict_threshold = pre_len - sliding_window_size
|
||||
else:
|
||||
evict_threshold = pre_len - sliding_window_size - page_size
|
||||
# Radix cache: keep max(window, page). The trailing floor page-aligns the
|
||||
# frontier, and subtracting at least one page keeps it below the insert
|
||||
# boundary (page_floor(seq_len)) so the last leaf is never all-tombstone.
|
||||
# No extra page margin is needed.
|
||||
evict_threshold = pre_len - max(sliding_window_size, page_size)
|
||||
new_swa_evicted_seqlen = max(
|
||||
req.swa_evicted_seqlen,
|
||||
evict_threshold,
|
||||
|
||||
Reference in New Issue
Block a user