From 22b209ad6ea05121b231583366370e411dea0978 Mon Sep 17 00:00:00 2001 From: yhzhuang Date: Tue, 25 Aug 2026 18:02:19 -0700 Subject: [PATCH 01/55] Support unified memory with speculative decoding Co-authored-by: Yonghao Zhuang --- python/sglang/srt/mem_cache/allocator/swa.py | 10 +++ .../srt/mem_cache/kv_cache_configurator.py | 85 +++++++++++-------- .../srt/mem_cache/multi_ended_allocator.py | 34 ++++++-- .../srt/model_executor/pool_configurator.py | 80 +++++++++++++++-- 4 files changed, 161 insertions(+), 48 deletions(-) diff --git a/python/sglang/srt/mem_cache/allocator/swa.py b/python/sglang/srt/mem_cache/allocator/swa.py index 785f9b83a30c..d2c8bed5eb2c 100644 --- a/python/sglang/srt/mem_cache/allocator/swa.py +++ b/python/sglang/srt/mem_cache/allocator/swa.py @@ -343,6 +343,16 @@ def set_full_to_swa_mapping( swa_indices = swa_indices.to(self.full_to_swa_index_mapping.dtype) self.full_to_swa_index_mapping[full_indices] = swa_indices + def recover_swa_with_locked_full( + self, kept_full: torch.Tensor, incoming_full: torch.Tensor + ) -> torch.Tensor: + """Keep locked FULL slots and transfer the incoming SWA slots to them.""" + swa_value = self.translate_loc_from_full_to_swa(incoming_full) + self.set_full_to_swa_mapping(kept_full, swa_value) + self.clear_full_to_swa_mapping(incoming_full) + self.full_attn_allocator.free(incoming_full) + return swa_value + def clear_full_to_swa_mapping(self, full_indices: torch.Tensor) -> None: if full_indices.numel() == 0: return diff --git a/python/sglang/srt/mem_cache/kv_cache_configurator.py b/python/sglang/srt/mem_cache/kv_cache_configurator.py index 98634cc8d220..7c40094db503 100644 --- a/python/sglang/srt/mem_cache/kv_cache_configurator.py +++ b/python/sglang/srt/mem_cache/kv_cache_configurator.py @@ -64,6 +64,10 @@ PageMajorMHATokenToKVPool, ReqToTokenPool, ) +from sglang.srt.mem_cache.multi_ended_allocator import ( + UnifiedMambaTokenToKVPoolAllocator, + UnifiedSWATokenToKVPoolAllocator, +) from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool from sglang.srt.platforms import current_platform from sglang.srt.runtime_context import ( @@ -482,35 +486,43 @@ def _init_pools( # pool must be sized by that space. draft_virtual_id_space: Optional[int] = None if self.is_draft_worker and token_to_kv_pool_allocator is not None: - from sglang.srt.mem_cache.multi_ended_allocator import ( - UnifiedMambaTokenToKVPoolAllocator, - UnifiedSWATokenToKVPoolAllocator, - ) - - if isinstance(token_to_kv_pool_allocator, UnifiedSWATokenToKVPoolAllocator): - raise ValueError( - "Speculative decoding with --enable-unified-memory is only " - "supported for hybrid-Mamba targets; the unified hybrid-SWA " - "pool's draft sizing (virtual-id space) is not wired yet." - ) if isinstance( - token_to_kv_pool_allocator, UnifiedMambaTokenToKVPoolAllocator + token_to_kv_pool_allocator, + ( + UnifiedMambaTokenToKVPoolAllocator, + UnifiedSWATokenToKVPoolAllocator, + ), ): - draft_virtual_id_space = token_to_kv_pool_allocator.size_full + draft_virtual_id_space = ( + token_to_kv_pool_allocator.draft_virtual_id_space + ) assert draft_virtual_id_space >= sizes.max_total_num_tokens, ( "unified allocator virtual space smaller than the token " - f"budget: size_full={draft_virtual_id_space} < " + f"budget: virtual_id_space={draft_virtual_id_space} < " f"max_total_num_tokens={sizes.max_total_num_tokens}" ) # Round UP to page alignment (paged draft backends view the - # pool as (-1, page_size, H, D); size_full is not aligned). + # pool as (-1, page_size, H, D); the virtual space is not aligned). page = max(int(self.pool_page_size or 1), 1) draft_virtual_id_space = ( (draft_virtual_id_space + page - 1) // page * page ) - sizes = msgspec.structs.replace( - sizes, max_total_num_tokens=draft_virtual_id_space - ) + size_overrides = { + "max_total_num_tokens": draft_virtual_id_space, + } + if ( + isinstance( + token_to_kv_pool_allocator, + UnifiedSWATokenToKVPoolAllocator, + ) + and self.is_hybrid_swa + ): + size_overrides["full_max_total_num_tokens"] = draft_virtual_id_space + if not self.is_inkling_mtp_draft or self.draft_swa_full_capacity: + size_overrides["swa_max_total_num_tokens"] = ( + draft_virtual_id_space + ) + sizes = msgspec.structs.replace(sizes, **size_overrides) # Initialize req_to_token_pool if req_to_token_pool is None: @@ -558,7 +570,8 @@ def _init_pools( assert token_to_kv_pool.size >= draft_virtual_id_space, ( "draft token_to_kv_pool smaller than the shared unified " f"allocator's virtual-id space: pool size=" - f"{token_to_kv_pool.size} < size_full={draft_virtual_id_space}; " + f"{token_to_kv_pool.size} < " + f"virtual_id_space={draft_virtual_id_space}; " "verify-window writes at high virtual ids would go out of " "bounds." ) @@ -1960,27 +1973,31 @@ def _build_token_to_kv_pool_allocator( else: assert self.is_draft_worker if self.is_hybrid_swa: - if self.draft_swa_full_capacity: - # Banded depth: the SWA ring is full draft capacity, so use - # an IDENTITY full->swa mapping — store and read locs both - # equal out_cache_loc, and a slot is never evicted before - # the request frees it. The window itself is enforced by the - # FA sliding-window kernel, not by the ring. Layout mirrors - # SWATokenToKVPoolAllocator's mapping (size + page_size - # entries + trailing -1 sentinel so a -1 last_loc maps - # to -1). + swa_allocator = getattr( + token_to_kv_pool_allocator, + "logical_attn_allocator", + token_to_kv_pool_allocator, + ) + uses_unified_virtual_ids = isinstance( + swa_allocator, UnifiedSWATokenToKVPoolAllocator + ) + has_draft_swa_layers = ( + not self.is_inkling_mtp_draft or self.draft_swa_full_capacity + ) + if self.draft_swa_full_capacity or ( + uses_unified_virtual_ids and has_draft_swa_layers + ): + # The draft pool owns independent KV but consumes the target + # allocator's virtual ids directly. Size its SWA side for that + # whole space and use an identity mapping. The trailing -1 + # sentinel keeps a -1 last_loc mapped to -1. n = sizes.full_max_total_num_tokens + self.page_size identity_mapping = torch.arange( n + 1, dtype=torch.int64, device=self.device ) identity_mapping[-1] = -1 token_to_kv_pool.register_mapping(identity_mapping) - else: - swa_allocator = getattr( - token_to_kv_pool_allocator, - "logical_attn_allocator", - token_to_kv_pool_allocator, - ) + elif not uses_unified_virtual_ids: assert isinstance(swa_allocator, SWATokenToKVPoolAllocator) token_to_kv_pool.register_mapping( swa_allocator.full_to_swa_index_mapping diff --git a/python/sglang/srt/mem_cache/multi_ended_allocator.py b/python/sglang/srt/mem_cache/multi_ended_allocator.py index 06e026af79a8..606a3db89e28 100644 --- a/python/sglang/srt/mem_cache/multi_ended_allocator.py +++ b/python/sglang/srt/mem_cache/multi_ended_allocator.py @@ -269,6 +269,7 @@ def __init__( forward_stream: Optional[torch.cuda.Stream] = None, lazy_compaction: bool = False, kernel_page_multiplier: Optional[int] = None, + virtual_num_pages: Optional[int] = None, ): spec = unified_buffer.spec(sub_pool_name) max_slots = unified_buffer.max_slots(sub_pool_name) @@ -323,15 +324,28 @@ def __init__( self.num_pages = max_slots // self.pool_page_size # `min_page_index` = ceil(min_slot_index / pool_page_size), keeping the # reserved-sink invariant (min_page_index * entry_bytes_per_page >= entry_max). + assert ( + virtual_num_pages is None or not is_id_owner + ), "only a non-owner allocator may use another pool's virtual-id space" + # An ID owner has one virtual page per physical page. A non-owner may + # deliberately have a different physical capacity while indexing the + # owner's virtual page space (the unified SWA allocator does this). + self.num_virtual_ids = ( + virtual_num_pages if virtual_num_pages is not None else self.num_pages + ) + if is_id_owner: + assert self.num_virtual_ids == self.num_pages + assert self.num_virtual_ids > 0, "virtual page count must be positive" self.min_page_index = ( self.min_slot_index + self.pool_page_size - 1 ) // self.pool_page_size self.entry_bytes_per_page = self.entry_bytes * self.pool_page_size - # v2p / p2v sized by PAGES. Page 0 is the padding anchor; trailing row is - # the -1 sentinel. + # v2p follows the shared virtual-id owner; p2v follows this pool's + # physical capacity. Page 0 is the padding anchor; the last row is a + # sentinel. self.virtual_to_physical = torch.full( - (self.num_pages + 1,), + (self.num_virtual_ids + 1,), -1, dtype=torch.int64, device=device, @@ -342,8 +356,6 @@ def __init__( dtype=torch.int64, device=device, ) - # Back-compat alias (count of virtual PAGES) consulted by is_slot_allocated. - self.num_virtual_ids = self.num_pages # Chain neighbours: `low_peer` toward byte 0, `high_peer` toward # `total_bytes`. Ends have one (`bind_peer`), float middles have both. @@ -536,7 +548,7 @@ def allocated_count(self) -> int: def is_slot_allocated(self, slot: int) -> bool: """Whether the PAGE containing this virtual id is in use.""" virt_page = slot // self.page_size - if virt_page < 0 or virt_page >= self.num_pages: + if virt_page < 0 or virt_page >= self.num_virtual_ids: return False return int(self.virtual_to_physical[virt_page].item()) != -1 @@ -546,6 +558,7 @@ def allocator_state_str(self) -> str: f"is_id_owner={self.is_id_owner}, page_size={self.page_size}, " f"min_page_index={self.min_page_index}, " f"num_pages={self.num_pages}, " + f"num_virtual_ids={self.num_virtual_ids}, " f"watermark_physical={self.watermark_physical}, " f"allocated_pages={self._allocated_pages()}" ) @@ -2890,6 +2903,10 @@ def size_full(self) -> int: # Widened like `size`: a logical token capacity, not a row count. return (self.full_attn_allocator.max_slots - 1) * get_parallel().attn_dcp_size + @property + def draft_virtual_id_space(self) -> int: + return self.size_full + @property def size_mamba(self) -> int: return self.mamba_allocator.max_slots - 1 @@ -3208,6 +3225,7 @@ def __init__( need_sort=need_sort, forward_stream=forward_stream, lazy_compaction=lazy_compaction, + virtual_num_pages=self.full_attn_allocator.num_virtual_ids, ) self._wire_peers() @@ -3390,6 +3408,10 @@ def _ask_float_for_room(self, need_tokens: int) -> None: # (set to the static caps). We do NOT report `max_slots - 1`: under unified # memory pool that ~= full_max + swa_max and would over-promise. + @property + def draft_virtual_id_space(self) -> int: + return self.full_attn_allocator.max_slots - 1 + def debug_print(self) -> str: return ( f"#full-available={self.full_attn_allocator.available_size()}, " diff --git a/python/sglang/srt/model_executor/pool_configurator.py b/python/sglang/srt/model_executor/pool_configurator.py index 38b7fccd42d4..407afea59de6 100644 --- a/python/sglang/srt/model_executor/pool_configurator.py +++ b/python/sglang/srt/model_executor/pool_configurator.py @@ -14,6 +14,7 @@ from __future__ import annotations import logging +from bisect import bisect_right from dataclasses import dataclass from typing import TYPE_CHECKING, Optional @@ -465,6 +466,7 @@ def __init__(self, kvc: KVCacheConfigurator): self._swa_full_tokens_ratio = get_schedule().swa_full_tokens_ratio self._sliding_window_size = kvc.sliding_window_size self._page_size = kvc.page_size + self._enable_unified_memory = get_memory().enable_unified_memory if model_config.attention_arch == AttentionArch.MLA: # MLA pool sizing uses latent dimensions rather than MHA heads. @@ -551,6 +553,9 @@ def __init__(self, kvc: KVCacheConfigurator): self._draft_cell_size = _dflash_draft_cell_size(kvc) + self._recompute_cell_size() + + def _recompute_cell_size(self) -> None: # Bytes per token of max_total_num_tokens. # # Hybrid (full_layers > 0): max_total = full_tokens, so cell_size accounts @@ -579,6 +584,50 @@ def __init__(self, kvc: KVCacheConfigurator): + self._draft_cell_size ) + def _draft_pool_bytes_per_token(self) -> int: + return int( + self._full_per_token * self._draft_full_layers_num + + self._swa_per_token + * (self._draft_swa_layers_num + self._draft_swa_full_layers_num) + + self._draft_cell_size + ) + + def _max_unified_full_tokens( + self, + available_bytes: int, + page_size: int, + fixed_swa_tokens: Optional[int] = None, + ) -> int: + """Find the largest page-aligned full capacity whose allocations fit.""" + draft_bytes_per_token = self._draft_pool_bytes_per_token() + target_full_bytes_per_token = self._full_per_token * self._full_layers_num + target_swa_bytes_per_token = self._swa_per_token * self._swa_layers_num + assert target_full_bytes_per_token > 0 + + def allocation_bytes(full_pages: int) -> int: + full_tokens = full_pages * page_size + swa_tokens = ( + fixed_swa_tokens + if fixed_swa_tokens is not None + else int(full_tokens * self._swa_full_tokens_ratio) + // page_size + * page_size + ) + target_bytes = ( + full_tokens * target_full_bytes_per_token + + swa_tokens * target_swa_bytes_per_token + ) + virtual_span = max(target_bytes // target_full_bytes_per_token - 1, 0) + draft_tokens = ceil_align(virtual_span, page_size) + page_size + return target_bytes + draft_tokens * draft_bytes_per_token + + max_pages = available_bytes // target_full_bytes_per_token // page_size + full_pages = ( + bisect_right(range(max_pages + 1), available_bytes, key=allocation_bytes) + - 1 + ) + return max(full_pages, 0) * page_size + def _solve_pool_sizes( self, max_total_num_tokens: int, page_size: int ) -> MemoryPoolConfig: @@ -630,7 +679,16 @@ def align_page_size(x: int) -> int: def calculate_pool_sizes( self, available_bytes: int, page_size: int ) -> MemoryPoolConfig: - max_total_num_tokens = int(available_bytes // self._cell_size) + if ( + self._enable_unified_memory + and self._full_layers_num > 0 + and self._draft_pool_bytes_per_token() > 0 + ): + max_total_num_tokens = self._max_unified_full_tokens( + available_bytes, page_size + ) + else: + max_total_num_tokens = int(available_bytes // self._cell_size) return self._solve_pool_sizes(max_total_num_tokens, page_size) def calculate_pool_sizes_from_max_tokens( @@ -716,13 +774,19 @@ def calculate_pool_sizes( * self._swa_per_token * (self._swa_layers_num + self._draft_swa_layers_num) ) - full_cell_size = ( - self._full_per_token * (self._full_layers_num + self._draft_full_layers_num) - + self._swa_per_token * self._draft_swa_full_layers_num - ) - full_tokens = ( - int((available_bytes - fixed_swa_bytes) // full_cell_size) // page_size - ) * page_size + if self._enable_unified_memory and self._draft_pool_bytes_per_token() > 0: + full_tokens = self._max_unified_full_tokens( + available_bytes, page_size, fixed_swa_tokens=swa_tokens + ) + else: + full_cell_size = ( + self._full_per_token + * (self._full_layers_num + self._draft_full_layers_num) + + self._swa_per_token * self._draft_swa_full_layers_num + ) + full_tokens = ( + int((available_bytes - fixed_swa_bytes) // full_cell_size) // page_size + ) * page_size if full_tokens <= 0: raise RuntimeError( f"SWA pool cap ({swa_tokens} tokens, " From e887a18fbad2c336f5ff64f7230900c42022358c Mon Sep 17 00:00:00 2001 From: Lianmin Zheng Date: Tue, 1 Sep 2026 02:16:30 -0700 Subject: [PATCH 02/55] fix: use hybrid SWA draft flag for pool setup --- python/sglang/srt/mem_cache/kv_cache_configurator.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/mem_cache/kv_cache_configurator.py b/python/sglang/srt/mem_cache/kv_cache_configurator.py index 7c40094db503..a8d887494135 100644 --- a/python/sglang/srt/mem_cache/kv_cache_configurator.py +++ b/python/sglang/srt/mem_cache/kv_cache_configurator.py @@ -518,7 +518,7 @@ def _init_pools( and self.is_hybrid_swa ): size_overrides["full_max_total_num_tokens"] = draft_virtual_id_space - if not self.is_inkling_mtp_draft or self.draft_swa_full_capacity: + if not self.is_hybrid_swa_mtp_draft or self.draft_swa_full_capacity: size_overrides["swa_max_total_num_tokens"] = ( draft_virtual_id_space ) @@ -1982,7 +1982,7 @@ def _build_token_to_kv_pool_allocator( swa_allocator, UnifiedSWATokenToKVPoolAllocator ) has_draft_swa_layers = ( - not self.is_inkling_mtp_draft or self.draft_swa_full_capacity + not self.is_hybrid_swa_mtp_draft or self.draft_swa_full_capacity ) if self.draft_swa_full_capacity or ( uses_unified_virtual_ids and has_draft_swa_layers From ad6c3b44f69d8dd7cbe0fb27248f0186afb95229 Mon Sep 17 00:00:00 2001 From: yhzhuang Date: Thu, 27 Aug 2026 13:37:21 -0700 Subject: [PATCH 03/55] Use shared byte capacity for unified SWA pools Co-authored-by: Yonghao Zhuang --- python/sglang/srt/disaggregation/decode.py | 202 +++++++++-- python/sglang/srt/managers/schedule_policy.py | 340 +++++++++++++----- python/sglang/srt/managers/scheduler.py | 60 +++- .../scheduler_components/invariant_checker.py | 25 +- .../pool_stats_observer.py | 37 +- python/sglang/srt/mem_cache/allocator/base.py | 10 + python/sglang/srt/mem_cache/common.py | 83 ++++- .../srt/mem_cache/multi_ended_allocator.py | 270 ++++++++++++-- .../srt/mem_cache/unified_memory_pool.py | 26 +- .../sglang/srt/model_executor/model_runner.py | 9 +- .../srt/model_executor/pool_configurator.py | 50 ++- .../test_decode_queue_cleanup.py | 35 +- .../unit/managers/test_prefill_adder.py | 45 +++ .../test_scheduler_init_req_max_new_tokens.py | 1 + .../mem_cache/test_decode_radix_lock_ref.py | 14 +- .../model_executor/test_pool_configurator.py | 20 +- 16 files changed, 969 insertions(+), 258 deletions(-) diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index ed5f836e2272..348dc103df79 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -78,6 +78,7 @@ EvictParams, ) from sglang.srt.mem_cache.common import ( + evict_from_tree_cache, kv_to_page_indices, page_align_floor, release_kv_cache, @@ -90,6 +91,9 @@ KVCache, ReqToTokenPool, ) +from sglang.srt.mem_cache.multi_ended_allocator import ( + UnifiedSWATokenToKVPoolAllocator, +) from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool from sglang.srt.observability.req_time_stats import ( set_schedule_time_batch, @@ -409,6 +413,74 @@ def _uses_swa_tail_prealloc(self) -> bool: and hasattr(self.token_to_kv_pool_allocator, "alloc_extend_swa_tail") ) + def _unified_swa_reservation_fits( + self, + full_tokens: int, + swa_tokens: int, + *, + full_allocatable_tokens: Optional[int] = None, + swa_allocatable_tokens: Optional[int] = None, + empty_pool: bool = False, + ) -> bool: + allocator = self.token_to_kv_pool_allocator + if not ( + isinstance(allocator, UnifiedSWATokenToKVPoolAllocator) + and allocator.supports_asymmetric_reservation + ): + return True + + full_evictable = swa_evictable = 0 + if not empty_pool: + full_evictable = self._radix_full_evictable() + swa_evictable = self.tree_cache.swa_evictable_size() + if full_allocatable_tokens is not None: + full_tokens += ( + allocator.full_available_size() + + full_evictable + - full_allocatable_tokens + ) + if swa_allocatable_tokens is not None: + swa_tokens += ( + allocator.swa_available_size() + + swa_evictable + - swa_allocatable_tokens + ) + + return allocator.can_reserve( + full_tokens, + swa_tokens, + full_evictable_tokens=full_evictable, + swa_evictable_tokens=swa_evictable, + empty_pool=empty_pool, + ) + + def _prealloc_reservation_fits( + self, + full_tokens: int, + swa_tokens: int, + *, + full_allocatable_tokens: int, + swa_allocatable_tokens: Optional[int], + uses_swa_tail_prealloc: bool, + ) -> bool: + if not uses_swa_tail_prealloc: + return full_tokens <= full_allocatable_tokens + assert swa_allocatable_tokens is not None + if isinstance( + self.token_to_kv_pool_allocator, + UnifiedSWATokenToKVPoolAllocator, + ) and self.token_to_kv_pool_allocator.supports_asymmetric_reservation: + return self._unified_swa_reservation_fits( + full_tokens, + swa_tokens, + full_allocatable_tokens=full_allocatable_tokens, + swa_allocatable_tokens=swa_allocatable_tokens, + ) + return ( + full_tokens <= full_allocatable_tokens + and swa_tokens <= swa_allocatable_tokens + ) + def _release_matched_prefix_lock(self, req: Req) -> None: params = DecLockRefParams(swa_uuid_for_lock=req.swa_uuid_for_lock) if req.swa_prefix_lock_released: @@ -418,17 +490,35 @@ def _release_matched_prefix_lock(self, req: Req) -> None: self.tree_cache.dec_lock_ref(req.last_node, params) def _reclaim_swa_tail_capacity( - self, swa_tail_len: int, req_id: str + self, swa_tail_len: int, req_id: str, *, full_len: int = 0 ) -> Optional[str]: - page_size = self.token_to_kv_pool_allocator.page_size + allocator = self.token_to_kv_pool_allocator + page_size = allocator.page_size + if ( + isinstance(allocator, UnifiedSWATokenToKVPoolAllocator) + and allocator.supports_asymmetric_reservation + ): + full_required = ceil_align(full_len, page_size) + swa_required = ceil_align(swa_tail_len, page_size) + evict_from_tree_cache( + self.tree_cache, + full_required, + swa_num_tokens=swa_required, + ) + if allocator.ensure_capacity(full_required, swa_required): + return None + return ( + "Unified FULL/SWA byte reclamation insufficient: " + f"needed=({full_required}, {swa_required}), req={req_id}" + ) + required = ceil_align(swa_tail_len, page_size) - available = self.token_to_kv_pool_allocator.swa_available_size() + available = allocator.swa_available_size() if available < required: self.tree_cache.evict_for_alloc( EvictParams(swa_num_tokens=required - available) ) - available = self.token_to_kv_pool_allocator.swa_available_size() - + available = allocator.swa_available_size() if available < required: return ( f"SWA eviction insufficient: needed={required}, " @@ -759,6 +849,22 @@ def _check_if_req_exceed_kv_capacity(self, req: Req) -> bool: else: capacity = self.max_total_num_tokens input_len = self._rebootstrap_prefill_len(req) + if isinstance( + self.token_to_kv_pool_allocator, UnifiedSWATokenToKVPoolAllocator + ) and self.token_to_kv_pool_allocator.supports_asymmetric_reservation: + full_required, swa_required = self._prealloc_required_tokens(req) + if not self._unified_swa_reservation_fits( + full_required, swa_required, empty_pool=True + ): + message = ( + f"Request {req.rid} exceeds the unified FULL/SWA KV byte " + f"budget: full={full_required}, swa={swa_required}" + ) + logger.error(message) + prepare_abort(req, message, status_code=HTTPStatus.BAD_REQUEST) + self.scheduler.output_streamer.stream_output([req], req.return_logprob) + return True + return False if input_len > capacity: message = f"Request {req.rid} exceeds the maximum number of tokens: {input_len} > {capacity}" logger.error(message) @@ -827,11 +933,28 @@ def resume_retracted_reqs( break full_required, swa_required = self._prealloc_required_tokens(req) - if full_required > full_allocatable_tokens: - break - if uses_swa_tail_prealloc and swa_required > swa_allocatable_tokens: + if not self._prealloc_reservation_fits( + full_required, + swa_required, + full_allocatable_tokens=full_allocatable_tokens, + swa_allocatable_tokens=( + swa_allocatable_tokens if uses_swa_tail_prealloc else None + ), + uses_swa_tail_prealloc=uses_swa_tail_prealloc, + ): break + if uses_swa_tail_prealloc and isinstance( + self.token_to_kv_pool_allocator, + UnifiedSWATokenToKVPoolAllocator, + ) and self.token_to_kv_pool_allocator.supports_asymmetric_reservation: + full_len, swa_len = self._prealloc_kv_lens(req) + if ( + self._reclaim_swa_tail_capacity(swa_len, req.rid, full_len=full_len) + is not None + ): + break + resumed_reqs.append(req) indices_to_remove.add(i) req.is_retracted = False @@ -1270,27 +1393,18 @@ def pop_preallocated( required_alloc_tokens + self.num_reserved_decode_tokens ) - if ( - max( - required_tokens_for_request, - origin_input_len - - prefix_len - + min( - decode_req.req.sampling_params.max_new_tokens, - CLIP_MAX_NEW_TOKEN, - ) - - retractable_tokens, + full_required_for_admission = max( + required_tokens_for_request, + origin_input_len + - prefix_len + + min( + decode_req.req.sampling_params.max_new_tokens, + CLIP_MAX_NEW_TOKEN, ) - > full_allocatable_tokens - ): - if prefix_match is not None and prefix_match.l1_prefix_len > 0: - self._release_matched_prefix_lock(decode_req.req) - break - if required_tokens_for_request > full_allocatable_tokens: - if prefix_match is not None and prefix_match.l1_prefix_len > 0: - self._release_matched_prefix_lock(decode_req.req) - break - + - retractable_tokens, + ) + swa_required_for_admission = 0 + swa_len = 0 if uses_swa_tail_prealloc: _, swa_required = self._prealloc_required_tokens(decode_req.req) _, swa_len = self._prealloc_kv_lens(decode_req.req) @@ -1298,19 +1412,29 @@ def pop_preallocated( decode_req.req.sampling_params.max_new_tokens, CLIP_MAX_NEW_TOKEN, ) - if ( - max( - swa_required, - swa_len + max_new_tokens - retractable_swa_tokens, - ) - > swa_allocatable_tokens - ): - if prefix_match is not None and prefix_match.l1_prefix_len > 0: - self._release_matched_prefix_lock(decode_req.req) - break + swa_required_for_admission = max( + swa_required, + swa_len + max_new_tokens - retractable_swa_tokens, + ) + if not self._prealloc_reservation_fits( + full_required_for_admission, + swa_required_for_admission, + full_allocatable_tokens=full_allocatable_tokens, + swa_allocatable_tokens=( + swa_allocatable_tokens if uses_swa_tail_prealloc else None + ), + uses_swa_tail_prealloc=uses_swa_tail_prealloc, + ): + if prefix_match is not None and prefix_match.l1_prefix_len > 0: + self._release_matched_prefix_lock(decode_req.req) + break + + if uses_swa_tail_prealloc: reclaim_error = self._reclaim_swa_tail_capacity( - swa_len, decode_req.req.rid + swa_len, + decode_req.req.rid, + full_len=required_alloc_tokens, ) if reclaim_error is not None: if prefix_match is not None and prefix_match.l1_prefix_len > 0: diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index ae9fa0dc9551..44490f9300dd 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -63,6 +63,7 @@ from sglang.srt.mem_cache.multi_ended_allocator import ( UnifiedMambaSWATokenToKVPoolAllocator, UnifiedMambaTokenToKVPoolAllocator, + UnifiedSWATokenToKVPoolAllocator, ) from sglang.srt.mem_cache.radix_cache import RadixCache, RadixKey, TreeNode @@ -77,6 +78,35 @@ os.environ.get("SGLANG_CLIP_MAX_NEW_TOKENS_ESTIMATION", "4096") ) + +def estimate_swa_kv_tokens( + extend_input_len: int, + max_new_tokens: int, + *, + sliding_window_size: Optional[int], + page_size: int, + allocation_limit: Optional[int] = None, + host_hit_length: int = 0, +) -> int: + """Peak SWA reservation for one prefill/decode request.""" + if sliding_window_size is None or sliding_window_size <= 0: + reserved = extend_input_len + max_new_tokens + page_size + else: + allocated = ( + extend_input_len + if allocation_limit is None + else min(extend_input_len, allocation_limit) + ) + reserved = ( + max(allocated - sliding_window_size, 0) + + min(extend_input_len + max_new_tokens, sliding_window_size) + + page_size + ) + if host_hit_length > 0: + reserved += -(-host_hit_length // page_size) * page_size + return reserved + + # Threshold for in-batch prefix cache. # If a request has a matched prefix length (against existing cache) less than this value, # the scheduler runs the in-batch prefix caching check for this request. @@ -545,8 +575,13 @@ def __init__( self.token_to_kv_pool_allocator, PureSWATokenToKVPoolAllocator ) self.is_hybrid_ssm_cache = self.tree_cache.supports_mamba() + self.is_unified_swa = isinstance( + self.token_to_kv_pool_allocator, UnifiedSWATokenToKVPoolAllocator + ) and self.token_to_kv_pool_allocator.supports_asymmetric_reservation - self.rem_swa_token_offset = 0 + self.rem_swa_token_offset = ( + num_mixed_decode_tokens if self.is_unified_swa else 0 + ) # A new state slot eats shared-gap bytes that `rem_total_tokens` counts # as free, so reserve per slot or admission over-commits. Gate on the @@ -710,33 +745,28 @@ def _swa_budget_for_req( where alloc = min(extend, rem_chunk); the min() cap keeps the two terms from double-counting extend, so budget <= extend + max_new_tokens + page. """ - if self.rem_chunk_tokens is not None: - alloc = min(extend_input_len, self.rem_chunk_tokens) - else: - alloc = extend_input_len - window = self.tree_cache.sliding_window_size - return max(alloc - window, 0) + self._swa_reserved_tokens( - extend_input_len, max_new_tokens, swa_host_hit_length + allocated = ( + extend_input_len + if self.rem_chunk_tokens is None + else min(extend_input_len, self.rem_chunk_tokens) + ) + return max(allocated - self.tree_cache.sliding_window_size, 0) + ( + self._swa_reserved_tokens( + extend_input_len, max_new_tokens, swa_host_hit_length + ) ) def _swa_reserved_tokens( self, extend_input_len: int, max_new_tokens: int, swa_host_hit_length: int = 0 ) -> int: - """SWA slots a request adds to its own sliding window + page slack + the - load-back charge. Shared floor of _swa_budget_for_req and _swa_chunk_cap. - - The headroom is min(extend + decode, window), not a constant window: a - request contributes only extend + decode fresh tokens to its window and - a cached SWA prefix funds the rest. Charging a full window double-counted - a short cached-prefix resume and livelocked admission at a ~2-window - pool; keeping extend in the min() holds the reservation >= the prefill - allocation so an admitted request cannot OOM.""" - window = self.tree_cache.sliding_window_size - headroom = min(extend_input_len + max_new_tokens, window) - reserved = headroom + self.page_size - if swa_host_hit_length > 0: - reserved += self.ceil_paged_tokens(swa_host_hit_length) - return reserved + return estimate_swa_kv_tokens( + extend_input_len, + max_new_tokens, + sliding_window_size=self.tree_cache.sliding_window_size, + page_size=self.page_size, + allocation_limit=0, + host_hit_length=swa_host_hit_length, + ) def _swa_new_tokens(self, req: Req) -> int: """Tokens a request may still decode, for SWA headroom sizing. Mirrors @@ -748,10 +778,45 @@ def _swa_new_tokens(self, req: Req) -> int: CLIP_MAX_NEW_TOKENS, ) - def _swa_chunk_cap(self, max_new_tokens: int, swa_host_hit_length: int = 0) -> int: - """Largest page-aligned extend chunk the SWA pool can admit right now, - keeping a sliding window of headroom below rem_swa_tokens; 0 if not - even one page fits. Only valid when is_hybrid_swa is True. + def _unified_swa_reservation_fits( + self, + full_tokens: int, + swa_tokens: int, + *, + empty_pool: bool = False, + ) -> bool: + if not self.is_unified_swa: + return True + allocator = self.token_to_kv_pool_allocator + if empty_pool: + full_evictable_tokens = 0 + swa_evictable_tokens = 0 + else: + full_tokens += self.rem_total_token_offset + swa_tokens += self.rem_swa_token_offset + full_evictable_tokens = self.tree_cache.full_evictable_size() + swa_evictable_tokens = self.tree_cache.swa_evictable_size() + return allocator.can_reserve( + full_tokens, + swa_tokens, + full_evictable_tokens=full_evictable_tokens, + swa_evictable_tokens=swa_evictable_tokens, + empty_pool=empty_pool, + require_token_slack=empty_pool, + ) + + def _swa_chunk_cap( + self, + max_new_tokens: int, + swa_host_hit_length: int = 0, + *, + max_chunk_tokens: Optional[int] = None, + ) -> int: + """Largest page-aligned extend chunk the SWA pool can admit right now. + + Split pools reserve headroom below rem_swa_tokens; unified pools use + the joint byte budget. Returns 0 if not even one page fits. Only valid + when is_hybrid_swa is True. Escape hatch for a request whose budget can never pass the _swa_budget_for_req gate (extend near/above the pool size, or a large @@ -759,17 +824,41 @@ def _swa_chunk_cap(self, max_new_tokens: int, swa_host_hit_length: int = 0) -> i forever (head-of-line livelock). Shrinking is sound because past a chunk boundary only the sliding window stays locked — the rest turns evictable — so each pass's transient footprint fits the pool.""" - # extend_input_len=0: this solves for the extend chunk itself, so the - # reserved headroom is the post-chunk decode window only. - cap = int(self.rem_swa_tokens) - self._swa_reserved_tokens( - 0, max_new_tokens, swa_host_hit_length - ) - if cap <= 0: - return 0 - return cap // self.page_size * self.page_size + if not self.is_unified_swa: + # extend_input_len=0: this solves for the extend chunk itself, so + # the reserved headroom is the post-chunk decode window only. + cap = int(self.rem_swa_tokens) - self._swa_reserved_tokens( + 0, max_new_tokens, swa_host_hit_length + ) + if cap <= 0: + return 0 + cap = cap // self.page_size * self.page_size + return cap + + assert max_chunk_tokens is not None + cap = max_chunk_tokens // self.page_size * self.page_size + lo, hi = 0, cap // self.page_size + while lo < hi: + mid = (lo + hi + 1) // 2 + chunk_tokens = mid * self.page_size + full_tokens = self.ceil_paged_tokens(chunk_tokens) + self.page_size + swa_tokens = self._swa_budget_for_req( + chunk_tokens, max_new_tokens, swa_host_hit_length + ) + if self._unified_swa_reservation_fits(full_tokens, swa_tokens): + lo = mid + else: + hi = mid - 1 + return lo * self.page_size def _swa_req_never_fits( - self, extend_input_len: int, max_new_tokens: int, swa_host_hit_length: int = 0 + self, + extend_input_len: int, + max_new_tokens: int, + swa_host_hit_length: int = 0, + *, + full_tokens: Optional[int] = None, + unified_swa_tokens: Optional[int] = None, ) -> bool: """True when a request's SWA budget exceeds the *entire* SWA pool, so it can never be admitted whole no matter how far the pool drains. @@ -781,13 +870,72 @@ def _swa_req_never_fits( decode headroom collapses the SWA evictable cushion and forces running requests to retract (observed as a severe retraction/re-prefill storm on hybrid-SWA models at high concurrency).""" - capacity = self.token_to_kv_pool_allocator.size_swa - return ( - self._swa_budget_for_req( - extend_input_len, max_new_tokens, swa_host_hit_length + swa_tokens = self._swa_budget_for_req( + extend_input_len, max_new_tokens, swa_host_hit_length + ) + if self.is_unified_swa: + assert full_tokens is not None + assert unified_swa_tokens is not None + return not self._unified_swa_reservation_fits( + full_tokens, unified_swa_tokens, empty_pool=True + ) + return swa_tokens >= self.token_to_kv_pool_allocator.size_swa + + def _check_swa_admission( + self, + req: Req, + *, + real_input_tokens: int, + total_tokens: int, + ) -> tuple[bool, Optional[int]]: + """Check SWA pressure and return the usable chunk limit.""" + max_new_tokens = self._swa_new_tokens(req) + swa_needed = self._swa_budget_for_req( + real_input_tokens, + max_new_tokens, + swa_host_hit_length=req.swa_host_hit_length, + ) + reservation_fits = ( + self._unified_swa_reservation_fits(total_tokens, swa_needed) + if self.is_unified_swa + else swa_needed < self.rem_swa_tokens + ) + if reservation_fits: + return True, self.rem_chunk_tokens + + full_ever_tokens = unified_swa_tokens = None + if self.is_unified_swa: + full_input_tokens = len(req.full_untruncated_fill_ids) + full_ever_tokens = ( + self.ceil_paged_tokens(full_input_tokens) + + max_new_tokens + + self.page_size + ) + unified_swa_tokens = self._swa_budget_for_req( + full_input_tokens, max_new_tokens ) - >= capacity + if not self._swa_req_never_fits( + real_input_tokens, + max_new_tokens, + req.swa_host_hit_length, + full_tokens=full_ever_tokens, + unified_swa_tokens=unified_swa_tokens, + ): + return False, None + + max_chunk_tokens = self.rem_chunk_tokens + chunk_max_new = max_new_tokens + if self.is_unified_swa: + max_chunk_tokens = min(max_chunk_tokens or 0, max(0, real_input_tokens - 1)) + chunk_max_new = 0 + swa_cap = self._swa_chunk_cap( + chunk_max_new, + req.swa_host_hit_length, + max_chunk_tokens=max_chunk_tokens, ) + if self.rem_chunk_tokens is None or swa_cap <= 0: + return False, None + return True, min(self.rem_chunk_tokens, swa_cap) def _mamba_gap_budget_for_req(self, req: Req) -> int: """Shared-gap reservation (full-token-equivalents) for a request's new @@ -808,9 +956,12 @@ def ceil_paged_tokens(self, tokens: int) -> int: return -(-tokens // self.page_size) * self.page_size def budget_state(self): - no_token = self.rem_total_tokens <= 0 or self.cur_rem_tokens <= 0 - if not no_token and self.is_hybrid_swa: - no_token = self.rem_swa_tokens <= 0 + if self.is_unified_swa: + no_token = not self._unified_swa_reservation_fits(0, 0) + else: + no_token = self.rem_total_tokens <= 0 or self.cur_rem_tokens <= 0 + if not no_token and self.is_hybrid_swa: + no_token = self.rem_swa_tokens <= 0 # Gate new mamba slots separately: rem_total_tokens' full_evictable can't # cover a mamba slot, which needs mamba-recoverable bytes (see __init__). if not no_token and self.rem_mamba_slots is not None: @@ -974,8 +1125,11 @@ def add_chunked_req(self, req: Req): if self.dllm_config is not None: _rem_tokens = self._get_dllm_remain_tokens() else: - _rem_tokens = min(self.rem_chunk_tokens, int(self.rem_total_tokens)) - if self.is_hybrid_swa: + if self.is_unified_swa: + _rem_tokens = self.rem_chunk_tokens + else: + _rem_tokens = min(self.rem_chunk_tokens, int(self.rem_total_tokens)) + if self.is_hybrid_swa and not self.is_unified_swa: # alloc_extend needs extend_num_tokens + page_size per request, # so reserve one page here to avoid OOM _rem_tokens = min( @@ -1002,6 +1156,23 @@ def add_chunked_req(self, req: Req): cand_extend_input_len = len(req.full_untruncated_fill_ids) - len( req.prefix_indices ) + if self.is_unified_swa: + candidate_tokens = min(cand_extend_input_len, _rem_tokens) + candidate_finishes = candidate_tokens >= cand_extend_input_len + candidate_max_new = self._swa_new_tokens(req) if candidate_finishes else 0 + if not self._unified_swa_reservation_fits( + self.ceil_paged_tokens(candidate_tokens) + + candidate_max_new + + self.page_size, + self._swa_budget_for_req(candidate_tokens, candidate_max_new), + ): + max_partial = candidate_tokens + if candidate_finishes: + max_partial = max(0, candidate_tokens - 1) + joint_cap = self._swa_chunk_cap(0, max_chunk_tokens=max_partial) + if joint_cap <= 0: + return req + _rem_tokens = min(_rem_tokens, joint_cap) truncated = cand_extend_input_len > _rem_tokens new_len = min(cand_extend_input_len, _rem_tokens) req.set_extend_range(len(req.prefix_indices), len(req.prefix_indices) + new_len) @@ -1046,16 +1217,22 @@ def add_one_req_ignore_eos(self, req: Req): # Shared Mamba pool: fold the new mamba state's shared-gap cost into the # budget gate so admission can't over-commit (0 for baseline / non-Mamba). paged_input += self._mamba_gap_budget_for_req(req) - if paged_input > min(self.cur_rem_tokens, self.rem_total_tokens): - return AddReqResult.NO_TOKEN if self.is_hybrid_swa: - if ( - self._swa_budget_for_req( - cand_extend_input_len, self._swa_new_tokens(req) - ) - > self.rem_swa_tokens + max_new_tokens = self._swa_new_tokens(req) + swa_needed = self._swa_budget_for_req(cand_extend_input_len, max_new_tokens) + if self.is_unified_swa: + if not self._unified_swa_reservation_fits( + paged_input + max_new_tokens + self.page_size, + swa_needed, + ): + return AddReqResult.NO_TOKEN + elif ( + paged_input > min(self.cur_rem_tokens, self.rem_total_tokens) + or swa_needed > self.rem_swa_tokens ): return AddReqResult.NO_TOKEN + elif paged_input > min(self.cur_rem_tokens, self.rem_total_tokens): + return AddReqResult.NO_TOKEN def add_req_state(r, insert_sort=False): new_token_ratio = ( @@ -1203,38 +1380,23 @@ def add_one_req( # Shared Mamba pool: fold the new mamba state's shared-gap cost into # `total_tokens` so both `rem_total_tokens` gates reflect the joint budget. total_tokens += self._mamba_gap_budget_for_req(req) - # adjusting the input_tokens based on host_hit_length and page_size real_input_tokens = cand_extend_input_len - req.host_hit_length real_input_tokens = self.ceil_paged_tokens(real_input_tokens) prefix_len = len(req.prefix_indices) - if total_tokens >= self.rem_total_tokens: + if not self.is_unified_swa and total_tokens >= self.rem_total_tokens: return AddReqResult.NO_TOKEN chunk_tokens_limit = self.rem_chunk_tokens if self.is_hybrid_swa: - # host-hit prefix is loaded back, not re-prefilled, so the SWA peak is - # driven only by the freshly-prefilled tail (the loaded window is - # charged separately via swa_host_hit_length). - swa_needed = self._swa_budget_for_req( - real_input_tokens, - self._swa_new_tokens(req), - swa_host_hit_length=req.swa_host_hit_length, + can_admit, chunk_tokens_limit = self._check_swa_admission( + req, + real_input_tokens=real_input_tokens, + total_tokens=total_tokens, ) - if swa_needed >= self.rem_swa_tokens: - if not self._swa_req_never_fits( - real_input_tokens, - self._swa_new_tokens(req), - req.swa_host_hit_length, - ): - return AddReqResult.NO_TOKEN - swa_cap = self._swa_chunk_cap( - self._swa_new_tokens(req), req.swa_host_hit_length - ) - if self.rem_chunk_tokens is None or swa_cap <= 0: - return AddReqResult.NO_TOKEN - chunk_tokens_limit = min(self.rem_chunk_tokens, swa_cap) + if not can_admit: + return AddReqResult.NO_TOKEN if ( self.rem_chunk_tokens is None @@ -1247,30 +1409,18 @@ def add_one_req( return AddReqResult.OTHER with self._lock_node(req.last_node): - # self.rem_total_tokens may decrease after the lock acquisition - if total_tokens >= self.rem_total_tokens: + # Capacity may decrease after the lock acquisition. + if not self.is_unified_swa and total_tokens >= self.rem_total_tokens: return AddReqResult.NO_TOKEN if self.is_hybrid_swa: - # self.rem_swa_tokens may decrease after the lock acquisition - swa_needed = self._swa_budget_for_req( - real_input_tokens, - self._swa_new_tokens(req), - swa_host_hit_length=req.swa_host_hit_length, + can_admit, chunk_tokens_limit = self._check_swa_admission( + req, + real_input_tokens=real_input_tokens, + total_tokens=total_tokens, ) - if swa_needed >= self.rem_swa_tokens: - if not self._swa_req_never_fits( - real_input_tokens, - self._swa_new_tokens(req), - req.swa_host_hit_length, - ): - return AddReqResult.NO_TOKEN - swa_cap = self._swa_chunk_cap( - self._swa_new_tokens(req), req.swa_host_hit_length - ) - if self.rem_chunk_tokens is None or swa_cap <= 0: - return AddReqResult.NO_TOKEN - chunk_tokens_limit = min(self.rem_chunk_tokens, swa_cap) + if not can_admit: + return AddReqResult.NO_TOKEN # Negotiate only after every KV-budget gate (a NO_TOKEN rank must # report not-prefillable via finalize()) and before init_load_back diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 454f75454452..91864864525c 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -211,6 +211,7 @@ AddReqResult, PrefillAdder, SchedulePolicy, + estimate_swa_kv_tokens, ) from sglang.srt.managers.scheduler_components.batch_result_processor import ( SchedulerBatchResultProcessor, @@ -280,6 +281,9 @@ release_kv_cache, retraction_discard, ) +from sglang.srt.mem_cache.multi_ended_allocator import ( + UnifiedSWATokenToKVPoolAllocator, +) from sglang.srt.model_executor.forward_batch_info import PPProxyTensors from sglang.srt.model_loader.utils import get_resolved_model_impl from sglang.srt.multiplex.multiplexing_mixin import SchedulerMultiplexMixin @@ -2392,23 +2396,63 @@ def init_req_max_new_tokens(self, req): ) max_new_tokens = min(max_new_tokens, self.max_new_tokens_limit) - # Keep this bound consistent with PrefillAdder's admission budget: - # ceil_page(input_len) + max_new_tokens + page_size must be strictly - # smaller than max_total_num_tokens. Otherwise a request can be accepted - # into the waiting queue but can never be scheduled, blocking the queue - # and eventually making health checks fail. + # Keep this bound consistent with PrefillAdder's admission budget. paged_input_len = -(-input_len // self.page_size) * self.page_size - req.sampling_params.max_new_tokens = max( + max_new_tokens = max( 0, min( max_new_tokens, self.max_req_len - input_len - 1, + ), + ) + allocator = self.token_to_kv_pool_allocator + if ( + isinstance(allocator, UnifiedSWATokenToKVPoolAllocator) + and allocator.supports_asymmetric_reservation + ): + + def fits_shared_pool(candidate: int) -> bool: + full_tokens = paged_input_len + candidate + self.page_size + swa_tokens = estimate_swa_kv_tokens( + input_len, + candidate, + sliding_window_size=self.sliding_window_size, + page_size=self.page_size, + allocation_limit=self.chunked_prefill_size, + ) + return allocator.can_reserve( + full_tokens, + swa_tokens, + empty_pool=True, + require_token_slack=True, + ) + + if not fits_shared_pool(0): + req.set_finish_with_abort( + "Request prompt exceeds the unified FULL/SWA KV byte budget: " + f"input_len={input_len}." + ) + max_new_tokens = 0 + elif not fits_shared_pool(max_new_tokens): + lo, hi = 0, max_new_tokens + while lo < hi: + mid = (lo + hi + 1) // 2 + if fits_shared_pool(mid): + lo = mid + else: + hi = mid - 1 + max_new_tokens = lo + else: + # Non-unified pools keep their fixed token partition. + max_new_tokens = min( + max_new_tokens, self.max_total_num_tokens * get_parallel().attn_dcp_size - paged_input_len - self.page_size - 1, - ), - ) + ) + + req.sampling_params.max_new_tokens = max(0, max_new_tokens) # Clipping above can push max_new_tokens below min_new_tokens, which # would suppress EOS for the whole generation. Restore the invariant. if req.sampling_params.min_new_tokens > req.sampling_params.max_new_tokens: diff --git a/python/sglang/srt/managers/scheduler_components/invariant_checker.py b/python/sglang/srt/managers/scheduler_components/invariant_checker.py index e2af23cfb938..c619f8aa77b1 100644 --- a/python/sglang/srt/managers/scheduler_components/invariant_checker.py +++ b/python/sglang/srt/managers/scheduler_components/invariant_checker.py @@ -23,9 +23,6 @@ from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache from sglang.srt.mem_cache.memory_pool import ReqToTokenPool -from sglang.srt.mem_cache.multi_ended_allocator import ( - UnifiedMambaSWATokenToKVPoolAllocator, -) from sglang.srt.runtime_context import get_parallel from sglang.srt.utils.common import ( ceil_align, @@ -84,12 +81,13 @@ def _check_pool_invariant( return leak, msg def _check_full_pool(self, ps: PoolStats, uncached: int = 0) -> Tuple[bool, str]: - if self.is_hybrid_swa and not self.full_tokens_per_layer: + full_capacity = self.token_to_kv_pool_allocator.current_full_capacity + if self.is_hybrid_swa and not full_capacity: return False, "" if self.is_hybrid_swa: protected = self.tree_cache.full_protected_size() session_held = self.pool_stats_observer.session_held_full_tokens() - total = self.full_tokens_per_layer + total = full_capacity elif self.is_hybrid_ssm: # Branch on cache type for the protected accessor (MambaRadixCache # splits full/mamba; ChunkCache only has the single protected_size). @@ -119,14 +117,9 @@ def _check_full_pool(self, ps: PoolStats, uncached: int = 0) -> Tuple[bool, str] // allocator.page_size * allocator.page_size ) - full_available = ps.full_available_size - if isinstance(allocator, UnifiedMambaSWATokenToKVPoolAllocator): - # Pair the static per-layer total with the conserve view, never the - # byte-coordinated one -- see `conserve_full_available_size`. - full_available = allocator.conserve_full_available_size() leak, msg = self._check_pool_invariant( "full", - full_available, + ps.full_available_size, full_evictable_size, protected, session_held, @@ -142,19 +135,13 @@ def _check_full_pool(self, ps: PoolStats, uncached: int = 0) -> Tuple[bool, str] return leak, msg def _check_swa_pool(self, ps: PoolStats, uncached: int = 0) -> Tuple[bool, str]: - allocator = self.token_to_kv_pool_allocator - swa_available = ps.swa_available_size - if isinstance(allocator, UnifiedMambaSWATokenToKVPoolAllocator): - # Tri-pool: same floating-boundary phantom as the full pool -- use the - # slot-conservation view, not the byte-coordinated min (see _check_full_pool). - swa_available = allocator.conserve_swa_available_size() return self._check_pool_invariant( "swa", - swa_available, + ps.swa_available_size, ps.swa_evictable_size, self.tree_cache.swa_protected_size(), self.pool_stats_observer.session_held_swa_tokens(), - self.swa_tokens_per_layer, + self.token_to_kv_pool_allocator.current_swa_capacity, uncached, ) diff --git a/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py b/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py index 7dc5f63b8093..1d6fc836bbee 100644 --- a/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py +++ b/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py @@ -11,10 +11,6 @@ Tuple, ) -from sglang.srt.mem_cache.multi_ended_allocator import ( - UnifiedMambaSWATokenToKVPoolAllocator, -) - if TYPE_CHECKING: from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache @@ -288,25 +284,15 @@ def _get_mamba_token_info(self): ) def _get_swa_token_info(self) -> PoolStats: - # `*_num_used` is `static_cap - (available + evictable)`, so the - # available term must match the static cap's denomination: the conserve - # view, never the byte-coordinated one (see - # `conserve_full_available_size`). Measured ~25-90x inflated otherwise. allocator = self.token_to_kv_pool_allocator - if isinstance(allocator, UnifiedMambaSWATokenToKVPoolAllocator): - full_available_size = allocator.conserve_full_available_size() - swa_available_size = allocator.conserve_swa_available_size() - else: - full_available_size = allocator.full_available_size() - swa_available_size = allocator.swa_available_size() + full_capacity = allocator.current_full_capacity + swa_capacity = allocator.current_swa_capacity + full_available_size = allocator.full_available_size() + swa_available_size = allocator.swa_available_size() full_evictable_size = self.tree_cache.full_evictable_size() swa_evictable_size = self.tree_cache.swa_evictable_size() - full_num_used = self.full_tokens_per_layer - ( - full_available_size + full_evictable_size - ) - swa_num_used = self.swa_tokens_per_layer - ( - swa_available_size + swa_evictable_size - ) + full_num_used = full_capacity - (full_available_size + full_evictable_size) + swa_num_used = swa_capacity - (swa_available_size + swa_evictable_size) # FIXME(hisparse): host-backup transiently over-releases the device pool # counter, producing negative full_num_used / swa_num_used. We clamp to 0 # to keep token_usage / leak checks sane, but the underlying accounting @@ -314,13 +300,18 @@ def _get_swa_token_info(self) -> PoolStats: if self.enable_hisparse: full_num_used = max(0, full_num_used) swa_num_used = max(0, swa_num_used) - if not self.full_tokens_per_layer: + if not full_capacity: full_num_used = 0 full_available_size = 0 full_token_usage = 0.0 else: - full_token_usage = full_num_used / self.full_tokens_per_layer - swa_token_usage = swa_num_used / self.swa_tokens_per_layer + full_token_usage = full_num_used / full_capacity + if not swa_capacity: + swa_num_used = 0 + swa_available_size = 0 + swa_token_usage = 0.0 + else: + swa_token_usage = swa_num_used / swa_capacity return PoolStats( is_hybrid_swa=True, diff --git a/python/sglang/srt/mem_cache/allocator/base.py b/python/sglang/srt/mem_cache/allocator/base.py index c9fb67c42f74..79c4df5e9934 100644 --- a/python/sglang/srt/mem_cache/allocator/base.py +++ b/python/sglang/srt/mem_cache/allocator/base.py @@ -84,6 +84,16 @@ def verify_byte_accounting(self) -> list: """ return [] + @property + def current_full_capacity(self) -> int: + """Capacity represented by full_available_size plus live allocations.""" + return self.size_full + + @property + def current_swa_capacity(self) -> int: + """SWA counterpart of current_full_capacity for hybrid allocators.""" + return self.size_swa + def debug_print(self) -> str: return "" diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index b1497c6f5c93..ac07bdeeb282 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/python/sglang/srt/mem_cache/common.py @@ -165,7 +165,11 @@ def maybe_cache_unfinished_req(req: Req, tree_cache: BasePrefixCache, **kwargs): tree_cache.cache_unfinished_req(req, **kwargs) -def evict_from_tree_cache(tree_cache: BasePrefixCache | None, num_tokens: int): +def evict_from_tree_cache( + tree_cache: BasePrefixCache | None, + num_tokens: int, + swa_num_tokens: Optional[int] = None, +): if tree_cache is None: return @@ -174,14 +178,85 @@ def evict_from_tree_cache(tree_cache: BasePrefixCache | None, num_tokens: int): allocator = tree_cache.token_to_kv_pool_allocator - if isinstance(allocator, SWATokenToKVPoolAllocator): + from sglang.srt.mem_cache.multi_ended_allocator import ( + UnifiedSWATokenToKVPoolAllocator, + ) + + if ( + isinstance(allocator, UnifiedSWATokenToKVPoolAllocator) + and allocator.supports_asymmetric_reservation + ): + required_swa = num_tokens if swa_num_tokens is None else swa_num_tokens + if allocator.can_reserve(num_tokens, required_swa): + return + + page_size = allocator.page_size + full_evictable = tree_cache.full_evictable_size() // page_size * page_size + swa_evictable = tree_cache.swa_evictable_size() // page_size * page_size + if not allocator.can_reserve( + num_tokens, + required_swa, + full_evictable_tokens=full_evictable, + swa_evictable_tokens=swa_evictable, + ): + return + + def minimum_reclaim(max_tokens: int, fits) -> int: + lo, hi = 0, max_tokens // page_size + while lo < hi: + mid = (lo + hi) // 2 + if fits(mid * page_size): + hi = mid + else: + lo = mid + 1 + return lo * page_size + + if allocator.can_reserve( + num_tokens, + required_swa, + swa_evictable_tokens=swa_evictable, + ): + full_reclaim = 0 + swa_reclaim = minimum_reclaim( + swa_evictable, + lambda value: allocator.can_reserve( + num_tokens, required_swa, swa_evictable_tokens=value + ), + ) + else: + swa_reclaim = swa_evictable + full_reclaim = minimum_reclaim( + full_evictable, + lambda value: allocator.can_reserve( + num_tokens, + required_swa, + full_evictable_tokens=value, + swa_evictable_tokens=swa_reclaim, + ), + ) + evicted = tree_cache.evict( + EvictParams( + num_tokens=full_reclaim, + swa_num_tokens=swa_reclaim, + ) + ) + if not allocator.can_reserve(num_tokens, required_swa): + logger.warning( + "Unified SWA eviction did not satisfy the planned reservation: " + f"requested=({num_tokens}, {required_swa}), " + f"planned=({full_reclaim}, {swa_reclaim}), " + f"evicted=({evicted.num_tokens_evicted}, " + f"{evicted.swa_num_tokens_evicted})" + ) + elif isinstance(allocator, SWATokenToKVPoolAllocator): # Hybrid allocator + required_swa = num_tokens if swa_num_tokens is None else swa_num_tokens full_available_size = allocator.full_available_size() swa_available_size = allocator.swa_available_size() - if full_available_size < num_tokens or swa_available_size < num_tokens: + if full_available_size < num_tokens or swa_available_size < required_swa: full_num_tokens = max(0, num_tokens - full_available_size) - swa_num_tokens = max(0, num_tokens - swa_available_size) + swa_num_tokens = max(0, required_swa - swa_available_size) tree_cache.evict_for_alloc( EvictParams(num_tokens=full_num_tokens, swa_num_tokens=swa_num_tokens) ) diff --git a/python/sglang/srt/mem_cache/multi_ended_allocator.py b/python/sglang/srt/mem_cache/multi_ended_allocator.py index 606a3db89e28..d4b40a126ead 100644 --- a/python/sglang/srt/mem_cache/multi_ended_allocator.py +++ b/python/sglang/srt/mem_cache/multi_ended_allocator.py @@ -24,6 +24,7 @@ import inspect import logging +import math import os from typing import ( Callable, @@ -3152,14 +3153,16 @@ class UnifiedSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator): replaces). Capacity views: - - `available_size()`: joint byte-budget, the only safe `alloc(N)` pre-check - (N slots cost N*(entry_full + entry_swa) shared-gap bytes). - - `_conserve_*`: slot-conservation, for the LEAK invariant only. + - `available_size()`: joint byte-budget for symmetric FULL/SWA demand. + - `can_reserve()`: read-only admission check. + - `ensure_capacity()`: allocation-time check, with compaction on shortfall. - `schedulable_*`: byte-coordinated, realizable-with-compaction. - `full_available_size()` / `swa_available_size()`: per-side scheduler view - = min(conserve, schedulable). + bounded by physical/index space and the shared byte gap. """ + supports_asymmetric_reservation = True + # Parent's `size` property has no setter but base init does `self.size = size`; # override with a no-op setter. Reading returns `min(_size_full, _size_swa)`. @property @@ -3176,26 +3179,36 @@ def __init__( unified_buffer: UnifiedKVPool, kvcache, # UnifiedSWAKVPool device: str, - full_max_total_num_tokens: int, - swa_max_total_num_tokens: int, + full_max_total_num_tokens: Optional[int] = None, + swa_max_total_num_tokens: Optional[int] = None, page_size: int = 1, need_sort: bool = False, forward_stream: Optional[torch.cuda.Stream] = None, lazy_compaction: bool = False, ): - # Set _size_full / _size_swa BEFORE base init (read during it). STATIC - # partition caps — the slot-conservation value the leak invariant expects. - self._size_full = full_max_total_num_tokens - self._size_swa = swa_max_total_num_tokens - self._full_max_total_num_tokens = full_max_total_num_tokens - self._swa_max_total_num_tokens = swa_max_total_num_tokens + if (full_max_total_num_tokens is None) != (swa_max_total_num_tokens is None): + raise ValueError( + "full_max_total_num_tokens and swa_max_total_num_tokens must " + "either both be set or both be omitted" + ) + legacy_capacities = full_max_total_num_tokens is not None + self._size_full = ( + int(full_max_total_num_tokens) + if legacy_capacities + else unified_buffer.max_slots("full") - 1 + ) + self._size_swa = ( + int(swa_max_total_num_tokens) + if legacy_capacities + else unified_buffer.max_slots("swa") - 1 + ) self.page_size = page_size # Skip SWATokenToKVPoolAllocator.__init__; call grand-parent base init # directly (its `self.size = size` is absorbed by our no-op setter). BaseTokenToKVPoolAllocator.__init__( self, - size=full_max_total_num_tokens, + size=self._size_full, page_size=page_size, dtype=unified_buffer.mha_spec("full").store_dtype, device=device, @@ -3229,6 +3242,18 @@ def __init__( ) self._wire_peers() + self._empty_shared_gap_bytes = self.full_attn_allocator._current_gap_bytes() + if not legacy_capacities: + self._size_full = self.full_attn_allocator.available_size() + self._size_swa = min( + self.swa_attn_allocator.available_size(), + len(self.full_attn_allocator.free_virtual_ids) * page_size, + ) + # Preserve these private aliases for callers still constructing the + # allocator through the legacy split-capacity API. + self._full_max_total_num_tokens = self._size_full + self._swa_max_total_num_tokens = self._size_swa + # Epoch-keyed memo for the joint capacity view (any chain member's # mutation invalidates -- see `MultiEndedAllocator._chain_capacity_epoch`). self._joint_avail_memo_epoch: Optional[int] = None @@ -3251,15 +3276,15 @@ def __init__( "[unified-memory-pool] UnifiedSWATokenToKVPoolAllocator ready: " "full max_slots=%d (min_slot_index=%d, entry_bytes=%d), " "swa max_slots=%d (min_slot_index=%d, entry_bytes=%d), " - "static caps full=%d swa=%d, joint available=%d", + "max capacity full=%d swa=%d, joint available=%d", self.full_attn_allocator.max_slots, self.full_attn_allocator.min_slot_index, self.full_attn_allocator.entry_bytes, self.swa_attn_allocator.max_slots, self.swa_attn_allocator.min_slot_index, self.swa_attn_allocator.entry_bytes, - self._full_max_total_num_tokens, - self._swa_max_total_num_tokens, + self._size_full, + self._size_swa, self.available_size(), ) @@ -3343,23 +3368,16 @@ def _compute_available_size(self) -> int: K_total = min(K_total, H_f + R_f, H_s + R_s) # index-space caps return K_total * self.page_size - # Slot-conservation views — the ONLY views the leak invariant should see - # (returning the byte-coordinated value would flag spurious leaks). - # `allocated_count()` is in TOKENS (the unit the leak check expects). + # Slot-conservation views for the leak invariant. def _conserve_full_available_size(self) -> int: - return ( - self._full_max_total_num_tokens - self.full_attn_allocator.allocated_count() - ) + return self._size_full - self.full_attn_allocator.allocated_count() def _conserve_swa_available_size(self) -> int: - return ( - self._swa_max_total_num_tokens - self.swa_attn_allocator.allocated_count() - ) + return self._size_swa - self.swa_attn_allocator.allocated_count() # PHYSICAL per-side views read by scheduling / eviction consumers. The - # `min(...)` is sound under dynamic borrowing: the static-conserve cap bounds - # the lending side, the byte-coordinated `schedulable_*` bounds the side that - # has grown into the shared gap; whichever is tighter wins. + # `min(...)` combines each side's physical/index limit with the current + # shared-gap limit. def full_available_size(self) -> int: return min( self._conserve_full_available_size(), @@ -3408,6 +3426,183 @@ def _ask_float_for_room(self, need_tokens: int) -> None: # (set to the static caps). We do NOT report `max_slots - 1`: under unified # memory pool that ~= full_max + swa_max and would over-promise. + @property + def current_full_capacity(self) -> int: + return self.full_available_size() + self.full_attn_allocator.allocated_count() + + @property + def current_swa_capacity(self) -> int: + return self.swa_available_size() + self.swa_attn_allocator.allocated_count() + + def can_reserve( + self, + full_tokens: int | float, + swa_tokens: int | float, + *, + full_evictable_tokens: int = 0, + swa_evictable_tokens: int = 0, + empty_pool: bool = False, + require_token_slack: bool = False, + ) -> bool: + """Check pending FULL/SWA demand against the shared byte envelope. + + Scheduler admission keeps the historical one-token strict slack at an + empty-pool boundary. Actual allocation and eviction planning may consume + the exact physical capacity, so they leave ``require_token_slack`` off. + """ + if not self.supports_asymmetric_reservation: + if ( + full_tokens != swa_tokens + or full_evictable_tokens + or swa_evictable_tokens + or empty_pool + ): + return False + return full_tokens <= self.available_size() + if full_tokens < 0 or swa_tokens < 0: + return False + if require_token_slack and ( + full_tokens >= self.size_full or swa_tokens >= self.size_swa + ): + return False + + page_size = self.page_size + full_pages = (math.ceil(full_tokens) + page_size - 1) // page_size + swa_pages = (math.ceil(swa_tokens) + page_size - 1) // page_size + compacted = empty_pool or not self.lazy_compaction + if not compacted: + compacted = self._compaction_allowed() + return self._fits_page_demand( + full_pages, + swa_pages, + full_reclaim_pages=max(0, int(full_evictable_tokens)) // page_size, + swa_reclaim_pages=max(0, int(swa_evictable_tokens)) // page_size, + compacted=compacted, + empty_pool=empty_pool, + ) + + def _compaction_allowed(self) -> bool: + return all( + allocator.disagg_move_gate is None or allocator.disagg_move_gate() + for allocator in (self.full_attn_allocator, self.swa_attn_allocator) + ) + + def _fits_page_demand( + self, + num_full_pages: int, + num_swa_pages: int, + *, + full_reclaim_pages: int = 0, + swa_reclaim_pages: int = 0, + compacted: bool, + empty_pool: bool = False, + ) -> bool: + """Check one FULL/SWA page demand against a single allocator snapshot.""" + if min(num_full_pages, num_swa_pages) < 0: + return False + + fa, sa = self.full_attn_allocator, self.swa_attn_allocator + if empty_pool: + full_live_pages = swa_live_pages = 0 + full_reclaim_pages = swa_reclaim_pages = 0 + else: + full_live_pages = fa.allocated_count() // self.page_size + swa_live_pages = sa.allocated_count() // self.page_size + full_reclaim_pages = min(full_live_pages, max(0, int(full_reclaim_pages))) + swa_reclaim_pages = min(swa_live_pages, max(0, int(swa_reclaim_pages))) + + full_live_pages -= full_reclaim_pages + swa_live_pages -= swa_reclaim_pages + full_total_pages = full_live_pages + num_full_pages + swa_total_pages = swa_live_pages + num_swa_pages + virtual_page_capacity = fa.num_virtual_ids - fa.min_page_index + full_page_capacity = min( + virtual_page_capacity, + fa.num_pages - fa.min_page_index, + ) + swa_page_capacity = min( + virtual_page_capacity, + sa.num_pages - sa.min_page_index, + ) + if ( + full_total_pages > full_page_capacity + or swa_total_pages > swa_page_capacity + or swa_total_pages > full_total_pages + ): + return False + + if compacted: + full_virtual_room = virtual_page_capacity - full_live_pages + full_holes = swa_holes = 0 + full_index_room = full_page_capacity - full_live_pages + swa_index_room = swa_page_capacity - swa_live_pages + gap_bytes = ( + self._empty_shared_gap_bytes + - full_live_pages * fa.entry_bytes_per_page + - swa_live_pages * sa.entry_bytes_per_page + ) + else: + full_virtual_room = len(fa.free_virtual_ids) + full_reclaim_pages + full_holes = len(fa._free_phys_pages) + full_reclaim_pages + swa_holes = len(sa._free_phys_pages) + swa_reclaim_pages + full_index_room = fa.num_pages - fa.min_page_index - fa._allocated_pages() + swa_index_room = sa.num_pages - sa.min_page_index - sa._allocated_pages() + gap_bytes = fa._current_gap_bytes() + + if num_full_pages > full_virtual_room: + return False + if num_full_pages > full_holes + full_index_room: + return False + if num_swa_pages > swa_holes + swa_index_room: + return False + full_extensions = max(0, num_full_pages - full_holes) + swa_extensions = max(0, num_swa_pages - swa_holes) + return ( + full_extensions * fa.entry_bytes_per_page + + swa_extensions * sa.entry_bytes_per_page + <= max(0, gap_bytes) + ) + + def ensure_capacity(self, full_tokens: int, swa_tokens: int) -> bool: + """Gate one allocation and compact both sides on shortfall.""" + if full_tokens < 0 or swa_tokens < 0: + return False + if full_tokens == 0 and swa_tokens == 0: + return True + if not self.supports_asymmetric_reservation: + if full_tokens != swa_tokens: + return False + need_tokens = int(full_tokens) + if need_tokens <= self.available_size(): + return True + return _relieve_for_alloc(self, need_tokens) + page_size = self.page_size + num_full_pages = (int(full_tokens) + page_size - 1) // page_size + num_swa_pages = (int(swa_tokens) + page_size - 1) // page_size + if num_swa_pages > num_full_pages: + return False + if self._fits_page_demand( + num_full_pages, + num_swa_pages, + compacted=False, + ): + return True + if not self.lazy_compaction or not self._compaction_allowed(): + return False + if not self._fits_page_demand( + num_full_pages, + num_swa_pages, + compacted=True, + ): + return False + self.full_attn_allocator._flush(urgent=True) + self.swa_attn_allocator._flush(urgent=True) + return self._fits_page_demand( + num_full_pages, + num_swa_pages, + compacted=False, + ) + @property def draft_virtual_id_space(self) -> int: return self.full_attn_allocator.max_slots - 1 @@ -3490,11 +3685,8 @@ def swa_v2p_page_table(self) -> torch.Tensor: def alloc(self, need_size: int) -> Optional[torch.Tensor]: with record_function("UnifiedSWAAlloc.alloc"): - # Joint pre-check. Both sides are mutual peers (each side's compaction - # opens gap for the other), so flush BOTH on shortfall. - if need_size > self.available_size(): - if not _relieve_for_alloc(self, need_size): - return None + if not self.ensure_capacity(need_size, need_size): + return None # Snapshot the virtual PAGES full will consume, to bind them on swa too. num_pages = need_size // self.page_size fa = self.full_attn_allocator @@ -3530,9 +3722,8 @@ def alloc_extend( prefix_lens=prefix_lens_cpu, ) need_tokens = num_new_pages * self.page_size - if need_tokens > self.available_size(): - if not _relieve_for_alloc(self, need_tokens): - return None + if not self.ensure_capacity(need_tokens, need_tokens): + return None # Snapshot the virtual PAGES the kernel will consume; clone so swa keeps # its view after the slice is consumed. @@ -3569,9 +3760,8 @@ def alloc_decode( seq_lens=seq_lens_cpu, page_size=self.page_size, decode=True ) need_tokens = num_new_pages * self.page_size - if need_tokens > self.available_size(): - if not _relieve_for_alloc(self, need_tokens): - return None + if not self.ensure_capacity(need_tokens, need_tokens): + return None fa = self.full_attn_allocator new_virtual_pages = fa.free_virtual_ids[:num_new_pages].clone() @@ -3828,6 +4018,8 @@ class UnifiedMambaSWATokenToKVPoolAllocator(UnifiedSWATokenToKVPoolAllocator): like the 2-pool mamba composite. """ + supports_asymmetric_reservation = False + def __init__( self, *, diff --git a/python/sglang/srt/mem_cache/unified_memory_pool.py b/python/sglang/srt/mem_cache/unified_memory_pool.py index abe5474500ff..58a09de84f93 100644 --- a/python/sglang/srt/mem_cache/unified_memory_pool.py +++ b/python/sglang/srt/mem_cache/unified_memory_pool.py @@ -1665,8 +1665,8 @@ def init_unified_swa_pools( end_layer: int, swa_attention_layer_ids: List[int], full_attention_layer_ids: List[int], - full_max_total_num_tokens: int, - swa_max_total_num_tokens: int, + full_max_total_num_tokens: Optional[int] = None, + swa_max_total_num_tokens: Optional[int] = None, enable_memory_saver: bool, need_sort: bool, forward_stream: Optional[torch.cuda.Stream] = None, @@ -1710,15 +1710,24 @@ def init_unified_swa_pools( store_dtype=store_dtype, grow_direction="up", ) + legacy_allocator_capacities = {} if unified_total_bytes is not None: # PROFILED byte budget, sized from directly: the re-sum's floor losses # stay out of the buffer, and the token counts remain boot labels. total_bytes = unified_total_bytes else: + if full_max_total_num_tokens is None or swa_max_total_num_tokens is None: + raise ValueError( + "unified_total_bytes or both legacy full/SWA capacities must be provided" + ) total_bytes = ( full_max_total_num_tokens * full_spec.entry_bytes() + swa_max_total_num_tokens * swa_spec.entry_bytes() ) + legacy_allocator_capacities = { + "full_max_total_num_tokens": full_max_total_num_tokens, + "swa_max_total_num_tokens": swa_max_total_num_tokens, + } if model_context_len is not None: # bs=1 floor: ONE sliding window of swa KV (+ a page of slack for the # page-granular walk) + the slot-0 sink. The full side is not charged @@ -1737,6 +1746,8 @@ def init_unified_swa_pools( ], factory="init_unified_swa_pools", ) + if total_bytes <= 0: + raise ValueError(f"total_bytes must be positive, got {total_bytes}") shared_pool = UnifiedKVPool( total_bytes=total_bytes, sub_pool_specs=[full_spec, swa_spec], @@ -1757,12 +1768,11 @@ def init_unified_swa_pools( unified_buffer=shared_pool, kvcache=token_to_kv_pool, device=device, - full_max_total_num_tokens=full_max_total_num_tokens, - swa_max_total_num_tokens=swa_max_total_num_tokens, page_size=page_size, need_sort=need_sort, forward_stream=forward_stream, lazy_compaction=lazy_compaction, + **legacy_allocator_capacities, ) logger.info( @@ -1793,12 +1803,12 @@ def init_unified_swa_pools( page_size, ) logger.info( - "[unified-memory-pool] total_bytes=%d (=%.2f GB), full_max_total_num_tokens=%d, " - "swa_max_total_num_tokens=%d, joint_available=%d slots", + "[unified-memory-pool] total_bytes=%d (=%.2f GB), " + "full_capacity=%d, swa_capacity=%d, joint_available=%d slots", total_bytes, total_bytes / GB, - full_max_total_num_tokens, - swa_max_total_num_tokens, + allocator.size_full, + allocator.size_swa, allocator.available_size(), ) logger.info( diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 95b4118b17c1..bc49a99e95b2 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -1350,7 +1350,14 @@ def unload_lora_adapter(self, lora_ref: LoRARef): def effective_max_total_num_tokens(self): """Return the max token pool size considering hybrid swa settings.""" if self.is_hybrid_swa: - capacity = self.full_max_total_num_tokens or self.swa_max_total_num_tokens + if self.server_args.enable_unified_memory: + capacity = self.token_to_kv_pool_allocator.size_full + if self.server_args.max_total_tokens is not None: + capacity = min(capacity, self.server_args.max_total_tokens) + else: + capacity = ( + self.full_max_total_num_tokens or self.swa_max_total_num_tokens + ) else: capacity = self.max_total_num_tokens if (req_to_token_pool := getattr(self, "req_to_token_pool", None)) is not None: diff --git a/python/sglang/srt/model_executor/pool_configurator.py b/python/sglang/srt/model_executor/pool_configurator.py index 407afea59de6..e9918338f260 100644 --- a/python/sglang/srt/model_executor/pool_configurator.py +++ b/python/sglang/srt/model_executor/pool_configurator.py @@ -63,6 +63,7 @@ class MemoryPoolConfig: max_running_requests: Optional[int] = None full_max_total_num_tokens: Optional[int] = None swa_max_total_num_tokens: Optional[int] = None + unified_memory_pool_bytes: Optional[int] = None # DSV4 compressed-attention pool sizes (target only; draft workers leave at 0). c4_max_total_num_tokens: int = 0 @@ -459,6 +460,14 @@ def __init__(self, kvc: KVCacheConfigurator): self._full_layers_num = len(model_config.full_attention_layer_ids) self._swa_layers_num = len(model_config.swa_attention_layer_ids) + self._target_full_layers_num = sum( + kvc.layer_info.start_layer <= layer_id < kvc.layer_info.end_layer + for layer_id in model_config.full_attention_layer_ids + ) + self._target_swa_layers_num = sum( + kvc.layer_info.start_layer <= layer_id < kvc.layer_info.end_layer + for layer_id in model_config.swa_attention_layer_ids + ) assert ( self._swa_layers_num > 0 ), "Hybrid SWA model must have at least one SWA layer" @@ -592,6 +601,12 @@ def _draft_pool_bytes_per_token(self) -> int: + self._draft_cell_size ) + def _target_unified_bytes(self, full_tokens: int, swa_tokens: int) -> int: + return ( + full_tokens * self._full_per_token * self._target_full_layers_num + + swa_tokens * self._swa_per_token * self._target_swa_layers_num + ) + def _max_unified_full_tokens( self, available_bytes: int, @@ -600,8 +615,9 @@ def _max_unified_full_tokens( ) -> int: """Find the largest page-aligned full capacity whose allocations fit.""" draft_bytes_per_token = self._draft_pool_bytes_per_token() - target_full_bytes_per_token = self._full_per_token * self._full_layers_num - target_swa_bytes_per_token = self._swa_per_token * self._swa_layers_num + target_full_bytes_per_token = ( + self._full_per_token * self._target_full_layers_num + ) assert target_full_bytes_per_token > 0 def allocation_bytes(full_pages: int) -> int: @@ -613,10 +629,7 @@ def allocation_bytes(full_pages: int) -> int: // page_size * page_size ) - target_bytes = ( - full_tokens * target_full_bytes_per_token - + swa_tokens * target_swa_bytes_per_token - ) + target_bytes = self._target_unified_bytes(full_tokens, swa_tokens) virtual_span = max(target_bytes // target_full_bytes_per_token - 1, 0) draft_tokens = ceil_align(virtual_span, page_size) + page_size return target_bytes + draft_tokens * draft_bytes_per_token @@ -674,16 +687,17 @@ def align_page_size(x: int) -> int: max_total_num_tokens=full_tokens, full_max_total_num_tokens=full_tokens, swa_max_total_num_tokens=swa_tokens, + unified_memory_pool_bytes=( + self._target_unified_bytes(full_tokens, swa_tokens) + if self._enable_unified_memory + else None + ), ) def calculate_pool_sizes( self, available_bytes: int, page_size: int ) -> MemoryPoolConfig: - if ( - self._enable_unified_memory - and self._full_layers_num > 0 - and self._draft_pool_bytes_per_token() > 0 - ): + if self._enable_unified_memory and self._full_layers_num > 0: max_total_num_tokens = self._max_unified_full_tokens( available_bytes, page_size ) @@ -774,7 +788,7 @@ def calculate_pool_sizes( * self._swa_per_token * (self._swa_layers_num + self._draft_swa_layers_num) ) - if self._enable_unified_memory and self._draft_pool_bytes_per_token() > 0: + if self._enable_unified_memory: full_tokens = self._max_unified_full_tokens( available_bytes, page_size, fixed_swa_tokens=swa_tokens ) @@ -799,6 +813,11 @@ def calculate_pool_sizes( max_total_num_tokens=full_tokens, full_max_total_num_tokens=full_tokens, swa_max_total_num_tokens=swa_tokens, + unified_memory_pool_bytes=( + self._target_unified_bytes(full_tokens, swa_tokens) + if self._enable_unified_memory + else None + ), ) def calculate_pool_sizes_from_max_tokens( @@ -811,6 +830,13 @@ def calculate_pool_sizes_from_max_tokens( max_total_num_tokens=full_tokens, full_max_total_num_tokens=full_tokens, swa_max_total_num_tokens=min(swa_tokens, max_total_num_tokens), + unified_memory_pool_bytes=( + self._target_unified_bytes( + full_tokens, min(swa_tokens, max_total_num_tokens) + ) + if self._enable_unified_memory + else None + ), ) diff --git a/test/registered/unit/disaggregation/test_decode_queue_cleanup.py b/test/registered/unit/disaggregation/test_decode_queue_cleanup.py index 7b5b9f8a06a0..b037216213cb 100644 --- a/test/registered/unit/disaggregation/test_decode_queue_cleanup.py +++ b/test/registered/unit/disaggregation/test_decode_queue_cleanup.py @@ -13,6 +13,9 @@ from sglang.srt.distributed.parallel_state_wrapper import ParallelState from sglang.srt.managers.schedule_batch import FINISH_ABORT from sglang.srt.managers.scheduler import Scheduler +from sglang.srt.mem_cache.multi_ended_allocator import ( + UnifiedSWATokenToKVPoolAllocator, +) from sglang.srt.runtime_context import get_context from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -44,7 +47,8 @@ def test_paged_swa_retraction_resume_uses_physical_page_budget(self): page_size = 128 fill_len = 574 physical_tokens_per_req = 5 * page_size - physical_available = 18 * page_size + full_available = 18 * page_size + swa_available = 4 * page_size reqs = [ SimpleNamespace( @@ -62,7 +66,10 @@ def test_paged_swa_retraction_resume_uses_physical_page_budget(self): queue.retracted_queue = reqs.copy() queue.num_reserved_decode_tokens = 0 queue.req_to_token_pool = SimpleNamespace(available_size=lambda: len(reqs)) - queue.token_to_kv_pool_allocator = SimpleNamespace(page_size=page_size) + queue.token_to_kv_pool_allocator = MagicMock( + spec=UnifiedSWATokenToKVPoolAllocator + ) + queue.token_to_kv_pool_allocator.page_size = page_size queue.tree_cache = MagicMock() queue.scheduler = SimpleNamespace( sliding_window_size=2047, @@ -70,16 +77,21 @@ def test_paged_swa_retraction_resume_uses_physical_page_budget(self): ) queue._uses_swa_tail_prealloc = MagicMock(return_value=True) queue._swa_aware_allocatable_token_budgets = MagicMock( - return_value=(physical_available, physical_available) + return_value=(full_available, swa_available) ) queue._swa_tail_allocatable_token_budget = MagicMock( - side_effect=lambda **_: physical_available + side_effect=lambda **_: swa_available + ) + queue._unified_swa_reservation_fits = MagicMock( + side_effect=lambda full_tokens, _swa_tokens, *, full_allocatable_tokens, **_: full_tokens + <= full_allocatable_tokens ) + queue._reclaim_swa_tail_capacity = MagicMock(return_value=None) def pre_alloc(_req): - nonlocal physical_available - self.assertGreaterEqual(physical_available, physical_tokens_per_req) - physical_available -= physical_tokens_per_req + nonlocal full_available + self.assertGreaterEqual(full_available, physical_tokens_per_req) + full_available -= physical_tokens_per_req queue._pre_alloc = MagicMock(side_effect=pre_alloc) @@ -87,8 +99,14 @@ def pre_alloc(_req): self.assertEqual(resumed, reqs[:3]) self.assertEqual(queue.retracted_queue, reqs[3:]) - self.assertEqual(physical_available, 3 * page_size) + self.assertEqual(full_available, 3 * page_size) self.assertEqual(queue._pre_alloc.call_count, 3) + self.assertTrue( + all( + call.args[1] > call.kwargs["swa_allocatable_tokens"] + for call in queue._unified_swa_reservation_fits.call_args_list + ) + ) def test_prealloc_abort_clears_receiver_before_removing_request(self): receiver = FakeReceiver() @@ -212,6 +230,7 @@ def test_swa_reclaim_failure_rejects_only_request(self): ) queue._hicache_pending_restore_tokens = MagicMock(return_value=0) queue._pre_alloc = MagicMock() + queue.token_to_kv_pool_allocator = MagicMock() queue.req_to_token_pool = MagicMock() queue.req_to_token_pool.available_size.return_value = 1 queue.req_to_metadata_buffer_idx_allocator = MagicMock() diff --git a/test/registered/unit/managers/test_prefill_adder.py b/test/registered/unit/managers/test_prefill_adder.py index 6525c3ce9bf5..3edf87fdbc11 100644 --- a/test/registered/unit/managers/test_prefill_adder.py +++ b/test/registered/unit/managers/test_prefill_adder.py @@ -14,6 +14,9 @@ IncLockRefResult, ) from sglang.srt.runtime_context import get_context +from sglang.srt.mem_cache.multi_ended_allocator import ( + UnifiedSWATokenToKVPoolAllocator, +) from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler from sglang.srt.utils.common import Range from sglang.test.ci.ci_register import register_cpu_ci @@ -521,6 +524,48 @@ def test_add_chunked_req_hybrid_swa_defers_when_swa_below_page(self): req.set_extend_range.assert_not_called() self.assertEqual(len(adder.can_run_list), 0) + def test_add_one_req_unified_swa_uses_shared_capacity(self): + allocator = MagicMock(spec=UnifiedSWATokenToKVPoolAllocator) + allocator.available_size.return_value = 0 + allocator.full_available_size.return_value = 0 + allocator.swa_available_size.return_value = 0 + allocator.size_swa = 1_000_000 + allocator.can_reserve.side_effect = ( + lambda full_tokens, swa_tokens, *, full_evictable_tokens=0, swa_evictable_tokens=0, **_: max( + full_tokens, swa_tokens + ) + <= full_evictable_tokens + swa_evictable_tokens + ) + self.mock_token_allocator = allocator + self.mock_tree_cache.swa_evictable_size.return_value = 64 + self.mock_tree_cache.sliding_window_size = 128 + adder = self.create_adder(self.create_running_batch(), page_size=8) + + req = self.create_mock_req("unified", priority=0, max_new_tokens=16) + req.origin_input_ids = list(range(24)) + req.full_untruncated_fill_ids = list(range(24)) + req.swa_host_hit_length = 0 + req.last_node = MagicMock() + req.sampling_params.ignore_eos = False + req.set_extend_range = MagicMock( + side_effect=lambda start, end: setattr( + req, "extend_range", Range(start, end) + ) + ) + + result = adder.add_one_req( + req, has_chunked_req=False, truncation_align_size=None + ) + + self.assertEqual(result, AddReqResult.CONTINUE) + self.assertIn(req, adder.can_run_list) + self.assertTrue( + any( + call.kwargs["swa_evictable_tokens"] == 64 + for call in allocator.can_reserve.call_args_list + ) + ) + def test_swa_budget_for_req(self): # budget = max(alloc - window, 0) + min(extend + max_new, window) + page, # where alloc = min(extend, rem_chunk). The decode headroom is the SWA the diff --git a/test/registered/unit/managers/test_scheduler_init_req_max_new_tokens.py b/test/registered/unit/managers/test_scheduler_init_req_max_new_tokens.py index 93f41fb5cec1..9e3c62eccbe5 100644 --- a/test/registered/unit/managers/test_scheduler_init_req_max_new_tokens.py +++ b/test/registered/unit/managers/test_scheduler_init_req_max_new_tokens.py @@ -56,6 +56,7 @@ def _new_scheduler( scheduler.max_total_num_tokens = max_total_num_tokens scheduler.page_size = page_size scheduler.max_new_tokens_limit = envs.SGLANG_MAX_NEW_TOKENS_LIMIT.get() + scheduler.token_to_kv_pool_allocator = None return scheduler def _new_req(self, max_new_tokens, input_len: int = 8, min_new_tokens: int = 0): diff --git a/test/registered/unit/mem_cache/test_decode_radix_lock_ref.py b/test/registered/unit/mem_cache/test_decode_radix_lock_ref.py index 73528e65bf91..3cba68499020 100644 --- a/test/registered/unit/mem_cache/test_decode_radix_lock_ref.py +++ b/test/registered/unit/mem_cache/test_decode_radix_lock_ref.py @@ -38,6 +38,9 @@ InsertParams, MatchPrefixParams, ) +from sglang.srt.mem_cache.multi_ended_allocator import ( + UnifiedSWATokenToKVPoolAllocator, +) from sglang.srt.mem_cache.radix_cache import RadixCache, RadixKey from sglang.srt.utils.common import Range @@ -151,6 +154,15 @@ def test_reclaim_swa_tail_capacity_page_rounds(self): self.assertEqual(params.num_tokens, 0) self.assertEqual(params.swa_num_tokens, 128) + unified_allocator = MagicMock(spec=UnifiedSWATokenToKVPoolAllocator) + unified_allocator.page_size = 64 + unified_allocator.ensure_capacity.return_value = True + queue.token_to_kv_pool_allocator = unified_allocator + queue.tree_cache = MagicMock() + + self.assertIsNone(queue._reclaim_swa_tail_capacity(129, "unified", full_len=65)) + unified_allocator.ensure_capacity.assert_called_once_with(128, 192) + def test_reclaim_swa_tail_capacity_fails_before_allocation(self): queue = DecodePreallocQueue.__new__(DecodePreallocQueue) queue.token_to_kv_pool_allocator = MagicMock(page_size=64) @@ -449,7 +461,7 @@ def test_pop_preallocated_rechecks_budget_after_lock(self): skip_swa=True, ) self.assertFalse(req.swa_prefix_lock_released) - queue._swa_tail_len.assert_called_once_with(8) + queue._swa_tail_len.assert_called_with(8) queue._allocatable_token_budgets.assert_called_once() def test_repeated_incremental_no_leak(self): diff --git a/test/registered/unit/model_executor/test_pool_configurator.py b/test/registered/unit/model_executor/test_pool_configurator.py index 3412a18d2748..915dcf822693 100644 --- a/test/registered/unit/model_executor/test_pool_configurator.py +++ b/test/registered/unit/model_executor/test_pool_configurator.py @@ -83,6 +83,7 @@ def _make_model_runner( disaggregation_mode="null", max_running_requests=None, disaggregation_decode_extra_slots=0, + enable_unified_memory=False, kv_lora_rank=512, qk_rope_head_dim=64, swa_kv_lora_rank=128, @@ -148,6 +149,7 @@ def _make_model_runner( disaggregation_mode=disaggregation_mode, max_running_requests=max_running_requests, disaggregation_decode_extra_slots=disaggregation_decode_extra_slots, + enable_unified_memory=enable_unified_memory, enable_hisparse=False, enable_hierarchical_cache=False, enable_dsa_cache_layer_split=False, @@ -297,7 +299,14 @@ def test_dsa_mla_cell_size_uses_backend_kv_layout( class TestHybridSWAConfigurator(CustomTestCase): """Hybrid SWA: full/swa split, ratio, memory invariant.""" - def _make_swa_runner(self, full_layers=16, swa_layers=16, ratio=0.5, page_size=1): + def _make_swa_runner( + self, + full_layers=16, + swa_layers=16, + ratio=0.5, + page_size=1, + enable_unified_memory=False, + ): return _make_model_runner( self, is_hybrid_swa=True, @@ -306,6 +315,7 @@ def _make_swa_runner(self, full_layers=16, swa_layers=16, ratio=0.5, page_size=1 swa_num_kv_heads=4, page_size=page_size, swa_full_tokens_ratio=ratio, + enable_unified_memory=enable_unified_memory, ) def _run(self, available_bytes, **kwargs): @@ -327,6 +337,14 @@ def test_memory_utilization(self): self.assertLessEqual(used, available) self.assertGreater(used, available * 0.99) + def test_unified_memory_records_exact_target_bytes(self): + available = 10_000_000 + mr, _, config = self._run(available, enable_unified_memory=True) + used = _actual_memory_used(mr, config) + self.assertEqual(config.unified_memory_pool_bytes, used) + self.assertLessEqual(used, available) + self.assertGreater(used, available * 0.99) + @patch( "sglang.srt.mem_cache.kv_cache_configurator.calculate_mla_kv_cache_dim", return_value=576, From 1aac878722116da3b299d1d9818aece1dd309d33 Mon Sep 17 00:00:00 2001 From: Yonghao Zhuang Date: Thu, 27 Aug 2026 16:12:48 -0700 Subject: [PATCH 04/55] Avoid double-rounding unified SWA demand --- python/sglang/srt/managers/schedule_policy.py | 16 +++++++------- .../unit/managers/test_prefill_adder.py | 21 ++++++++++++------- 2 files changed, 22 insertions(+), 15 deletions(-) diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index 44490f9300dd..13ca997369dd 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -885,13 +885,13 @@ def _check_swa_admission( self, req: Req, *, - real_input_tokens: int, + extend_input_len: int, total_tokens: int, ) -> tuple[bool, Optional[int]]: """Check SWA pressure and return the usable chunk limit.""" max_new_tokens = self._swa_new_tokens(req) swa_needed = self._swa_budget_for_req( - real_input_tokens, + extend_input_len, max_new_tokens, swa_host_hit_length=req.swa_host_hit_length, ) @@ -915,7 +915,7 @@ def _check_swa_admission( full_input_tokens, max_new_tokens ) if not self._swa_req_never_fits( - real_input_tokens, + extend_input_len, max_new_tokens, req.swa_host_hit_length, full_tokens=full_ever_tokens, @@ -926,7 +926,7 @@ def _check_swa_admission( max_chunk_tokens = self.rem_chunk_tokens chunk_max_new = max_new_tokens if self.is_unified_swa: - max_chunk_tokens = min(max_chunk_tokens or 0, max(0, real_input_tokens - 1)) + max_chunk_tokens = min(max_chunk_tokens or 0, max(0, extend_input_len - 1)) chunk_max_new = 0 swa_cap = self._swa_chunk_cap( chunk_max_new, @@ -1381,8 +1381,8 @@ def add_one_req( # `total_tokens` so both `rem_total_tokens` gates reflect the joint budget. total_tokens += self._mamba_gap_budget_for_req(req) # adjusting the input_tokens based on host_hit_length and page_size - real_input_tokens = cand_extend_input_len - req.host_hit_length - real_input_tokens = self.ceil_paged_tokens(real_input_tokens) + swa_extend_input_len = cand_extend_input_len - req.host_hit_length + real_input_tokens = self.ceil_paged_tokens(swa_extend_input_len) prefix_len = len(req.prefix_indices) if not self.is_unified_swa and total_tokens >= self.rem_total_tokens: @@ -1392,7 +1392,7 @@ def add_one_req( if self.is_hybrid_swa: can_admit, chunk_tokens_limit = self._check_swa_admission( req, - real_input_tokens=real_input_tokens, + extend_input_len=swa_extend_input_len, total_tokens=total_tokens, ) if not can_admit: @@ -1416,7 +1416,7 @@ def add_one_req( if self.is_hybrid_swa: can_admit, chunk_tokens_limit = self._check_swa_admission( req, - real_input_tokens=real_input_tokens, + extend_input_len=swa_extend_input_len, total_tokens=total_tokens, ) if not can_admit: diff --git a/test/registered/unit/managers/test_prefill_adder.py b/test/registered/unit/managers/test_prefill_adder.py index 3edf87fdbc11..8512a0a3cf90 100644 --- a/test/registered/unit/managers/test_prefill_adder.py +++ b/test/registered/unit/managers/test_prefill_adder.py @@ -525,6 +525,10 @@ def test_add_chunked_req_hybrid_swa_defers_when_swa_below_page(self): self.assertEqual(len(adder.can_run_list), 0) def test_add_one_req_unified_swa_uses_shared_capacity(self): + # The real warmup is shorter than one page. Feeding an already + # page-rounded input into the SWA budget would round it a second time + # after adding decode headroom and reject the first request forever. + page_size = 128 allocator = MagicMock(spec=UnifiedSWATokenToKVPoolAllocator) allocator.available_size.return_value = 0 allocator.full_available_size.return_value = 0 @@ -537,13 +541,13 @@ def test_add_one_req_unified_swa_uses_shared_capacity(self): <= full_evictable_tokens + swa_evictable_tokens ) self.mock_token_allocator = allocator - self.mock_tree_cache.swa_evictable_size.return_value = 64 - self.mock_tree_cache.sliding_window_size = 128 - adder = self.create_adder(self.create_running_batch(), page_size=8) + self.mock_tree_cache.swa_evictable_size.return_value = 3 * page_size + self.mock_tree_cache.sliding_window_size = 4096 + adder = self.create_adder(self.create_running_batch(), page_size=page_size) - req = self.create_mock_req("unified", priority=0, max_new_tokens=16) - req.origin_input_ids = list(range(24)) - req.full_untruncated_fill_ids = list(range(24)) + req = self.create_mock_req("unified", priority=0, max_new_tokens=10) + req.origin_input_ids = list(range(3)) + req.full_untruncated_fill_ids = list(range(3)) req.swa_host_hit_length = 0 req.last_node = MagicMock() req.sampling_params.ignore_eos = False @@ -559,9 +563,12 @@ def test_add_one_req_unified_swa_uses_shared_capacity(self): self.assertEqual(result, AddReqResult.CONTINUE) self.assertIn(req, adder.can_run_list) + full_tokens, swa_tokens = allocator.can_reserve.call_args_list[0].args + self.assertEqual(full_tokens, 3 + 10 + page_size) + self.assertEqual(swa_tokens, 3 + 10 + page_size) self.assertTrue( any( - call.kwargs["swa_evictable_tokens"] == 64 + call.kwargs["swa_evictable_tokens"] == 3 * page_size for call in allocator.can_reserve.call_args_list ) ) From 7ee2c747e3cde07f512bf5e9648815e6ea02d84f Mon Sep 17 00:00:00 2001 From: yhzhuang Date: Thu, 27 Aug 2026 16:50:15 -0700 Subject: [PATCH 05/55] Reject unsupported unified pool scheduler modes Co-authored-by: yhzhuang --- python/sglang/srt/arg_groups/kv_cache_hook.py | 12 +++++++----- python/sglang/srt/mem_cache/common.py | 2 +- .../srt/mem_cache/kv_cache_configurator.py | 16 +++++++--------- .../srt/mem_cache/multi_ended_allocator.py | 15 +++++++++------ .../sglang/srt/mem_cache/unified_memory_pool.py | 14 ++++++++------ .../unit/server_args/test_server_args.py | 17 +++++++++++++++++ 6 files changed, 49 insertions(+), 27 deletions(-) diff --git a/python/sglang/srt/arg_groups/kv_cache_hook.py b/python/sglang/srt/arg_groups/kv_cache_hook.py index 6567f0a8e08e..a25f4214389b 100644 --- a/python/sglang/srt/arg_groups/kv_cache_hook.py +++ b/python/sglang/srt/arg_groups/kv_cache_hook.py @@ -198,6 +198,13 @@ def handle_unified_memory_pool(server_args: Any) -> None: cfg = resolving_view(server_args) if not cfg.enable_unified_memory: return + assert cfg.pp_size == 1, ( + "--enable-unified-memory does not support pipeline parallelism " + "(--pp-size > 1)." + ) + assert ( + not cfg.enable_mixed_chunk + ), "--enable-unified-memory does not support --enable-mixed-chunk." if cfg.disaggregation_mode != "null": # Constraints of the whole-envelope transfer; see # UnifiedMLATokenToKVPool.get_contiguous_buf_infos. @@ -206,11 +213,6 @@ def handle_unified_memory_pool(server_args: Any) -> None: "the mooncake transfer backend; got " f"{cfg.disaggregation_transfer_backend!r}." ) - assert cfg.pp_size == 1, ( - "--enable-unified-memory with PD disaggregation does not support " - "pipeline parallelism (whole-envelope transfer has no per-layer " - "entries to subset)." - ) assert not envs.SGLANG_DISABLE_LAZY_COMPACTION.get(), ( "--enable-unified-memory with PD disaggregation requires lazy " "compaction; unset SGLANG_DISABLE_LAZY_COMPACTION." diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index ac07bdeeb282..1d0fcc231075 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/python/sglang/srt/mem_cache/common.py @@ -234,7 +234,7 @@ def minimum_reclaim(max_tokens: int, fits) -> int: swa_evictable_tokens=swa_reclaim, ), ) - evicted = tree_cache.evict( + evicted = tree_cache.evict_for_alloc( EvictParams( num_tokens=full_reclaim, swa_num_tokens=swa_reclaim, diff --git a/python/sglang/srt/mem_cache/kv_cache_configurator.py b/python/sglang/srt/mem_cache/kv_cache_configurator.py index a8d887494135..515919c2ba92 100644 --- a/python/sglang/srt/mem_cache/kv_cache_configurator.py +++ b/python/sglang/srt/mem_cache/kv_cache_configurator.py @@ -234,6 +234,7 @@ class _PoolSizes(msgspec.Struct, frozen=True, kw_only=True): c128_state_pool_size: int c4_state_dtype: Optional[torch.dtype] c128_state_dtype: Optional[torch.dtype] + unified_memory_pool_bytes: Optional[int] = None unified_total_bytes: Optional[int] = None @@ -389,6 +390,7 @@ def _derive_pool_sizes(self, *, config: MemoryPoolConfig) -> _PoolSizes: max_running_requests=max_running_requests, full_max_total_num_tokens=full_max_total_num_tokens, swa_max_total_num_tokens=swa_max_total_num_tokens, + unified_memory_pool_bytes=config.unified_memory_pool_bytes, c4_max_total_num_tokens=c4_max_total_num_tokens, c128_max_total_num_tokens=c128_max_total_num_tokens, c4_state_pool_size=c4_state_pool_size, @@ -460,9 +462,7 @@ def _init_pools( ) bundle = self._init_unified_swa_pools( max_num_reqs=sizes.max_running_requests, - full_max_total_num_tokens=sizes.full_max_total_num_tokens, - swa_max_total_num_tokens=sizes.swa_max_total_num_tokens, - unified_total_bytes=sizes.unified_total_bytes, + unified_memory_pool_bytes=sizes.unified_memory_pool_bytes, ) else: # Fail loud, not silently fall through to the normal pools (which would @@ -807,9 +807,9 @@ def _init_unified_swa_pools( self, *, max_num_reqs: int, - full_max_total_num_tokens: Optional[int], - swa_max_total_num_tokens: Optional[int], - unified_total_bytes: Optional[int] = None, + full_max_total_num_tokens: Optional[int] = None, + swa_max_total_num_tokens: Optional[int] = None, + unified_memory_pool_bytes: Optional[int] = None, ) -> UnifiedPoolBundle: """Build the unified-pool stack for a hybrid-SWA model (Triton): one byte buffer split between the full-attention and SWA KV pools.""" @@ -885,6 +885,7 @@ def _init_unified_swa_pools( full_attention_layer_ids=full_attention_layer_ids, full_max_total_num_tokens=full_max_total_num_tokens, swa_max_total_num_tokens=swa_max_total_num_tokens, + total_bytes=unified_memory_pool_bytes, enable_memory_saver=get_exec().features.enable_memory_saver, need_sort=get_disagg().disaggregation_mode in ("decode", "prefill"), # Overlap mode: same wait_stream(forward_stream) rationale as @@ -892,9 +893,6 @@ def _init_unified_swa_pools( forward_stream=self.forward_stream, # Lazy compaction: default ON, with env var escape hatch for rollback / A/B. lazy_compaction=_should_enable_lazy_compaction(), - # Draft workers keep the token-count byte sum (spec is asserted - # off under unified; belt only). - unified_total_bytes=(None if self.is_draft_worker else unified_total_bytes), # bs=1 feasibility floor inputs. `model_context_len` bounds the # sliding window term only -- the full-attention side is not # charged, see `_check_bs1_feasibility_floor`. diff --git a/python/sglang/srt/mem_cache/multi_ended_allocator.py b/python/sglang/srt/mem_cache/multi_ended_allocator.py index d4b40a126ead..479babc6366e 100644 --- a/python/sglang/srt/mem_cache/multi_ended_allocator.py +++ b/python/sglang/srt/mem_cache/multi_ended_allocator.py @@ -3685,8 +3685,9 @@ def swa_v2p_page_table(self) -> torch.Tensor: def alloc(self, need_size: int) -> Optional[torch.Tensor]: with record_function("UnifiedSWAAlloc.alloc"): - if not self.ensure_capacity(need_size, need_size): - return None + if need_size > self.available_size(): + if not _relieve_for_alloc(self, need_size): + return None # Snapshot the virtual PAGES full will consume, to bind them on swa too. num_pages = need_size // self.page_size fa = self.full_attn_allocator @@ -3722,8 +3723,9 @@ def alloc_extend( prefix_lens=prefix_lens_cpu, ) need_tokens = num_new_pages * self.page_size - if not self.ensure_capacity(need_tokens, need_tokens): - return None + if need_tokens > self.available_size(): + if not _relieve_for_alloc(self, need_tokens): + return None # Snapshot the virtual PAGES the kernel will consume; clone so swa keeps # its view after the slice is consumed. @@ -3760,8 +3762,9 @@ def alloc_decode( seq_lens=seq_lens_cpu, page_size=self.page_size, decode=True ) need_tokens = num_new_pages * self.page_size - if not self.ensure_capacity(need_tokens, need_tokens): - return None + if need_tokens > self.available_size(): + if not _relieve_for_alloc(self, need_tokens): + return None fa = self.full_attn_allocator new_virtual_pages = fa.free_virtual_ids[:num_new_pages].clone() diff --git a/python/sglang/srt/mem_cache/unified_memory_pool.py b/python/sglang/srt/mem_cache/unified_memory_pool.py index 58a09de84f93..aa96cbdd216b 100644 --- a/python/sglang/srt/mem_cache/unified_memory_pool.py +++ b/python/sglang/srt/mem_cache/unified_memory_pool.py @@ -1667,6 +1667,7 @@ def init_unified_swa_pools( full_attention_layer_ids: List[int], full_max_total_num_tokens: Optional[int] = None, swa_max_total_num_tokens: Optional[int] = None, + total_bytes: Optional[int] = None, enable_memory_saver: bool, need_sort: bool, forward_stream: Optional[torch.cuda.Stream] = None, @@ -1710,15 +1711,16 @@ def init_unified_swa_pools( store_dtype=store_dtype, grow_direction="up", ) - legacy_allocator_capacities = {} - if unified_total_bytes is not None: - # PROFILED byte budget, sized from directly: the re-sum's floor losses - # stay out of the buffer, and the token counts remain boot labels. + if total_bytes is not None and unified_total_bytes is not None: + raise ValueError("total_bytes and unified_total_bytes are mutually exclusive") + if total_bytes is None: total_bytes = unified_total_bytes - else: + + legacy_allocator_capacities = {} + if total_bytes is None: if full_max_total_num_tokens is None or swa_max_total_num_tokens is None: raise ValueError( - "unified_total_bytes or both legacy full/SWA capacities must be provided" + "total_bytes or both legacy full/SWA capacities must be provided" ) total_bytes = ( full_max_total_num_tokens * full_spec.entry_bytes() diff --git a/test/registered/unit/server_args/test_server_args.py b/test/registered/unit/server_args/test_server_args.py index b81bdd2dfd24..9620dfa45af7 100644 --- a/test/registered/unit/server_args/test_server_args.py +++ b/test/registered/unit/server_args/test_server_args.py @@ -29,6 +29,7 @@ ) from sglang.srt.arg_groups.kv_cache_hook import ( handle_cache_compatibility, + handle_unified_memory_pool, validate_prefill_only_disable_kv_cache_args, ) from sglang.srt.arg_groups.mamba_hook import handle_mamba_backend @@ -971,6 +972,22 @@ def test_hisparse_rejects_unsupported_kv_cache_dtype(self): validate_hisparse_kv_cache_dtype(server_args) +class TestUnifiedMemoryPoolCompatibility(CustomTestCase): + def test_rejects_unsupported_scheduler_configurations(self): + cases = ( + ({"pp_size": 2, "enable_mixed_chunk": False}, "pipeline parallelism"), + ({"pp_size": 1, "enable_mixed_chunk": True}, "enable-mixed-chunk"), + ) + for overrides, error in cases: + with self.subTest(**overrides): + args = object.__new__(ServerArgs) + args.enable_unified_memory = True + args.pp_size = overrides["pp_size"] + args.enable_mixed_chunk = overrides["enable_mixed_chunk"] + with self.assertRaisesRegex(AssertionError, error): + handle_unified_memory_pool(args) + + class TestFa4PageSizeAutoForce(CustomTestCase): """FA4 requires page_size 128 for non-MLA models on SM100. The auto-force must trigger for `--attention-backend fa4` (combined) too, not only for the From c54b97d9703cdbc9d7c3b4c7bdce4fdae91ea8fe Mon Sep 17 00:00:00 2001 From: yhzhuang Date: Thu, 27 Aug 2026 17:11:40 -0700 Subject: [PATCH 06/55] Remove unified pool gate-only tests --- .../unit/server_args/test_server_args.py | 17 ----------------- 1 file changed, 17 deletions(-) diff --git a/test/registered/unit/server_args/test_server_args.py b/test/registered/unit/server_args/test_server_args.py index 9620dfa45af7..b81bdd2dfd24 100644 --- a/test/registered/unit/server_args/test_server_args.py +++ b/test/registered/unit/server_args/test_server_args.py @@ -29,7 +29,6 @@ ) from sglang.srt.arg_groups.kv_cache_hook import ( handle_cache_compatibility, - handle_unified_memory_pool, validate_prefill_only_disable_kv_cache_args, ) from sglang.srt.arg_groups.mamba_hook import handle_mamba_backend @@ -972,22 +971,6 @@ def test_hisparse_rejects_unsupported_kv_cache_dtype(self): validate_hisparse_kv_cache_dtype(server_args) -class TestUnifiedMemoryPoolCompatibility(CustomTestCase): - def test_rejects_unsupported_scheduler_configurations(self): - cases = ( - ({"pp_size": 2, "enable_mixed_chunk": False}, "pipeline parallelism"), - ({"pp_size": 1, "enable_mixed_chunk": True}, "enable-mixed-chunk"), - ) - for overrides, error in cases: - with self.subTest(**overrides): - args = object.__new__(ServerArgs) - args.enable_unified_memory = True - args.pp_size = overrides["pp_size"] - args.enable_mixed_chunk = overrides["enable_mixed_chunk"] - with self.assertRaisesRegex(AssertionError, error): - handle_unified_memory_pool(args) - - class TestFa4PageSizeAutoForce(CustomTestCase): """FA4 requires page_size 128 for non-MLA models on SM100. The auto-force must trigger for `--attention-backend fa4` (combined) too, not only for the From 8892c52615df94787b6d24e37871efce4a8f3b65 Mon Sep 17 00:00:00 2001 From: yhzhuang Date: Thu, 27 Aug 2026 19:00:58 -0700 Subject: [PATCH 07/55] Limit shared-capacity changes to enabled paths Co-authored-by: Yonghao Zhuang --- python/sglang/srt/disaggregation/decode.py | 202 ++++-------------- python/sglang/srt/mem_cache/common.py | 26 +-- .../srt/model_executor/pool_configurator.py | 26 +-- .../test_decode_queue_cleanup.py | 35 +-- .../mem_cache/test_decode_radix_lock_ref.py | 14 +- 5 files changed, 65 insertions(+), 238 deletions(-) diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 348dc103df79..917a6700f626 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -78,7 +78,6 @@ EvictParams, ) from sglang.srt.mem_cache.common import ( - evict_from_tree_cache, kv_to_page_indices, page_align_floor, release_kv_cache, @@ -91,9 +90,6 @@ KVCache, ReqToTokenPool, ) -from sglang.srt.mem_cache.multi_ended_allocator import ( - UnifiedSWATokenToKVPoolAllocator, -) from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool from sglang.srt.observability.req_time_stats import ( set_schedule_time_batch, @@ -413,74 +409,6 @@ def _uses_swa_tail_prealloc(self) -> bool: and hasattr(self.token_to_kv_pool_allocator, "alloc_extend_swa_tail") ) - def _unified_swa_reservation_fits( - self, - full_tokens: int, - swa_tokens: int, - *, - full_allocatable_tokens: Optional[int] = None, - swa_allocatable_tokens: Optional[int] = None, - empty_pool: bool = False, - ) -> bool: - allocator = self.token_to_kv_pool_allocator - if not ( - isinstance(allocator, UnifiedSWATokenToKVPoolAllocator) - and allocator.supports_asymmetric_reservation - ): - return True - - full_evictable = swa_evictable = 0 - if not empty_pool: - full_evictable = self._radix_full_evictable() - swa_evictable = self.tree_cache.swa_evictable_size() - if full_allocatable_tokens is not None: - full_tokens += ( - allocator.full_available_size() - + full_evictable - - full_allocatable_tokens - ) - if swa_allocatable_tokens is not None: - swa_tokens += ( - allocator.swa_available_size() - + swa_evictable - - swa_allocatable_tokens - ) - - return allocator.can_reserve( - full_tokens, - swa_tokens, - full_evictable_tokens=full_evictable, - swa_evictable_tokens=swa_evictable, - empty_pool=empty_pool, - ) - - def _prealloc_reservation_fits( - self, - full_tokens: int, - swa_tokens: int, - *, - full_allocatable_tokens: int, - swa_allocatable_tokens: Optional[int], - uses_swa_tail_prealloc: bool, - ) -> bool: - if not uses_swa_tail_prealloc: - return full_tokens <= full_allocatable_tokens - assert swa_allocatable_tokens is not None - if isinstance( - self.token_to_kv_pool_allocator, - UnifiedSWATokenToKVPoolAllocator, - ) and self.token_to_kv_pool_allocator.supports_asymmetric_reservation: - return self._unified_swa_reservation_fits( - full_tokens, - swa_tokens, - full_allocatable_tokens=full_allocatable_tokens, - swa_allocatable_tokens=swa_allocatable_tokens, - ) - return ( - full_tokens <= full_allocatable_tokens - and swa_tokens <= swa_allocatable_tokens - ) - def _release_matched_prefix_lock(self, req: Req) -> None: params = DecLockRefParams(swa_uuid_for_lock=req.swa_uuid_for_lock) if req.swa_prefix_lock_released: @@ -490,35 +418,16 @@ def _release_matched_prefix_lock(self, req: Req) -> None: self.tree_cache.dec_lock_ref(req.last_node, params) def _reclaim_swa_tail_capacity( - self, swa_tail_len: int, req_id: str, *, full_len: int = 0 + self, swa_tail_len: int, req_id: str ) -> Optional[str]: - allocator = self.token_to_kv_pool_allocator - page_size = allocator.page_size - if ( - isinstance(allocator, UnifiedSWATokenToKVPoolAllocator) - and allocator.supports_asymmetric_reservation - ): - full_required = ceil_align(full_len, page_size) - swa_required = ceil_align(swa_tail_len, page_size) - evict_from_tree_cache( - self.tree_cache, - full_required, - swa_num_tokens=swa_required, - ) - if allocator.ensure_capacity(full_required, swa_required): - return None - return ( - "Unified FULL/SWA byte reclamation insufficient: " - f"needed=({full_required}, {swa_required}), req={req_id}" - ) - + page_size = self.token_to_kv_pool_allocator.page_size required = ceil_align(swa_tail_len, page_size) - available = allocator.swa_available_size() + available = self.token_to_kv_pool_allocator.swa_available_size() if available < required: self.tree_cache.evict_for_alloc( EvictParams(swa_num_tokens=required - available) ) - available = allocator.swa_available_size() + available = self.token_to_kv_pool_allocator.swa_available_size() if available < required: return ( f"SWA eviction insufficient: needed={required}, " @@ -849,22 +758,6 @@ def _check_if_req_exceed_kv_capacity(self, req: Req) -> bool: else: capacity = self.max_total_num_tokens input_len = self._rebootstrap_prefill_len(req) - if isinstance( - self.token_to_kv_pool_allocator, UnifiedSWATokenToKVPoolAllocator - ) and self.token_to_kv_pool_allocator.supports_asymmetric_reservation: - full_required, swa_required = self._prealloc_required_tokens(req) - if not self._unified_swa_reservation_fits( - full_required, swa_required, empty_pool=True - ): - message = ( - f"Request {req.rid} exceeds the unified FULL/SWA KV byte " - f"budget: full={full_required}, swa={swa_required}" - ) - logger.error(message) - prepare_abort(req, message, status_code=HTTPStatus.BAD_REQUEST) - self.scheduler.output_streamer.stream_output([req], req.return_logprob) - return True - return False if input_len > capacity: message = f"Request {req.rid} exceeds the maximum number of tokens: {input_len} > {capacity}" logger.error(message) @@ -933,28 +826,10 @@ def resume_retracted_reqs( break full_required, swa_required = self._prealloc_required_tokens(req) - if not self._prealloc_reservation_fits( - full_required, - swa_required, - full_allocatable_tokens=full_allocatable_tokens, - swa_allocatable_tokens=( - swa_allocatable_tokens if uses_swa_tail_prealloc else None - ), - uses_swa_tail_prealloc=uses_swa_tail_prealloc, - ): + if full_required > full_allocatable_tokens: + break + if uses_swa_tail_prealloc and swa_required > swa_allocatable_tokens: break - - if uses_swa_tail_prealloc and isinstance( - self.token_to_kv_pool_allocator, - UnifiedSWATokenToKVPoolAllocator, - ) and self.token_to_kv_pool_allocator.supports_asymmetric_reservation: - full_len, swa_len = self._prealloc_kv_lens(req) - if ( - self._reclaim_swa_tail_capacity(swa_len, req.rid, full_len=full_len) - is not None - ): - break - resumed_reqs.append(req) indices_to_remove.add(i) req.is_retracted = False @@ -1393,18 +1268,27 @@ def pop_preallocated( required_alloc_tokens + self.num_reserved_decode_tokens ) - full_required_for_admission = max( - required_tokens_for_request, - origin_input_len - - prefix_len - + min( - decode_req.req.sampling_params.max_new_tokens, - CLIP_MAX_NEW_TOKEN, + if ( + max( + required_tokens_for_request, + origin_input_len + - prefix_len + + min( + decode_req.req.sampling_params.max_new_tokens, + CLIP_MAX_NEW_TOKEN, + ) + - retractable_tokens, ) - - retractable_tokens, - ) - swa_required_for_admission = 0 - swa_len = 0 + > full_allocatable_tokens + ): + if prefix_match is not None and prefix_match.l1_prefix_len > 0: + self._release_matched_prefix_lock(decode_req.req) + break + if required_tokens_for_request > full_allocatable_tokens: + if prefix_match is not None and prefix_match.l1_prefix_len > 0: + self._release_matched_prefix_lock(decode_req.req) + break + if uses_swa_tail_prealloc: _, swa_required = self._prealloc_required_tokens(decode_req.req) _, swa_len = self._prealloc_kv_lens(decode_req.req) @@ -1412,29 +1296,19 @@ def pop_preallocated( decode_req.req.sampling_params.max_new_tokens, CLIP_MAX_NEW_TOKEN, ) - swa_required_for_admission = max( - swa_required, - swa_len + max_new_tokens - retractable_swa_tokens, - ) - - if not self._prealloc_reservation_fits( - full_required_for_admission, - swa_required_for_admission, - full_allocatable_tokens=full_allocatable_tokens, - swa_allocatable_tokens=( - swa_allocatable_tokens if uses_swa_tail_prealloc else None - ), - uses_swa_tail_prealloc=uses_swa_tail_prealloc, - ): - if prefix_match is not None and prefix_match.l1_prefix_len > 0: - self._release_matched_prefix_lock(decode_req.req) - break + if ( + max( + swa_required, + swa_len + max_new_tokens - retractable_swa_tokens, + ) + > swa_allocatable_tokens + ): + if prefix_match is not None and prefix_match.l1_prefix_len > 0: + self._release_matched_prefix_lock(decode_req.req) + break - if uses_swa_tail_prealloc: reclaim_error = self._reclaim_swa_tail_capacity( - swa_len, - decode_req.req.rid, - full_len=required_alloc_tokens, + swa_len, decode_req.req.rid ) if reclaim_error is not None: if prefix_match is not None and prefix_match.l1_prefix_len > 0: diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index 1d0fcc231075..ec03782fdd1d 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/python/sglang/srt/mem_cache/common.py @@ -165,11 +165,7 @@ def maybe_cache_unfinished_req(req: Req, tree_cache: BasePrefixCache, **kwargs): tree_cache.cache_unfinished_req(req, **kwargs) -def evict_from_tree_cache( - tree_cache: BasePrefixCache | None, - num_tokens: int, - swa_num_tokens: Optional[int] = None, -): +def evict_from_tree_cache(tree_cache: BasePrefixCache | None, num_tokens: int): if tree_cache is None: return @@ -186,8 +182,7 @@ def evict_from_tree_cache( isinstance(allocator, UnifiedSWATokenToKVPoolAllocator) and allocator.supports_asymmetric_reservation ): - required_swa = num_tokens if swa_num_tokens is None else swa_num_tokens - if allocator.can_reserve(num_tokens, required_swa): + if allocator.can_reserve(num_tokens, num_tokens): return page_size = allocator.page_size @@ -195,7 +190,7 @@ def evict_from_tree_cache( swa_evictable = tree_cache.swa_evictable_size() // page_size * page_size if not allocator.can_reserve( num_tokens, - required_swa, + num_tokens, full_evictable_tokens=full_evictable, swa_evictable_tokens=swa_evictable, ): @@ -213,14 +208,14 @@ def minimum_reclaim(max_tokens: int, fits) -> int: if allocator.can_reserve( num_tokens, - required_swa, + num_tokens, swa_evictable_tokens=swa_evictable, ): full_reclaim = 0 swa_reclaim = minimum_reclaim( swa_evictable, lambda value: allocator.can_reserve( - num_tokens, required_swa, swa_evictable_tokens=value + num_tokens, num_tokens, swa_evictable_tokens=value ), ) else: @@ -229,7 +224,7 @@ def minimum_reclaim(max_tokens: int, fits) -> int: full_evictable, lambda value: allocator.can_reserve( num_tokens, - required_swa, + num_tokens, full_evictable_tokens=value, swa_evictable_tokens=swa_reclaim, ), @@ -240,23 +235,22 @@ def minimum_reclaim(max_tokens: int, fits) -> int: swa_num_tokens=swa_reclaim, ) ) - if not allocator.can_reserve(num_tokens, required_swa): + if not allocator.can_reserve(num_tokens, num_tokens): logger.warning( "Unified SWA eviction did not satisfy the planned reservation: " - f"requested=({num_tokens}, {required_swa}), " + f"requested=({num_tokens}, {num_tokens}), " f"planned=({full_reclaim}, {swa_reclaim}), " f"evicted=({evicted.num_tokens_evicted}, " f"{evicted.swa_num_tokens_evicted})" ) elif isinstance(allocator, SWATokenToKVPoolAllocator): # Hybrid allocator - required_swa = num_tokens if swa_num_tokens is None else swa_num_tokens full_available_size = allocator.full_available_size() swa_available_size = allocator.swa_available_size() - if full_available_size < num_tokens or swa_available_size < required_swa: + 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, required_swa - swa_available_size) + swa_num_tokens = max(0, num_tokens - swa_available_size) tree_cache.evict_for_alloc( EvictParams(num_tokens=full_num_tokens, swa_num_tokens=swa_num_tokens) ) diff --git a/python/sglang/srt/model_executor/pool_configurator.py b/python/sglang/srt/model_executor/pool_configurator.py index e9918338f260..70f7f33c466b 100644 --- a/python/sglang/srt/model_executor/pool_configurator.py +++ b/python/sglang/srt/model_executor/pool_configurator.py @@ -460,14 +460,6 @@ def __init__(self, kvc: KVCacheConfigurator): self._full_layers_num = len(model_config.full_attention_layer_ids) self._swa_layers_num = len(model_config.swa_attention_layer_ids) - self._target_full_layers_num = sum( - kvc.layer_info.start_layer <= layer_id < kvc.layer_info.end_layer - for layer_id in model_config.full_attention_layer_ids - ) - self._target_swa_layers_num = sum( - kvc.layer_info.start_layer <= layer_id < kvc.layer_info.end_layer - for layer_id in model_config.swa_attention_layer_ids - ) assert ( self._swa_layers_num > 0 ), "Hybrid SWA model must have at least one SWA layer" @@ -601,10 +593,10 @@ def _draft_pool_bytes_per_token(self) -> int: + self._draft_cell_size ) - def _target_unified_bytes(self, full_tokens: int, swa_tokens: int) -> int: + def _unified_pool_bytes(self, full_tokens: int, swa_tokens: int) -> int: return ( - full_tokens * self._full_per_token * self._target_full_layers_num - + swa_tokens * self._swa_per_token * self._target_swa_layers_num + full_tokens * self._full_per_token * self._full_layers_num + + swa_tokens * self._swa_per_token * self._swa_layers_num ) def _max_unified_full_tokens( @@ -615,9 +607,7 @@ def _max_unified_full_tokens( ) -> int: """Find the largest page-aligned full capacity whose allocations fit.""" draft_bytes_per_token = self._draft_pool_bytes_per_token() - target_full_bytes_per_token = ( - self._full_per_token * self._target_full_layers_num - ) + target_full_bytes_per_token = self._full_per_token * self._full_layers_num assert target_full_bytes_per_token > 0 def allocation_bytes(full_pages: int) -> int: @@ -629,7 +619,7 @@ def allocation_bytes(full_pages: int) -> int: // page_size * page_size ) - target_bytes = self._target_unified_bytes(full_tokens, swa_tokens) + target_bytes = self._unified_pool_bytes(full_tokens, swa_tokens) virtual_span = max(target_bytes // target_full_bytes_per_token - 1, 0) draft_tokens = ceil_align(virtual_span, page_size) + page_size return target_bytes + draft_tokens * draft_bytes_per_token @@ -688,7 +678,7 @@ def align_page_size(x: int) -> int: full_max_total_num_tokens=full_tokens, swa_max_total_num_tokens=swa_tokens, unified_memory_pool_bytes=( - self._target_unified_bytes(full_tokens, swa_tokens) + self._unified_pool_bytes(full_tokens, swa_tokens) if self._enable_unified_memory else None ), @@ -814,7 +804,7 @@ def calculate_pool_sizes( full_max_total_num_tokens=full_tokens, swa_max_total_num_tokens=swa_tokens, unified_memory_pool_bytes=( - self._target_unified_bytes(full_tokens, swa_tokens) + self._unified_pool_bytes(full_tokens, swa_tokens) if self._enable_unified_memory else None ), @@ -831,7 +821,7 @@ def calculate_pool_sizes_from_max_tokens( full_max_total_num_tokens=full_tokens, swa_max_total_num_tokens=min(swa_tokens, max_total_num_tokens), unified_memory_pool_bytes=( - self._target_unified_bytes( + self._unified_pool_bytes( full_tokens, min(swa_tokens, max_total_num_tokens) ) if self._enable_unified_memory diff --git a/test/registered/unit/disaggregation/test_decode_queue_cleanup.py b/test/registered/unit/disaggregation/test_decode_queue_cleanup.py index b037216213cb..7b5b9f8a06a0 100644 --- a/test/registered/unit/disaggregation/test_decode_queue_cleanup.py +++ b/test/registered/unit/disaggregation/test_decode_queue_cleanup.py @@ -13,9 +13,6 @@ from sglang.srt.distributed.parallel_state_wrapper import ParallelState from sglang.srt.managers.schedule_batch import FINISH_ABORT from sglang.srt.managers.scheduler import Scheduler -from sglang.srt.mem_cache.multi_ended_allocator import ( - UnifiedSWATokenToKVPoolAllocator, -) from sglang.srt.runtime_context import get_context from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -47,8 +44,7 @@ def test_paged_swa_retraction_resume_uses_physical_page_budget(self): page_size = 128 fill_len = 574 physical_tokens_per_req = 5 * page_size - full_available = 18 * page_size - swa_available = 4 * page_size + physical_available = 18 * page_size reqs = [ SimpleNamespace( @@ -66,10 +62,7 @@ def test_paged_swa_retraction_resume_uses_physical_page_budget(self): queue.retracted_queue = reqs.copy() queue.num_reserved_decode_tokens = 0 queue.req_to_token_pool = SimpleNamespace(available_size=lambda: len(reqs)) - queue.token_to_kv_pool_allocator = MagicMock( - spec=UnifiedSWATokenToKVPoolAllocator - ) - queue.token_to_kv_pool_allocator.page_size = page_size + queue.token_to_kv_pool_allocator = SimpleNamespace(page_size=page_size) queue.tree_cache = MagicMock() queue.scheduler = SimpleNamespace( sliding_window_size=2047, @@ -77,21 +70,16 @@ def test_paged_swa_retraction_resume_uses_physical_page_budget(self): ) queue._uses_swa_tail_prealloc = MagicMock(return_value=True) queue._swa_aware_allocatable_token_budgets = MagicMock( - return_value=(full_available, swa_available) + return_value=(physical_available, physical_available) ) queue._swa_tail_allocatable_token_budget = MagicMock( - side_effect=lambda **_: swa_available - ) - queue._unified_swa_reservation_fits = MagicMock( - side_effect=lambda full_tokens, _swa_tokens, *, full_allocatable_tokens, **_: full_tokens - <= full_allocatable_tokens + side_effect=lambda **_: physical_available ) - queue._reclaim_swa_tail_capacity = MagicMock(return_value=None) def pre_alloc(_req): - nonlocal full_available - self.assertGreaterEqual(full_available, physical_tokens_per_req) - full_available -= physical_tokens_per_req + nonlocal physical_available + self.assertGreaterEqual(physical_available, physical_tokens_per_req) + physical_available -= physical_tokens_per_req queue._pre_alloc = MagicMock(side_effect=pre_alloc) @@ -99,14 +87,8 @@ def pre_alloc(_req): self.assertEqual(resumed, reqs[:3]) self.assertEqual(queue.retracted_queue, reqs[3:]) - self.assertEqual(full_available, 3 * page_size) + self.assertEqual(physical_available, 3 * page_size) self.assertEqual(queue._pre_alloc.call_count, 3) - self.assertTrue( - all( - call.args[1] > call.kwargs["swa_allocatable_tokens"] - for call in queue._unified_swa_reservation_fits.call_args_list - ) - ) def test_prealloc_abort_clears_receiver_before_removing_request(self): receiver = FakeReceiver() @@ -230,7 +212,6 @@ def test_swa_reclaim_failure_rejects_only_request(self): ) queue._hicache_pending_restore_tokens = MagicMock(return_value=0) queue._pre_alloc = MagicMock() - queue.token_to_kv_pool_allocator = MagicMock() queue.req_to_token_pool = MagicMock() queue.req_to_token_pool.available_size.return_value = 1 queue.req_to_metadata_buffer_idx_allocator = MagicMock() diff --git a/test/registered/unit/mem_cache/test_decode_radix_lock_ref.py b/test/registered/unit/mem_cache/test_decode_radix_lock_ref.py index 3cba68499020..73528e65bf91 100644 --- a/test/registered/unit/mem_cache/test_decode_radix_lock_ref.py +++ b/test/registered/unit/mem_cache/test_decode_radix_lock_ref.py @@ -38,9 +38,6 @@ InsertParams, MatchPrefixParams, ) -from sglang.srt.mem_cache.multi_ended_allocator import ( - UnifiedSWATokenToKVPoolAllocator, -) from sglang.srt.mem_cache.radix_cache import RadixCache, RadixKey from sglang.srt.utils.common import Range @@ -154,15 +151,6 @@ def test_reclaim_swa_tail_capacity_page_rounds(self): self.assertEqual(params.num_tokens, 0) self.assertEqual(params.swa_num_tokens, 128) - unified_allocator = MagicMock(spec=UnifiedSWATokenToKVPoolAllocator) - unified_allocator.page_size = 64 - unified_allocator.ensure_capacity.return_value = True - queue.token_to_kv_pool_allocator = unified_allocator - queue.tree_cache = MagicMock() - - self.assertIsNone(queue._reclaim_swa_tail_capacity(129, "unified", full_len=65)) - unified_allocator.ensure_capacity.assert_called_once_with(128, 192) - def test_reclaim_swa_tail_capacity_fails_before_allocation(self): queue = DecodePreallocQueue.__new__(DecodePreallocQueue) queue.token_to_kv_pool_allocator = MagicMock(page_size=64) @@ -461,7 +449,7 @@ def test_pop_preallocated_rechecks_budget_after_lock(self): skip_swa=True, ) self.assertFalse(req.swa_prefix_lock_released) - queue._swa_tail_len.assert_called_with(8) + queue._swa_tail_len.assert_called_once_with(8) queue._allocatable_token_budgets.assert_called_once() def test_repeated_incremental_no_leak(self): From 48b2a24041775653364cb3c3c84c51f1fbe2357e Mon Sep 17 00:00:00 2001 From: yhzhuang Date: Fri, 28 Aug 2026 12:24:31 -0700 Subject: [PATCH 08/55] Fix unified SWA eviction over-reclaim Minimize both component reclaim quotas under the joint capacity predicate so a blocked compaction path cannot turn one allocation shortfall into an all-SWA eviction. Original prod_inference commit: 56c7082a92bc0c9024585387460e64b065652918 Co-authored-by: Yonghao Zhuang --- python/sglang/srt/mem_cache/common.py | 41 ++++++++----------- .../test_unified_radix_cache_unittest.py | 39 +++++++++++++++++- 2 files changed, 56 insertions(+), 24 deletions(-) diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index ec03782fdd1d..e4dac917811c 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/python/sglang/srt/mem_cache/common.py @@ -206,29 +206,24 @@ def minimum_reclaim(max_tokens: int, fits) -> int: lo = mid + 1 return lo * page_size - if allocator.can_reserve( - num_tokens, - num_tokens, - swa_evictable_tokens=swa_evictable, - ): - full_reclaim = 0 - swa_reclaim = minimum_reclaim( - swa_evictable, - lambda value: allocator.can_reserve( - num_tokens, num_tokens, swa_evictable_tokens=value - ), - ) - else: - swa_reclaim = swa_evictable - full_reclaim = minimum_reclaim( - full_evictable, - lambda value: allocator.can_reserve( - num_tokens, - num_tokens, - full_evictable_tokens=value, - swa_evictable_tokens=swa_reclaim, - ), - ) + full_reclaim = minimum_reclaim( + full_evictable, + lambda value: allocator.can_reserve( + num_tokens, + num_tokens, + full_evictable_tokens=value, + swa_evictable_tokens=swa_evictable, + ), + ) + swa_reclaim = minimum_reclaim( + swa_evictable, + lambda value: allocator.can_reserve( + num_tokens, + num_tokens, + full_evictable_tokens=full_reclaim, + swa_evictable_tokens=value, + ), + ) evicted = tree_cache.evict_for_alloc( EvictParams( num_tokens=full_reclaim, diff --git a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py index ae9f0f889ad6..fc6ece8ba757 100644 --- a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py +++ b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py @@ -38,7 +38,10 @@ zero_match_result, ) from sglang.srt.mem_cache.cache_init_params import CacheInitParams -from sglang.srt.mem_cache.common import available_and_evictable_str +from sglang.srt.mem_cache.common import ( + available_and_evictable_str, + evict_from_tree_cache, +) from sglang.srt.mem_cache.hicache_storage import ( PoolHitPolicy, PoolName, @@ -51,6 +54,9 @@ MHATokenToKVPool, ReqToTokenPool, ) +from sglang.srt.mem_cache.multi_ended_allocator import ( + UnifiedSWATokenToKVPoolAllocator, +) from sglang.srt.mem_cache.radix_cache import RadixKey from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool from sglang.srt.mem_cache.unified_cache.cache_action import ( @@ -121,6 +127,37 @@ def _session_radix_cache_test_values() -> tuple[bool, ...]: return False, True +class TestUnifiedSWAEvictionPlanning(CustomTestCase): + def test_reclaims_only_needed_pages_from_each_side(self): + allocator = object.__new__(UnifiedSWATokenToKVPoolAllocator) + allocator.page_size = 4 + tree_cache = mock.MagicMock() + tree_cache.is_chunk_cache.return_value = False + tree_cache.token_to_kv_pool_allocator = allocator + tree_cache.full_evictable_size.return_value = 16 + tree_cache.swa_evictable_size.return_value = 64 + + def can_reserve( + full_tokens, + swa_tokens, + *, + full_evictable_tokens=0, + swa_evictable_tokens=0, + **_, + ): + return tree_cache.evict.called or ( + full_evictable_tokens >= 8 and swa_evictable_tokens >= 8 + ) + + allocator.can_reserve = mock.Mock(side_effect=can_reserve) + + evict_from_tree_cache(tree_cache, 4) + + tree_cache.evict.assert_called_once_with( + EvictParams(num_tokens=8, swa_num_tokens=8) + ) + + @dataclass(frozen=True) class CacheConfig: # Tree From 139a9a1f6e7f8ecf1dad3588fe528ed3240b5c3a Mon Sep 17 00:00:00 2001 From: Yonghao Zhuang Date: Fri, 28 Aug 2026 13:52:16 -0700 Subject: [PATCH 09/55] Remove added unified eviction planner test --- .../test_unified_radix_cache_unittest.py | 39 +------------------ 1 file changed, 1 insertion(+), 38 deletions(-) diff --git a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py index fc6ece8ba757..ae9f0f889ad6 100644 --- a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py +++ b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py @@ -38,10 +38,7 @@ zero_match_result, ) from sglang.srt.mem_cache.cache_init_params import CacheInitParams -from sglang.srt.mem_cache.common import ( - available_and_evictable_str, - evict_from_tree_cache, -) +from sglang.srt.mem_cache.common import available_and_evictable_str from sglang.srt.mem_cache.hicache_storage import ( PoolHitPolicy, PoolName, @@ -54,9 +51,6 @@ MHATokenToKVPool, ReqToTokenPool, ) -from sglang.srt.mem_cache.multi_ended_allocator import ( - UnifiedSWATokenToKVPoolAllocator, -) from sglang.srt.mem_cache.radix_cache import RadixKey from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool from sglang.srt.mem_cache.unified_cache.cache_action import ( @@ -127,37 +121,6 @@ def _session_radix_cache_test_values() -> tuple[bool, ...]: return False, True -class TestUnifiedSWAEvictionPlanning(CustomTestCase): - def test_reclaims_only_needed_pages_from_each_side(self): - allocator = object.__new__(UnifiedSWATokenToKVPoolAllocator) - allocator.page_size = 4 - tree_cache = mock.MagicMock() - tree_cache.is_chunk_cache.return_value = False - tree_cache.token_to_kv_pool_allocator = allocator - tree_cache.full_evictable_size.return_value = 16 - tree_cache.swa_evictable_size.return_value = 64 - - def can_reserve( - full_tokens, - swa_tokens, - *, - full_evictable_tokens=0, - swa_evictable_tokens=0, - **_, - ): - return tree_cache.evict.called or ( - full_evictable_tokens >= 8 and swa_evictable_tokens >= 8 - ) - - allocator.can_reserve = mock.Mock(side_effect=can_reserve) - - evict_from_tree_cache(tree_cache, 4) - - tree_cache.evict.assert_called_once_with( - EvictParams(num_tokens=8, swa_num_tokens=8) - ) - - @dataclass(frozen=True) class CacheConfig: # Tree From 9c31d3d0d03162c970172b4295fe273da3f17b98 Mon Sep 17 00:00:00 2001 From: yhzhuang Date: Mon, 31 Aug 2026 17:03:50 -0700 Subject: [PATCH 10/55] Remove low-signal unified capacity tests Co-authored-by: Yonghao Zhuang --- .../unit/managers/test_prefill_adder.py | 52 ------------------- .../model_executor/test_pool_configurator.py | 8 --- 2 files changed, 60 deletions(-) diff --git a/test/registered/unit/managers/test_prefill_adder.py b/test/registered/unit/managers/test_prefill_adder.py index 8512a0a3cf90..6525c3ce9bf5 100644 --- a/test/registered/unit/managers/test_prefill_adder.py +++ b/test/registered/unit/managers/test_prefill_adder.py @@ -14,9 +14,6 @@ IncLockRefResult, ) from sglang.srt.runtime_context import get_context -from sglang.srt.mem_cache.multi_ended_allocator import ( - UnifiedSWATokenToKVPoolAllocator, -) from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler from sglang.srt.utils.common import Range from sglang.test.ci.ci_register import register_cpu_ci @@ -524,55 +521,6 @@ def test_add_chunked_req_hybrid_swa_defers_when_swa_below_page(self): req.set_extend_range.assert_not_called() self.assertEqual(len(adder.can_run_list), 0) - def test_add_one_req_unified_swa_uses_shared_capacity(self): - # The real warmup is shorter than one page. Feeding an already - # page-rounded input into the SWA budget would round it a second time - # after adding decode headroom and reject the first request forever. - page_size = 128 - allocator = MagicMock(spec=UnifiedSWATokenToKVPoolAllocator) - allocator.available_size.return_value = 0 - allocator.full_available_size.return_value = 0 - allocator.swa_available_size.return_value = 0 - allocator.size_swa = 1_000_000 - allocator.can_reserve.side_effect = ( - lambda full_tokens, swa_tokens, *, full_evictable_tokens=0, swa_evictable_tokens=0, **_: max( - full_tokens, swa_tokens - ) - <= full_evictable_tokens + swa_evictable_tokens - ) - self.mock_token_allocator = allocator - self.mock_tree_cache.swa_evictable_size.return_value = 3 * page_size - self.mock_tree_cache.sliding_window_size = 4096 - adder = self.create_adder(self.create_running_batch(), page_size=page_size) - - req = self.create_mock_req("unified", priority=0, max_new_tokens=10) - req.origin_input_ids = list(range(3)) - req.full_untruncated_fill_ids = list(range(3)) - req.swa_host_hit_length = 0 - req.last_node = MagicMock() - req.sampling_params.ignore_eos = False - req.set_extend_range = MagicMock( - side_effect=lambda start, end: setattr( - req, "extend_range", Range(start, end) - ) - ) - - result = adder.add_one_req( - req, has_chunked_req=False, truncation_align_size=None - ) - - self.assertEqual(result, AddReqResult.CONTINUE) - self.assertIn(req, adder.can_run_list) - full_tokens, swa_tokens = allocator.can_reserve.call_args_list[0].args - self.assertEqual(full_tokens, 3 + 10 + page_size) - self.assertEqual(swa_tokens, 3 + 10 + page_size) - self.assertTrue( - any( - call.kwargs["swa_evictable_tokens"] == 3 * page_size - for call in allocator.can_reserve.call_args_list - ) - ) - def test_swa_budget_for_req(self): # budget = max(alloc - window, 0) + min(extend + max_new, window) + page, # where alloc = min(extend, rem_chunk). The decode headroom is the SWA the diff --git a/test/registered/unit/model_executor/test_pool_configurator.py b/test/registered/unit/model_executor/test_pool_configurator.py index 915dcf822693..3e8cf29b283c 100644 --- a/test/registered/unit/model_executor/test_pool_configurator.py +++ b/test/registered/unit/model_executor/test_pool_configurator.py @@ -337,14 +337,6 @@ def test_memory_utilization(self): self.assertLessEqual(used, available) self.assertGreater(used, available * 0.99) - def test_unified_memory_records_exact_target_bytes(self): - available = 10_000_000 - mr, _, config = self._run(available, enable_unified_memory=True) - used = _actual_memory_used(mr, config) - self.assertEqual(config.unified_memory_pool_bytes, used) - self.assertLessEqual(used, available) - self.assertGreater(used, available * 0.99) - @patch( "sglang.srt.mem_cache.kv_cache_configurator.calculate_mla_kv_cache_dim", return_value=576, From 159935705473ce7ce8698801c00f5667518e0f16 Mon Sep 17 00:00:00 2001 From: yhzhuang Date: Tue, 1 Sep 2026 15:22:19 -0700 Subject: [PATCH 11/55] Adapt shared capacity to current main --- python/sglang/srt/managers/schedule_policy.py | 9 ++++++--- python/sglang/srt/mem_cache/kv_cache_configurator.py | 10 +++++++++- 2 files changed, 15 insertions(+), 4 deletions(-) diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index 13ca997369dd..3afcdd11b354 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -575,9 +575,12 @@ def __init__( self.token_to_kv_pool_allocator, PureSWATokenToKVPoolAllocator ) self.is_hybrid_ssm_cache = self.tree_cache.supports_mamba() - self.is_unified_swa = isinstance( - self.token_to_kv_pool_allocator, UnifiedSWATokenToKVPoolAllocator - ) and self.token_to_kv_pool_allocator.supports_asymmetric_reservation + self.is_unified_swa = ( + isinstance( + self.token_to_kv_pool_allocator, UnifiedSWATokenToKVPoolAllocator + ) + and self.token_to_kv_pool_allocator.supports_asymmetric_reservation + ) self.rem_swa_token_offset = ( num_mixed_decode_tokens if self.is_unified_swa else 0 diff --git a/python/sglang/srt/mem_cache/kv_cache_configurator.py b/python/sglang/srt/mem_cache/kv_cache_configurator.py index 515919c2ba92..558b75f41d8c 100644 --- a/python/sglang/srt/mem_cache/kv_cache_configurator.py +++ b/python/sglang/srt/mem_cache/kv_cache_configurator.py @@ -462,7 +462,10 @@ def _init_pools( ) bundle = self._init_unified_swa_pools( max_num_reqs=sizes.max_running_requests, + full_max_total_num_tokens=sizes.full_max_total_num_tokens, + swa_max_total_num_tokens=sizes.swa_max_total_num_tokens, unified_memory_pool_bytes=sizes.unified_memory_pool_bytes, + unified_total_bytes=sizes.unified_total_bytes, ) else: # Fail loud, not silently fall through to the normal pools (which would @@ -810,6 +813,7 @@ def _init_unified_swa_pools( full_max_total_num_tokens: Optional[int] = None, swa_max_total_num_tokens: Optional[int] = None, unified_memory_pool_bytes: Optional[int] = None, + unified_total_bytes: Optional[int] = None, ) -> UnifiedPoolBundle: """Build the unified-pool stack for a hybrid-SWA model (Triton): one byte buffer split between the full-attention and SWA KV pools.""" @@ -869,6 +873,10 @@ def _init_unified_swa_pools( if self.layer_info.start_layer <= i < self.layer_info.end_layer ] + total_bytes = unified_memory_pool_bytes + if total_bytes is None and self.spec_algorithm.is_none(): + total_bytes = unified_total_bytes + bundle = init_unified_swa_pools( device=self.device, kv_cache_dtype=self.kv_cache_dtype, @@ -885,7 +893,7 @@ def _init_unified_swa_pools( full_attention_layer_ids=full_attention_layer_ids, full_max_total_num_tokens=full_max_total_num_tokens, swa_max_total_num_tokens=swa_max_total_num_tokens, - total_bytes=unified_memory_pool_bytes, + total_bytes=total_bytes, enable_memory_saver=get_exec().features.enable_memory_saver, need_sort=get_disagg().disaggregation_mode in ("decode", "prefill"), # Overlap mode: same wait_stream(forward_stream) rationale as From c2296bec6231e1cb9bee8633f0976a03791ca39f Mon Sep 17 00:00:00 2001 From: yhzhuang Date: Tue, 1 Sep 2026 15:39:36 -0700 Subject: [PATCH 12/55] Address unified SWA admission feedback --- python/sglang/srt/arg_groups/kv_cache_hook.py | 11 ++++------- python/sglang/srt/managers/schedule_policy.py | 8 ++++++-- 2 files changed, 10 insertions(+), 9 deletions(-) diff --git a/python/sglang/srt/arg_groups/kv_cache_hook.py b/python/sglang/srt/arg_groups/kv_cache_hook.py index a25f4214389b..96c541d33b8f 100644 --- a/python/sglang/srt/arg_groups/kv_cache_hook.py +++ b/python/sglang/srt/arg_groups/kv_cache_hook.py @@ -198,14 +198,11 @@ def handle_unified_memory_pool(server_args: Any) -> None: cfg = resolving_view(server_args) if not cfg.enable_unified_memory: return - assert cfg.pp_size == 1, ( - "--enable-unified-memory does not support pipeline parallelism " - "(--pp-size > 1)." - ) - assert ( - not cfg.enable_mixed_chunk - ), "--enable-unified-memory does not support --enable-mixed-chunk." if cfg.disaggregation_mode != "null": + assert cfg.pp_size == 1, ( + "--enable-unified-memory with PD disaggregation does not support " + "pipeline parallelism (--pp-size > 1)." + ) # Constraints of the whole-envelope transfer; see # UnifiedMLATokenToKVPool.get_contiguous_buf_infos. assert cfg.disaggregation_transfer_backend == "mooncake", ( diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index 3afcdd11b354..4b4075506926 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -87,6 +87,7 @@ def estimate_swa_kv_tokens( page_size: int, allocation_limit: Optional[int] = None, host_hit_length: int = 0, + include_allocated_tail: bool = True, ) -> int: """Peak SWA reservation for one prefill/decode request.""" if sliding_window_size is None or sliding_window_size <= 0: @@ -97,8 +98,11 @@ def estimate_swa_kv_tokens( if allocation_limit is None else min(extend_input_len, allocation_limit) ) + allocated_tail = ( + max(allocated - sliding_window_size, 0) if include_allocated_tail else 0 + ) reserved = ( - max(allocated - sliding_window_size, 0) + allocated_tail + min(extend_input_len + max_new_tokens, sliding_window_size) + page_size ) @@ -767,8 +771,8 @@ def _swa_reserved_tokens( max_new_tokens, sliding_window_size=self.tree_cache.sliding_window_size, page_size=self.page_size, - allocation_limit=0, host_hit_length=swa_host_hit_length, + include_allocated_tail=False, ) def _swa_new_tokens(self, req: Req) -> int: From 05cfb811c22682b06ff3654feb4ea9148a8bc993 Mon Sep 17 00:00:00 2001 From: yhzhuang Date: Thu, 27 Aug 2026 13:41:25 -0700 Subject: [PATCH 13/55] Support unified memory page envelopes in PD transfers Co-authored-by: Yonghao Zhuang --- python/sglang/srt/arg_groups/kv_cache_hook.py | 9 ++ python/sglang/srt/disaggregation/base/conn.py | 28 ++++++ python/sglang/srt/disaggregation/decode.py | 39 ++++++-- python/sglang/srt/disaggregation/fake/conn.py | 19 ++++ python/sglang/srt/disaggregation/prefill.py | 28 ++++-- .../sglang/srt/mem_cache/kv_cache_builder.py | 14 ++- .../srt/mem_cache/kv_cache_configurator.py | 7 -- .../srt/mem_cache/multi_ended_allocator.py | 92 +++++++++++++++++-- .../srt/mem_cache/unified_memory_pool.py | 40 ++++++-- python/sglang/srt/server_args.py | 4 +- 10 files changed, 230 insertions(+), 50 deletions(-) diff --git a/python/sglang/srt/arg_groups/kv_cache_hook.py b/python/sglang/srt/arg_groups/kv_cache_hook.py index 96c541d33b8f..dbed1eaa1e28 100644 --- a/python/sglang/srt/arg_groups/kv_cache_hook.py +++ b/python/sglang/srt/arg_groups/kv_cache_hook.py @@ -220,6 +220,15 @@ def handle_unified_memory_pool(server_args: Any) -> None: "ships host/C4 rows straight from the allocator, bypassing the " "virtual->physical translation the unified pool needs." ) + assert cfg.disaggregation_decode_retraction_backup != "host_pool", ( + "--enable-unified-memory with PD disaggregation does not support " + "--disaggregation-decode-retraction-backup=host_pool; use " + "cpu_tensor (the automatic default for unified pools)." + ) + assert not cfg.disaggregation_decode_enable_offload_kvcache, ( + "--enable-unified-memory with PD disaggregation does not yet support " + "--disaggregation-decode-enable-offload-kvcache." + ) assert cfg.speculative_algorithm in (None, "DSPARK"), ( "--enable-unified-memory only supports --speculative-algorithm " "DSPARK (chain draft); other speculative algorithms are not yet " diff --git a/python/sglang/srt/disaggregation/base/conn.py b/python/sglang/srt/disaggregation/base/conn.py index 1dedb16ebed6..00e1d84ee775 100644 --- a/python/sglang/srt/disaggregation/base/conn.py +++ b/python/sglang/srt/disaggregation/base/conn.py @@ -45,6 +45,9 @@ class KVArgs: kv_data_ptrs: List[int] kv_data_lens: List[int] kv_item_lens: List[int] + # Target KV buffers come first; any remaining buffers belong to an + # independent draft pool and may require a different index vector. + num_target_kv_data_ptrs: int kv_layer_ids: List[int] kv_cache_dtype_str: str aux_data_ptrs: List[int] @@ -101,6 +104,8 @@ class KVPoll: class BaseKVManager(ABC): """Base class for managing transfer states""" + uses_separate_draft_kv_indices: bool = False + @abstractmethod def __init__( self, @@ -147,6 +152,17 @@ def send( """ ... + def send_with_draft_indices( + self, + kv_indices: npt.NDArray[np.int32], + draft_kv_indices: npt.NDArray[np.int32], + state_indices: Optional[List] = None, + num_kv_tokens: Optional[int] = None, + ): + raise NotImplementedError( + f"{type(self).__name__} does not support separate draft KV indices" + ) + def pop_decode_prefix_len(self) -> int: return 0 @@ -217,6 +233,18 @@ def send_metadata( """ ... + def send_metadata_with_draft_indices( + self, + kv_indices: npt.NDArray[np.int32], + draft_kv_indices: npt.NDArray[np.int32], + aux_index: Optional[int] = None, + state_indices: Optional[List] = None, + decode_prefix_len: Optional[int] = None, + ): + raise NotImplementedError( + f"{type(self).__name__} does not support separate draft KV indices" + ) + @abstractmethod def poll(self) -> KVPoll: """ diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 917a6700f626..2c08a8fad921 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -538,6 +538,7 @@ def _init_kv_manager(self) -> CommonKVManager: kv_item_lens += device_kv_item_lens[c4_layer_num:] kv_data_mem_kinds += ["VRAM"] * len(device_kv_data_ptrs[c4_layer_num:]) num_draft_entries = 0 + kv_args.num_target_kv_data_ptrs = len(kv_data_ptrs) if self.draft_token_to_kv_pool is not None: # We should also transfer draft model kv cache. The indices are # always shared with a target model. @@ -1358,6 +1359,9 @@ def pop_preallocated( page_size = self.token_to_kv_pool_allocator.page_size kv_transfer_page_size = page_size + raw_kv_indices = self.req_to_token_pool.req_to_token[ + decode_req.req.kv.req_pool_idx + ][total_prefix_len:origin_input_len] if self.scheduler.enable_hisparse: # Direct-to-host sends host/C4 rows; keep allocator.page_size # logical and use the compressed page size only for these indices. @@ -1369,12 +1373,9 @@ def pop_preallocated( kv_indices = dst_kv_indices[: origin_input_len - prefix_len] else: # Only send delta indices (beyond prefix) to prefill. - kv_indices = self.req_to_token_pool.req_to_token[ - decode_req.req.kv.req_pool_idx - ][total_prefix_len:origin_input_len] kv_indices = ( self.token_to_kv_pool_allocator.translate_kv_indices_for_transfer( - kv_indices + raw_kv_indices ) ) @@ -1477,6 +1478,15 @@ def _c128_state_payload(): page_indices = kv_to_page_indices(kv_indices, kv_transfer_page_size).astype( np.int32 ) + draft_page_indices = None + if self.kv_manager.uses_separate_draft_kv_indices: + if self.scheduler.enable_hisparse: + raise NotImplementedError( + "separate draft KV indices are not supported with HiSparse" + ) + draft_page_indices = kv_to_page_indices( + raw_kv_indices, page_size + ).astype(np.int32) device_page_indices = None if ( self.scheduler.enable_hisparse @@ -1513,12 +1523,21 @@ def _c128_state_payload(): self.transfer_queue.staging_handler.register_decode_req( decode_req.req.bootstrap_room, decode_req ) - decode_req.kv_receiver.send_metadata( - page_indices, - decode_req.metadata_buffer_index, - state_indices, - **metadata_kwargs, - ) + if draft_page_indices is not None: + decode_req.kv_receiver.send_metadata_with_draft_indices( + page_indices, + draft_page_indices, + decode_req.metadata_buffer_index, + state_indices, + **metadata_kwargs, + ) + else: + decode_req.kv_receiver.send_metadata( + page_indices, + decode_req.metadata_buffer_index, + state_indices, + **metadata_kwargs, + ) if decode_req.is_rebootstrap: self.kv_manager.submit_prefill_recompute( decode_req.kv_receiver, diff --git a/python/sglang/srt/disaggregation/fake/conn.py b/python/sglang/srt/disaggregation/fake/conn.py index 15a76c039aa4..53389c88a282 100644 --- a/python/sglang/srt/disaggregation/fake/conn.py +++ b/python/sglang/srt/disaggregation/fake/conn.py @@ -85,6 +85,15 @@ def send( f"FakeKVSender send with kv_indices: {kv_indices}, state_indices: {state_indices}" ) + def send_with_draft_indices( + self, + kv_indices: npt.NDArray[np.int32], + draft_kv_indices: npt.NDArray[np.int32], + state_indices: Optional[List] = None, + num_kv_tokens: Optional[int] = None, + ): + self.send(kv_indices, state_indices, num_kv_tokens) + def failure_exception(self): raise Exception("Fake KVSender Exception") @@ -133,6 +142,16 @@ def send_metadata( f"FakeKVReceiver send_metadata with kv_indices: {kv_indices}, aux_index: {aux_index}, state_indices: {state_indices}" ) + def send_metadata_with_draft_indices( + self, + kv_indices: npt.NDArray[np.int32], + draft_kv_indices: npt.NDArray[np.int32], + aux_index: Optional[int] = None, + state_indices: Optional[List] = None, + decode_prefix_len: Optional[int] = None, + ): + self.send_metadata(kv_indices, aux_index, state_indices, decode_prefix_len) + def failure_exception(self): raise Exception("Fake KVReceiver Exception") diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index 404a9d07f5fa..8ab0e77ec31a 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -234,6 +234,7 @@ def _init_kv_manager(self) -> CommonKVManager: kv_data_ptrs, kv_data_lens, kv_item_lens = ( self.token_to_kv_pool.get_contiguous_buf_infos() ) + kv_args.num_target_kv_data_ptrs = len(kv_data_ptrs) kv_args.prefill_end_layer = ( kv_args.prefill_start_layer + len(kv_data_ptrs) if layer_shard_enabled @@ -1318,16 +1319,19 @@ def _c128_state_payload(): else: segments = [(start_idx, end_idx)] + separate_draft_indices = ( + self.disagg_prefill_bootstrap_queue.kv_manager.uses_separate_draft_kv_indices + ) for seg_start, seg_end in segments: is_final_segment = seg_end == end_idx - kv_indices = self.req_to_token_pool.req_to_token[ + raw_kv_indices = self.req_to_token_pool.req_to_token[ req.kv.req_pool_idx, seg_start:seg_end ] # Unified memory: req_to_token holds VIRTUAL ids; the transfer needs # physical ones. Per segment, since each is its own gather. kv_indices = ( self.token_to_kv_pool_allocator.translate_kv_indices_for_transfer( - kv_indices + raw_kv_indices ) ) page_indices = kv_to_page_indices(kv_indices, page_size) @@ -1336,11 +1340,21 @@ def _c128_state_payload(): len(page_indices), segment_is_last ): continue - req.disagg_kv_sender.send( - page_indices, - state_indices if segment_is_last else None, - num_kv_tokens=seg_end - seg_start, - ) + send_state_indices = state_indices if segment_is_last else None + if separate_draft_indices: + draft_page_indices = kv_to_page_indices(raw_kv_indices, page_size) + req.disagg_kv_sender.send_with_draft_indices( + page_indices, + draft_page_indices, + send_state_indices, + num_kv_tokens=seg_end - seg_start, + ) + else: + req.disagg_kv_sender.send( + page_indices, + send_state_indices, + num_kv_tokens=seg_end - seg_start, + ) req.start_send_idx = end_idx # A last chunk needs no entry: every `last_chunk=True` call site has # already put the request on `disagg_prefill_inflight_queue`. diff --git a/python/sglang/srt/mem_cache/kv_cache_builder.py b/python/sglang/srt/mem_cache/kv_cache_builder.py index 18acefeb8311..b841dc312cec 100644 --- a/python/sglang/srt/mem_cache/kv_cache_builder.py +++ b/python/sglang/srt/mem_cache/kv_cache_builder.py @@ -150,13 +150,17 @@ def resolve_decode_retraction_backup(*, tp_worker: BaseTpWorker) -> str: if tp_worker.is_hybrid_swa else None ) + # Unified page-envelope host pools are enabled by a separate feature. + # Until then, keep decode retraction on the direct CPU-tensor backup. # Host-pool retraction transfers full and sliding-window components # only, so a model with recurrent state stays on cpu_tensor. - supports_host_pool = not uses_ssm_state( - tp_worker.model_runner.model_config - ) and ( - isinstance(kv_cache, MHATokenToKVPool) - or (isinstance(kv_cache, SWAKVPool) and full_tokens_per_layer > 0) + supports_host_pool = ( + not memory.enable_unified_memory + and not uses_ssm_state(tp_worker.model_runner.model_config) + and ( + isinstance(kv_cache, MHATokenToKVPool) + or (isinstance(kv_cache, SWAKVPool) and full_tokens_per_layer > 0) + ) ) schedule = get_schedule() priority_preemption = ( diff --git a/python/sglang/srt/mem_cache/kv_cache_configurator.py b/python/sglang/srt/mem_cache/kv_cache_configurator.py index 558b75f41d8c..9516c6f5cfbd 100644 --- a/python/sglang/srt/mem_cache/kv_cache_configurator.py +++ b/python/sglang/srt/mem_cache/kv_cache_configurator.py @@ -453,13 +453,6 @@ def _init_pools( unified_total_bytes=sizes.unified_total_bytes, ) elif self.is_hybrid_swa and not is_dsv4: - if pd_enabled: - raise ValueError( - "--enable-unified-memory with PD disaggregation does " - "not support hybrid-SWA models yet (no whole-envelope " - "transfer scheme for the SWA sub-pool). Drop " - "--enable-unified-memory or run without PD." - ) bundle = self._init_unified_swa_pools( max_num_reqs=sizes.max_running_requests, full_max_total_num_tokens=sizes.full_max_total_num_tokens, diff --git a/python/sglang/srt/mem_cache/multi_ended_allocator.py b/python/sglang/srt/mem_cache/multi_ended_allocator.py index 479babc6366e..816fbc64dbe8 100644 --- a/python/sglang/srt/mem_cache/multi_ended_allocator.py +++ b/python/sglang/srt/mem_cache/multi_ended_allocator.py @@ -246,7 +246,11 @@ def _relieve_for_alloc(short_pool, need_tokens: int) -> bool: class MultiEndedAllocator(BaseTokenToKVPoolAllocator): - """Allocator for one sub-pool over a `UnifiedKVPool`.""" + """Allocator for one sub-pool over a `UnifiedKVPool`. + + ``need_sort`` applies to transfer-facing physical ids, not virtual ids. + Physical free pages are sorted during compaction. + """ # Capacity-bearing state: any rebind bumps `_capacity_epoch`, invalidating # the epoch-keyed capacity memos across the whole chain (see @@ -1242,11 +1246,6 @@ def alloc_extend( if not _relieve_for_alloc(self, need_tokens): return None bs = len(prefix_lens) - if self.need_sort and extend_num_tokens // self.page_size + bs + 1 > len( - self.free_virtual_ids - ): - self.merge_and_sort_free() - # Snapshot the virtual pages the kernel will consume, to bind them to # physical pages afterward (else v2p stays -1 → CUDA OOB). if num_new_pages > 0: @@ -1309,9 +1308,6 @@ def alloc_decode( if need_tokens > self.available_size(): if not _relieve_for_alloc(self, need_tokens): return None - if self.need_sort and bs > len(self.free_virtual_ids): - self.merge_and_sort_free() - # Most decode steps reuse the prefix's tail page → num_new_pages == 0. if num_new_pages > 0: new_virtual_pages = self.free_virtual_ids[:num_new_pages].clone() @@ -3629,6 +3625,22 @@ def translate_kv_loc( result = self.full_attn_allocator.translate_kv_loc(loc, out=out) return result + def translate_kv_indices_for_transfer( + self, kv_indices: torch.Tensor + ) -> torch.Tensor: + """Translate virtual token ids to full-pool physical ids for PD.""" + return self.full_attn_allocator.translate_kv_loc(kv_indices.to(torch.int64)) + + def set_disagg_move_gate(self, gate: Callable[[], bool]) -> None: + """Install the PD compaction gate on both physical sub-allocators.""" + assert self.lazy_compaction, ( + "PD disaggregation with the unified memory pool requires lazy " + "compaction (eager free-path compaction moves pages under " + "in-flight transfers)." + ) + self.full_attn_allocator.disagg_move_gate = gate + self.swa_attn_allocator.disagg_move_gate = gate + def translate_loc_from_full_to_swa( self, kv_indices: torch.Tensor, @@ -3748,6 +3760,68 @@ def alloc_extend( self.swa_attn_allocator.alloc_with_virtual(new_virtual_pages) return out_indices # virtual TOKEN ids + def alloc_extend_swa_tail( + self, + prefix_lens: torch.Tensor, + prefix_lens_cpu: torch.Tensor, + seq_lens: torch.Tensor, + seq_lens_cpu: torch.Tensor, + last_loc: torch.Tensor, + extend_num_tokens: int, + swa_tail_len: int, + ) -> Optional[torch.Tensor]: + """Allocate full KV for an extend and SWA KV only for its aligned tail.""" + with record_function("UnifiedSWAAlloc.alloc_extend_swa_tail"): + assert self.page_size > 1 + assert len(seq_lens_cpu) == len(prefix_lens_cpu) == 1 + prefix_len = int(prefix_lens_cpu[0]) + seq_len = int(seq_lens_cpu[0]) + assert seq_len - prefix_len == extend_num_tokens + assert 0 <= swa_tail_len <= extend_num_tokens + tail_start = seq_len - swa_tail_len + assert prefix_len <= tail_start + assert swa_tail_len == 0 or tail_start % self.page_size == 0, ( + "unified SWA tail allocation requires a page-aligned tail; " + f"got extend_num_tokens={extend_num_tokens}, " + f"tail_len={swa_tail_len}, page_size={self.page_size}" + ) + + num_full_pages = get_num_new_pages( + seq_lens=seq_lens_cpu, + page_size=self.page_size, + prefix_lens=prefix_lens_cpu, + ) + num_swa_pages = (swa_tail_len + self.page_size - 1) // self.page_size + if not self.ensure_capacity( + num_full_pages * self.page_size, + num_swa_pages * self.page_size, + ): + return None + + fa = self.full_attn_allocator + new_virtual_pages = fa.free_virtual_ids[:num_full_pages].clone() + tail_virtual_pages = ( + new_virtual_pages[-num_swa_pages:] + if num_swa_pages > 0 + else new_virtual_pages[:0] + ) + out_indices = fa.alloc_extend( + prefix_lens, + prefix_lens_cpu, + seq_lens, + seq_lens_cpu, + last_loc, + extend_num_tokens, + num_new_pages=num_full_pages, + ) + assert out_indices is not None, ( + "UnifiedSWA.alloc_extend_swa_tail: full.alloc_extend returned " + "None after the capacity check passed" + ) + if num_swa_pages > 0: + self.swa_attn_allocator.alloc_with_virtual(tail_virtual_pages) + return out_indices + def alloc_decode( self, seq_lens: torch.Tensor, diff --git a/python/sglang/srt/mem_cache/unified_memory_pool.py b/python/sglang/srt/mem_cache/unified_memory_pool.py index aa96cbdd216b..d4e9c2f9a73c 100644 --- a/python/sglang/srt/mem_cache/unified_memory_pool.py +++ b/python/sglang/srt/mem_cache/unified_memory_pool.py @@ -624,19 +624,32 @@ def move_kv_cache(self, tgt_loc: torch.Tensor, src_loc: torch.Tensor): env[tgt_pages] = env[src_pages] def get_contiguous_buf_infos(self): - raise NotImplementedError( - "unified layout has no per-layer contiguous regions; " - "KV transfer / disaggregation is unsupported." + """Register the raw buffer as physical page envelopes for PD transfer. + + Full and SWA expose the same allocation with different envelope sizes; + the transfer backend preserves both logical entries while deduplicating + the underlying memory registration. + """ + assert self._unified_buffer.anchor_bytes(self._sub_pool_name) == 0 + raw = self._unified_buffer._raw + return [raw.data_ptr()], [raw.numel()], [self._page_bytes] + + def _physical_to_kernel_indices(self, indices: torch.Tensor) -> torch.Tensor: + """Map physical token ids to this pool's page-envelope row ids.""" + return ( + (indices // self.page_size) + * (self.page_size * self.kernel_page_blocks) + + indices % self.page_size ) def get_cpu_copy(self, indices, mamba_indices=None): - raise NotImplementedError( - "CPU offloading is unsupported under the unified layout." - ) + """Copy rows addressed by physical token ids from page envelopes.""" + return super().get_cpu_copy(self._physical_to_kernel_indices(indices)) def load_cpu_copy(self, kv_cache_cpu, indices, mamba_indices=None): - raise NotImplementedError( - "CPU offloading is unsupported under the unified layout." + """Restore rows addressed by physical token ids into page envelopes.""" + super().load_cpu_copy( + kv_cache_cpu, self._physical_to_kernel_indices(indices) ) def set_kv_buffer_prefix_valid(self, *args, **kwargs): @@ -1632,7 +1645,7 @@ def get_cpu_copy(self, indices, mamba_indices=None): swa_cpu = None if bool(valid.any().item()): swa_cpu = self.swa_kv_pool.get_cpu_copy(swa_phys[valid]) - return {"full": full_cpu, "swa": swa_cpu} + return {"full": full_cpu, "swa": swa_cpu, "swa_mask": valid.cpu()} def load_cpu_copy(self, kv_cache_cpu, indices, mamba_indices=None): assert self._full_allocator is not None @@ -1641,7 +1654,14 @@ def load_cpu_copy(self, kv_cache_cpu, indices, mamba_indices=None): if kv_cache_cpu.get("swa") is not None: assert self._swa_allocator is not None swa_phys = self._virt_tokens_to_phys_tokens(indices, self._swa_allocator) - self.swa_kv_pool.load_cpu_copy(kv_cache_cpu["swa"], swa_phys) + old_swa_mask = kv_cache_cpu["swa_mask"].to(indices.device) + assert old_swa_mask.shape == indices.shape + row_mask = (swa_phys >= 0)[old_swa_mask].cpu() + swa_phys = swa_phys[old_swa_mask][row_mask.to(indices.device)] + if swa_phys.numel() == 0: + return + swa_cpu = self._filter_swa_cpu_copy(kv_cache_cpu["swa"], row_mask) + self.swa_kv_pool.load_cpu_copy(swa_cpu, swa_phys) class UnifiedSWAPoolBundle(NamedTuple): diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index fa4b80ed3a0c..1bbe01e1877a 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -873,8 +873,8 @@ class ServerArgs: "Replace the statically-partitioned hybrid-model pools (full-attn KV + " "SWA/Mamba state) with one byte buffer split dynamically between " "sub-pools. Requires the Triton attention / linear-attn / Mamba " - "backends; not yet compatible with PD disaggregation or speculative " - "decoding.", + "backends. Supported PD-disaggregation and speculative-decoding " + "configurations are validated at startup.", NS("memory"), ] = False disable_chunked_prefix_cache: A[ From c4551a5f4d003f7eac7bcc3a3927b2216e5ea663 Mon Sep 17 00:00:00 2001 From: yhzhuang Date: Thu, 27 Aug 2026 15:16:44 -0700 Subject: [PATCH 14/55] Use decode request pool for unified SWA PD Co-authored-by: Yonghao Zhuang --- .../srt/mem_cache/kv_cache_configurator.py | 11 +---- .../model_executor/test_pool_configurator.py | 47 +++++++++++++++++++ 2 files changed, 48 insertions(+), 10 deletions(-) diff --git a/python/sglang/srt/mem_cache/kv_cache_configurator.py b/python/sglang/srt/mem_cache/kv_cache_configurator.py index 9516c6f5cfbd..53c7fb126e09 100644 --- a/python/sglang/srt/mem_cache/kv_cache_configurator.py +++ b/python/sglang/srt/mem_cache/kv_cache_configurator.py @@ -822,16 +822,7 @@ def _init_unified_swa_pools( assert ( not self.use_mla_backend ), "unified memory pool does not support MLA-SWA hybrid yet" - # Mirror the non-shared path's extra_max_context_len computation. - extra_max_context_len = 4 - if get_spec().speculative_num_draft_tokens is not None: - extra_max_context_len += get_spec().speculative_num_draft_tokens - req_to_token_pool = ReqToTokenPool( - size=max_num_reqs, - max_context_len=self.model_config.context_len + extra_max_context_len, - device=self.device, - enable_memory_saver=get_exec().features.enable_memory_saver, - ) + req_to_token_pool = self._build_req_to_token_pool(max_num_reqs=max_num_reqs) head_num = self.model_config.get_num_kv_heads( get_parallel().attn_tp_size, get_parallel().attn_dcp_size diff --git a/test/registered/unit/model_executor/test_pool_configurator.py b/test/registered/unit/model_executor/test_pool_configurator.py index 3e8cf29b283c..43fa30d592c7 100644 --- a/test/registered/unit/model_executor/test_pool_configurator.py +++ b/test/registered/unit/model_executor/test_pool_configurator.py @@ -337,6 +337,53 @@ def test_memory_utilization(self): self.assertLessEqual(used, available) self.assertGreater(used, available * 0.99) + def test_unified_pd_decode_uses_decode_req_pool(self): + import torch + + from sglang.srt.disaggregation.decode import DecodeReqToTokenPool + from sglang.srt.mem_cache.kv_cache_configurator import KVCacheConfigurator + + mr = _make_model_runner( + self, + is_hybrid_swa=True, + full_attention_layer_ids=[0], + swa_attention_layer_ids=[1], + disaggregation_mode="decode", + disaggregation_decode_extra_slots=2, + enable_unified_memory=True, + ) + configurator = object.__new__(KVCacheConfigurator) + configurator.device = "cpu" + configurator.model_config = mr.model_config + configurator.kv_cache_dtype = torch.bfloat16 + configurator.page_size = 1 + configurator.is_hybrid_swa = True + configurator.is_hybrid_swa_compress = False + configurator.use_mla_backend = False + configurator.layer_info = mr.layer_info + configurator.forward_stream = None + configurator.mambaish_config = None + + initialized_pools = SimpleNamespace( + unified_memory_pool=MagicMock(), + token_to_kv_pool=MagicMock(), + token_to_kv_pool_allocator=MagicMock(), + ) + with ( + mock_cpu_env(), + patch( + "sglang.srt.mem_cache.unified_memory_pool.init_unified_swa_pools", + return_value=initialized_pools, + ), + ): + pools = configurator._init_unified_swa_pools( + max_num_reqs=4, + unified_memory_pool_bytes=4096, + ) + + self.assertIsInstance(pools.req_to_token_pool, DecodeReqToTokenPool) + self.assertEqual(pools.req_to_token_pool.pre_alloc_size, 2) + @patch( "sglang.srt.mem_cache.kv_cache_configurator.calculate_mla_kv_cache_dim", return_value=576, From 623f7620ae51f5029306804ff763e001bbb64839 Mon Sep 17 00:00:00 2001 From: yhzhuang Date: Thu, 27 Aug 2026 16:55:18 -0700 Subject: [PATCH 15/55] Reject unified Mooncake speculative transfers Co-authored-by: yhzhuang --- python/sglang/srt/arg_groups/kv_cache_hook.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/python/sglang/srt/arg_groups/kv_cache_hook.py b/python/sglang/srt/arg_groups/kv_cache_hook.py index dbed1eaa1e28..5004af4faca7 100644 --- a/python/sglang/srt/arg_groups/kv_cache_hook.py +++ b/python/sglang/srt/arg_groups/kv_cache_hook.py @@ -210,6 +210,13 @@ def handle_unified_memory_pool(server_args: Any) -> None: "the mooncake transfer backend; got " f"{cfg.disaggregation_transfer_backend!r}." ) + assert not ( + cfg.disaggregation_transfer_backend == "mooncake" + and cfg.speculative_algorithm is not None + ), ( + "--enable-unified-memory with PD disaggregation does not support " + "speculative decoding with the Mooncake transfer backend." + ) assert not envs.SGLANG_DISABLE_LAZY_COMPACTION.get(), ( "--enable-unified-memory with PD disaggregation requires lazy " "compaction; unset SGLANG_DISABLE_LAZY_COMPACTION." From c9b4c29b4acc5cdc2686ffaac536a48794ec9bbd Mon Sep 17 00:00:00 2001 From: yhzhuang Date: Thu, 27 Aug 2026 18:04:21 -0700 Subject: [PATCH 16/55] Simplify unified draft index plumbing --- python/sglang/srt/disaggregation/base/conn.py | 3 -- python/sglang/srt/disaggregation/decode.py | 5 +- python/sglang/srt/disaggregation/fake/conn.py | 19 -------- python/sglang/srt/disaggregation/prefill.py | 5 +- .../srt/mem_cache/unified_memory_pool.py | 12 ++--- .../model_executor/test_pool_configurator.py | 47 ------------------- 6 files changed, 8 insertions(+), 83 deletions(-) diff --git a/python/sglang/srt/disaggregation/base/conn.py b/python/sglang/srt/disaggregation/base/conn.py index 00e1d84ee775..5c6717ca0c51 100644 --- a/python/sglang/srt/disaggregation/base/conn.py +++ b/python/sglang/srt/disaggregation/base/conn.py @@ -45,9 +45,6 @@ class KVArgs: kv_data_ptrs: List[int] kv_data_lens: List[int] kv_item_lens: List[int] - # Target KV buffers come first; any remaining buffers belong to an - # independent draft pool and may require a different index vector. - num_target_kv_data_ptrs: int kv_layer_ids: List[int] kv_cache_dtype_str: str aux_data_ptrs: List[int] diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 2c08a8fad921..a66a87739cb0 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -538,10 +538,9 @@ def _init_kv_manager(self) -> CommonKVManager: kv_item_lens += device_kv_item_lens[c4_layer_num:] kv_data_mem_kinds += ["VRAM"] * len(device_kv_data_ptrs[c4_layer_num:]) num_draft_entries = 0 - kv_args.num_target_kv_data_ptrs = len(kv_data_ptrs) if self.draft_token_to_kv_pool is not None: - # We should also transfer draft model kv cache. The indices are - # always shared with a target model. + # Draft KV shares target virtual ids. Unified target KV is transferred + # with physical ids, so it needs a separate draft index vector. draft_kv_data_ptrs, draft_kv_data_lens, draft_kv_item_lens = ( self.draft_token_to_kv_pool.get_contiguous_buf_infos() ) diff --git a/python/sglang/srt/disaggregation/fake/conn.py b/python/sglang/srt/disaggregation/fake/conn.py index 53389c88a282..15a76c039aa4 100644 --- a/python/sglang/srt/disaggregation/fake/conn.py +++ b/python/sglang/srt/disaggregation/fake/conn.py @@ -85,15 +85,6 @@ def send( f"FakeKVSender send with kv_indices: {kv_indices}, state_indices: {state_indices}" ) - def send_with_draft_indices( - self, - kv_indices: npt.NDArray[np.int32], - draft_kv_indices: npt.NDArray[np.int32], - state_indices: Optional[List] = None, - num_kv_tokens: Optional[int] = None, - ): - self.send(kv_indices, state_indices, num_kv_tokens) - def failure_exception(self): raise Exception("Fake KVSender Exception") @@ -142,16 +133,6 @@ def send_metadata( f"FakeKVReceiver send_metadata with kv_indices: {kv_indices}, aux_index: {aux_index}, state_indices: {state_indices}" ) - def send_metadata_with_draft_indices( - self, - kv_indices: npt.NDArray[np.int32], - draft_kv_indices: npt.NDArray[np.int32], - aux_index: Optional[int] = None, - state_indices: Optional[List] = None, - decode_prefix_len: Optional[int] = None, - ): - self.send_metadata(kv_indices, aux_index, state_indices, decode_prefix_len) - def failure_exception(self): raise Exception("Fake KVReceiver Exception") diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index 8ab0e77ec31a..8d08befecf0a 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -234,7 +234,6 @@ def _init_kv_manager(self) -> CommonKVManager: kv_data_ptrs, kv_data_lens, kv_item_lens = ( self.token_to_kv_pool.get_contiguous_buf_infos() ) - kv_args.num_target_kv_data_ptrs = len(kv_data_ptrs) kv_args.prefill_end_layer = ( kv_args.prefill_start_layer + len(kv_data_ptrs) if layer_shard_enabled @@ -244,8 +243,8 @@ def _init_kv_manager(self) -> CommonKVManager: draft_kv_pool = self.draft_token_to_kv_pool if transfer_draft_cache else None num_draft_entries = 0 if draft_kv_pool is not None: - # We should also transfer draft model kv cache. The indices are - # always shared with a target model. + # Draft KV shares target virtual ids. Unified target KV is transferred + # with physical ids, so it needs a separate draft index vector. draft_kv_data_ptrs, draft_kv_data_lens, draft_kv_item_lens = ( draft_kv_pool.get_contiguous_buf_infos() ) diff --git a/python/sglang/srt/mem_cache/unified_memory_pool.py b/python/sglang/srt/mem_cache/unified_memory_pool.py index d4e9c2f9a73c..ecceb0e82c0a 100644 --- a/python/sglang/srt/mem_cache/unified_memory_pool.py +++ b/python/sglang/srt/mem_cache/unified_memory_pool.py @@ -636,11 +636,9 @@ def get_contiguous_buf_infos(self): def _physical_to_kernel_indices(self, indices: torch.Tensor) -> torch.Tensor: """Map physical token ids to this pool's page-envelope row ids.""" - return ( - (indices // self.page_size) - * (self.page_size * self.kernel_page_blocks) - + indices % self.page_size - ) + return (indices // self.page_size) * ( + self.page_size * self.kernel_page_blocks + ) + indices % self.page_size def get_cpu_copy(self, indices, mamba_indices=None): """Copy rows addressed by physical token ids from page envelopes.""" @@ -648,9 +646,7 @@ def get_cpu_copy(self, indices, mamba_indices=None): def load_cpu_copy(self, kv_cache_cpu, indices, mamba_indices=None): """Restore rows addressed by physical token ids into page envelopes.""" - super().load_cpu_copy( - kv_cache_cpu, self._physical_to_kernel_indices(indices) - ) + super().load_cpu_copy(kv_cache_cpu, self._physical_to_kernel_indices(indices)) def set_kv_buffer_prefix_valid(self, *args, **kwargs): raise NotImplementedError( diff --git a/test/registered/unit/model_executor/test_pool_configurator.py b/test/registered/unit/model_executor/test_pool_configurator.py index 43fa30d592c7..3e8cf29b283c 100644 --- a/test/registered/unit/model_executor/test_pool_configurator.py +++ b/test/registered/unit/model_executor/test_pool_configurator.py @@ -337,53 +337,6 @@ def test_memory_utilization(self): self.assertLessEqual(used, available) self.assertGreater(used, available * 0.99) - def test_unified_pd_decode_uses_decode_req_pool(self): - import torch - - from sglang.srt.disaggregation.decode import DecodeReqToTokenPool - from sglang.srt.mem_cache.kv_cache_configurator import KVCacheConfigurator - - mr = _make_model_runner( - self, - is_hybrid_swa=True, - full_attention_layer_ids=[0], - swa_attention_layer_ids=[1], - disaggregation_mode="decode", - disaggregation_decode_extra_slots=2, - enable_unified_memory=True, - ) - configurator = object.__new__(KVCacheConfigurator) - configurator.device = "cpu" - configurator.model_config = mr.model_config - configurator.kv_cache_dtype = torch.bfloat16 - configurator.page_size = 1 - configurator.is_hybrid_swa = True - configurator.is_hybrid_swa_compress = False - configurator.use_mla_backend = False - configurator.layer_info = mr.layer_info - configurator.forward_stream = None - configurator.mambaish_config = None - - initialized_pools = SimpleNamespace( - unified_memory_pool=MagicMock(), - token_to_kv_pool=MagicMock(), - token_to_kv_pool_allocator=MagicMock(), - ) - with ( - mock_cpu_env(), - patch( - "sglang.srt.mem_cache.unified_memory_pool.init_unified_swa_pools", - return_value=initialized_pools, - ), - ): - pools = configurator._init_unified_swa_pools( - max_num_reqs=4, - unified_memory_pool_bytes=4096, - ) - - self.assertIsInstance(pools.req_to_token_pool, DecodeReqToTokenPool) - self.assertEqual(pools.req_to_token_pool.pre_alloc_size, 2) - @patch( "sglang.srt.mem_cache.kv_cache_configurator.calculate_mla_kv_cache_dim", return_value=576, From 5a7d1efe3b0916487c663a4e8fe401d32dbb8acd Mon Sep 17 00:00:00 2001 From: yhzhuang Date: Tue, 1 Sep 2026 15:51:23 -0700 Subject: [PATCH 17/55] Fix unified SWA PD capacity accounting --- python/sglang/srt/arg_groups/kv_cache_hook.py | 11 +- python/sglang/srt/disaggregation/decode.py | 210 ++++++++++++++---- python/sglang/srt/mem_cache/common.py | 24 +- .../sglang/srt/mem_cache/kv_cache_builder.py | 6 +- .../srt/mem_cache/multi_ended_allocator.py | 1 - .../srt/mem_cache/unified_memory_pool.py | 4 +- python/sglang/srt/server_args.py | 3 + .../test_decode_queue_cleanup.py | 1 + .../mem_cache/test_decode_radix_lock_ref.py | 3 +- .../unit/mem_cache/test_unified_mha_views.py | 12 +- 10 files changed, 204 insertions(+), 71 deletions(-) diff --git a/python/sglang/srt/arg_groups/kv_cache_hook.py b/python/sglang/srt/arg_groups/kv_cache_hook.py index 5004af4faca7..8fe90e7b8029 100644 --- a/python/sglang/srt/arg_groups/kv_cache_hook.py +++ b/python/sglang/srt/arg_groups/kv_cache_hook.py @@ -203,11 +203,12 @@ def handle_unified_memory_pool(server_args: Any) -> None: "--enable-unified-memory with PD disaggregation does not support " "pipeline parallelism (--pp-size > 1)." ) - # Constraints of the whole-envelope transfer; see - # UnifiedMLATokenToKVPool.get_contiguous_buf_infos. - assert cfg.disaggregation_transfer_backend == "mooncake", ( - "--enable-unified-memory with PD disaggregation supports only " - "the mooncake transfer backend; got " + # Constraints of the whole-envelope transfer; see the unified MHA and + # MLA pool get_contiguous_buf_infos implementations. + supported_backends = server_args._unified_memory_pd_transfer_backends() + assert cfg.disaggregation_transfer_backend in supported_backends, ( + "--enable-unified-memory with PD disaggregation supports only these " + f"transfer backends: {', '.join(sorted(supported_backends))}; got " f"{cfg.disaggregation_transfer_backend!r}." ) assert not ( diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index a66a87739cb0..9eb68a1e0e2e 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -78,6 +78,7 @@ EvictParams, ) from sglang.srt.mem_cache.common import ( + evict_from_tree_cache, kv_to_page_indices, page_align_floor, release_kv_cache, @@ -90,6 +91,9 @@ KVCache, ReqToTokenPool, ) +from sglang.srt.mem_cache.multi_ended_allocator import ( + UnifiedSWATokenToKVPoolAllocator, +) from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool from sglang.srt.observability.req_time_stats import ( set_schedule_time_batch, @@ -409,6 +413,87 @@ def _uses_swa_tail_prealloc(self) -> bool: and hasattr(self.token_to_kv_pool_allocator, "alloc_extend_swa_tail") ) + def _supports_unified_swa_reservation(self) -> bool: + allocator = self.token_to_kv_pool_allocator + return ( + isinstance(allocator, UnifiedSWATokenToKVPoolAllocator) + and allocator.supports_asymmetric_reservation + ) + + def _unified_swa_reservation_fits( + self, + full_tokens: int, + swa_tokens: int, + *, + full_allocatable_tokens: Optional[int] = None, + swa_allocatable_tokens: Optional[int] = None, + empty_pool: bool = False, + ) -> bool: + allocator = self.token_to_kv_pool_allocator + if not self._supports_unified_swa_reservation(): + return True + assert isinstance(allocator, UnifiedSWATokenToKVPoolAllocator) + + full_evictable = swa_evictable = 0 + if not empty_pool: + full_evictable = self._radix_full_evictable() + swa_evictable = self.tree_cache.swa_evictable_size() + if full_allocatable_tokens is not None: + full_tokens += ( + allocator.full_available_size() + + full_evictable + - full_allocatable_tokens + ) + if swa_allocatable_tokens is not None: + swa_tokens += ( + allocator.swa_available_size() + + swa_evictable + - swa_allocatable_tokens + ) + + return allocator.can_reserve( + full_tokens, + swa_tokens, + full_evictable_tokens=full_evictable, + swa_evictable_tokens=swa_evictable, + empty_pool=empty_pool, + ) + + def _prealloc_reservation_fits( + self, + full_tokens: int, + swa_tokens: int, + *, + full_allocatable_tokens: int, + swa_allocatable_tokens: Optional[int], + uses_swa_tail_prealloc: bool, + ) -> bool: + if self._supports_unified_swa_reservation(): + if not uses_swa_tail_prealloc: + # alloc()/alloc_extend() charge both sides symmetrically even + # when the page-size-one path does not preallocate an SWA tail. + return ( + full_tokens <= full_allocatable_tokens + and self._unified_swa_reservation_fits( + full_tokens, + full_tokens, + ) + ) + assert swa_allocatable_tokens is not None + return self._unified_swa_reservation_fits( + full_tokens, + swa_tokens, + full_allocatable_tokens=full_allocatable_tokens, + swa_allocatable_tokens=swa_allocatable_tokens, + ) + if not uses_swa_tail_prealloc: + return full_tokens <= full_allocatable_tokens + assert swa_allocatable_tokens is not None + return ( + full_tokens <= full_allocatable_tokens + and swa_tokens <= swa_allocatable_tokens + ) + def _release_matched_prefix_lock(self, req: Req) -> None: params = DecLockRefParams(swa_uuid_for_lock=req.swa_uuid_for_lock) if req.swa_prefix_lock_released: @@ -418,16 +503,33 @@ def _release_matched_prefix_lock(self, req: Req) -> None: self.tree_cache.dec_lock_ref(req.last_node, params) def _reclaim_swa_tail_capacity( - self, swa_tail_len: int, req_id: str + self, swa_tail_len: int, req_id: str, *, full_len: int = 0 ) -> Optional[str]: - page_size = self.token_to_kv_pool_allocator.page_size + allocator = self.token_to_kv_pool_allocator + page_size = allocator.page_size + if self._supports_unified_swa_reservation(): + assert isinstance(allocator, UnifiedSWATokenToKVPoolAllocator) + full_required = ceil_align(full_len, page_size) + swa_required = ceil_align(swa_tail_len, page_size) + evict_from_tree_cache( + self.tree_cache, + full_required, + swa_num_tokens=swa_required, + ) + if allocator.ensure_capacity(full_required, swa_required): + return None + return ( + "Unified FULL/SWA byte reclamation insufficient: " + f"needed=({full_required}, {swa_required}), req={req_id}" + ) + required = ceil_align(swa_tail_len, page_size) - available = self.token_to_kv_pool_allocator.swa_available_size() + available = allocator.swa_available_size() if available < required: self.tree_cache.evict_for_alloc( EvictParams(swa_num_tokens=required - available) ) - available = self.token_to_kv_pool_allocator.swa_available_size() + available = allocator.swa_available_size() if available < required: return ( f"SWA eviction insufficient: needed={required}, " @@ -758,6 +860,22 @@ def _check_if_req_exceed_kv_capacity(self, req: Req) -> bool: else: capacity = self.max_total_num_tokens input_len = self._rebootstrap_prefill_len(req) + if self._supports_unified_swa_reservation(): + full_required, swa_required = self._prealloc_required_tokens(req) + if not self._uses_swa_tail_prealloc(): + swa_required = full_required + if not self._unified_swa_reservation_fits( + full_required, swa_required, empty_pool=True + ): + message = ( + f"Request {req.rid} exceeds the unified FULL/SWA KV byte " + f"budget: full={full_required}, swa={swa_required}" + ) + logger.error(message) + prepare_abort(req, message, status_code=HTTPStatus.BAD_REQUEST) + self.scheduler.output_streamer.stream_output([req], req.return_logprob) + return True + return False if input_len > capacity: message = f"Request {req.rid} exceeds the maximum number of tokens: {input_len} > {capacity}" logger.error(message) @@ -826,10 +944,25 @@ def resume_retracted_reqs( break full_required, swa_required = self._prealloc_required_tokens(req) - if full_required > full_allocatable_tokens: - break - if uses_swa_tail_prealloc and swa_required > swa_allocatable_tokens: + if not self._prealloc_reservation_fits( + full_required, + swa_required, + full_allocatable_tokens=full_allocatable_tokens, + swa_allocatable_tokens=( + swa_allocatable_tokens if uses_swa_tail_prealloc else None + ), + uses_swa_tail_prealloc=uses_swa_tail_prealloc, + ): break + + if uses_swa_tail_prealloc and self._supports_unified_swa_reservation(): + full_len, swa_len = self._prealloc_kv_lens(req) + if ( + self._reclaim_swa_tail_capacity(swa_len, req.rid, full_len=full_len) + is not None + ): + break + resumed_reqs.append(req) indices_to_remove.add(i) req.is_retracted = False @@ -1268,27 +1401,18 @@ def pop_preallocated( required_alloc_tokens + self.num_reserved_decode_tokens ) - if ( - max( - required_tokens_for_request, - origin_input_len - - prefix_len - + min( - decode_req.req.sampling_params.max_new_tokens, - CLIP_MAX_NEW_TOKEN, - ) - - retractable_tokens, + full_required_for_admission = max( + required_tokens_for_request, + origin_input_len + - prefix_len + + min( + decode_req.req.sampling_params.max_new_tokens, + CLIP_MAX_NEW_TOKEN, ) - > full_allocatable_tokens - ): - if prefix_match is not None and prefix_match.l1_prefix_len > 0: - self._release_matched_prefix_lock(decode_req.req) - break - if required_tokens_for_request > full_allocatable_tokens: - if prefix_match is not None and prefix_match.l1_prefix_len > 0: - self._release_matched_prefix_lock(decode_req.req) - break - + - retractable_tokens, + ) + swa_required_for_admission = 0 + swa_len = 0 if uses_swa_tail_prealloc: _, swa_required = self._prealloc_required_tokens(decode_req.req) _, swa_len = self._prealloc_kv_lens(decode_req.req) @@ -1296,19 +1420,29 @@ def pop_preallocated( decode_req.req.sampling_params.max_new_tokens, CLIP_MAX_NEW_TOKEN, ) - if ( - max( - swa_required, - swa_len + max_new_tokens - retractable_swa_tokens, - ) - > swa_allocatable_tokens - ): - if prefix_match is not None and prefix_match.l1_prefix_len > 0: - self._release_matched_prefix_lock(decode_req.req) - break + swa_required_for_admission = max( + swa_required, + swa_len + max_new_tokens - retractable_swa_tokens, + ) + if not self._prealloc_reservation_fits( + full_required_for_admission, + swa_required_for_admission, + full_allocatable_tokens=full_allocatable_tokens, + swa_allocatable_tokens=( + swa_allocatable_tokens if uses_swa_tail_prealloc else None + ), + uses_swa_tail_prealloc=uses_swa_tail_prealloc, + ): + if prefix_match is not None and prefix_match.l1_prefix_len > 0: + self._release_matched_prefix_lock(decode_req.req) + break + + if uses_swa_tail_prealloc: reclaim_error = self._reclaim_swa_tail_capacity( - swa_len, decode_req.req.rid + swa_len, + decode_req.req.rid, + full_len=required_alloc_tokens, ) if reclaim_error is not None: if prefix_match is not None and prefix_match.l1_prefix_len > 0: diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index e4dac917811c..96722b0c127e 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/python/sglang/srt/mem_cache/common.py @@ -165,7 +165,11 @@ def maybe_cache_unfinished_req(req: Req, tree_cache: BasePrefixCache, **kwargs): tree_cache.cache_unfinished_req(req, **kwargs) -def evict_from_tree_cache(tree_cache: BasePrefixCache | None, num_tokens: int): +def evict_from_tree_cache( + tree_cache: BasePrefixCache | None, + num_tokens: int, + swa_num_tokens: Optional[int] = None, +): if tree_cache is None: return @@ -182,7 +186,8 @@ def evict_from_tree_cache(tree_cache: BasePrefixCache | None, num_tokens: int): isinstance(allocator, UnifiedSWATokenToKVPoolAllocator) and allocator.supports_asymmetric_reservation ): - if allocator.can_reserve(num_tokens, num_tokens): + required_swa = num_tokens if swa_num_tokens is None else swa_num_tokens + if allocator.can_reserve(num_tokens, required_swa): return page_size = allocator.page_size @@ -190,7 +195,7 @@ def evict_from_tree_cache(tree_cache: BasePrefixCache | None, num_tokens: int): swa_evictable = tree_cache.swa_evictable_size() // page_size * page_size if not allocator.can_reserve( num_tokens, - num_tokens, + required_swa, full_evictable_tokens=full_evictable, swa_evictable_tokens=swa_evictable, ): @@ -210,7 +215,7 @@ def minimum_reclaim(max_tokens: int, fits) -> int: full_evictable, lambda value: allocator.can_reserve( num_tokens, - num_tokens, + required_swa, full_evictable_tokens=value, swa_evictable_tokens=swa_evictable, ), @@ -219,7 +224,7 @@ def minimum_reclaim(max_tokens: int, fits) -> int: swa_evictable, lambda value: allocator.can_reserve( num_tokens, - num_tokens, + required_swa, full_evictable_tokens=full_reclaim, swa_evictable_tokens=value, ), @@ -230,22 +235,23 @@ def minimum_reclaim(max_tokens: int, fits) -> int: swa_num_tokens=swa_reclaim, ) ) - if not allocator.can_reserve(num_tokens, num_tokens): + if not allocator.can_reserve(num_tokens, required_swa): logger.warning( "Unified SWA eviction did not satisfy the planned reservation: " - f"requested=({num_tokens}, {num_tokens}), " + f"requested=({num_tokens}, {required_swa}), " f"planned=({full_reclaim}, {swa_reclaim}), " f"evicted=({evicted.num_tokens_evicted}, " f"{evicted.swa_num_tokens_evicted})" ) elif isinstance(allocator, SWATokenToKVPoolAllocator): # Hybrid allocator + required_swa = num_tokens if swa_num_tokens is None else swa_num_tokens full_available_size = allocator.full_available_size() swa_available_size = allocator.swa_available_size() - if full_available_size < num_tokens or swa_available_size < num_tokens: + if full_available_size < num_tokens or swa_available_size < required_swa: full_num_tokens = max(0, num_tokens - full_available_size) - swa_num_tokens = max(0, num_tokens - swa_available_size) + swa_num_tokens = max(0, required_swa - swa_available_size) tree_cache.evict_for_alloc( EvictParams(num_tokens=full_num_tokens, swa_num_tokens=swa_num_tokens) ) diff --git a/python/sglang/srt/mem_cache/kv_cache_builder.py b/python/sglang/srt/mem_cache/kv_cache_builder.py index b841dc312cec..5e69c060df24 100644 --- a/python/sglang/srt/mem_cache/kv_cache_builder.py +++ b/python/sglang/srt/mem_cache/kv_cache_builder.py @@ -150,10 +150,8 @@ def resolve_decode_retraction_backup(*, tp_worker: BaseTpWorker) -> str: if tp_worker.is_hybrid_swa else None ) - # Unified page-envelope host pools are enabled by a separate feature. - # Until then, keep decode retraction on the direct CPU-tensor backup. - # Host-pool retraction transfers full and sliding-window components - # only, so a model with recurrent state stays on cpu_tensor. + # Host-pool retraction does not address unified page envelopes or + # recurrent state, so those configurations stay on cpu_tensor. supports_host_pool = ( not memory.enable_unified_memory and not uses_ssm_state(tp_worker.model_runner.model_config) diff --git a/python/sglang/srt/mem_cache/multi_ended_allocator.py b/python/sglang/srt/mem_cache/multi_ended_allocator.py index 816fbc64dbe8..fb27ced8e960 100644 --- a/python/sglang/srt/mem_cache/multi_ended_allocator.py +++ b/python/sglang/srt/mem_cache/multi_ended_allocator.py @@ -3776,7 +3776,6 @@ def alloc_extend_swa_tail( assert len(seq_lens_cpu) == len(prefix_lens_cpu) == 1 prefix_len = int(prefix_lens_cpu[0]) seq_len = int(seq_lens_cpu[0]) - assert seq_len - prefix_len == extend_num_tokens assert 0 <= swa_tail_len <= extend_num_tokens tail_start = seq_len - swa_tail_len assert prefix_len <= tail_start diff --git a/python/sglang/srt/mem_cache/unified_memory_pool.py b/python/sglang/srt/mem_cache/unified_memory_pool.py index ecceb0e82c0a..1a4b9acf3aab 100644 --- a/python/sglang/srt/mem_cache/unified_memory_pool.py +++ b/python/sglang/srt/mem_cache/unified_memory_pool.py @@ -635,17 +635,15 @@ def get_contiguous_buf_infos(self): return [raw.data_ptr()], [raw.numel()], [self._page_bytes] def _physical_to_kernel_indices(self, indices: torch.Tensor) -> torch.Tensor: - """Map physical token ids to this pool's page-envelope row ids.""" return (indices // self.page_size) * ( self.page_size * self.kernel_page_blocks ) + indices % self.page_size def get_cpu_copy(self, indices, mamba_indices=None): - """Copy rows addressed by physical token ids from page envelopes.""" + """Translate physical host-pool ids for the page-major parent path.""" return super().get_cpu_copy(self._physical_to_kernel_indices(indices)) def load_cpu_copy(self, kv_cache_cpu, indices, mamba_indices=None): - """Restore rows addressed by physical token ids into page envelopes.""" super().load_cpu_copy(kv_cache_cpu, self._physical_to_kernel_indices(indices)) def set_kv_buffer_prefix_valid(self, *args, **kwargs): diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 1bbe01e1877a..957eb96f8286 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -3779,6 +3779,9 @@ def replace_resolved(self, source: str, **changes: Any) -> ServerArgs: # _handle_page_major_kv_layout); the model-family gate is enforced at pool # construction in model_runner_kv_cache_mixin._init_pools. + def _unified_memory_pd_transfer_backends(self) -> set[str]: + return {"mooncake"} + @staticmethod def add_cli_args(parser: argparse.ArgumentParser): diff --git a/test/registered/unit/disaggregation/test_decode_queue_cleanup.py b/test/registered/unit/disaggregation/test_decode_queue_cleanup.py index 7b5b9f8a06a0..61b2e52454c6 100644 --- a/test/registered/unit/disaggregation/test_decode_queue_cleanup.py +++ b/test/registered/unit/disaggregation/test_decode_queue_cleanup.py @@ -212,6 +212,7 @@ def test_swa_reclaim_failure_rejects_only_request(self): ) queue._hicache_pending_restore_tokens = MagicMock(return_value=0) queue._pre_alloc = MagicMock() + queue.token_to_kv_pool_allocator = MagicMock() queue.req_to_token_pool = MagicMock() queue.req_to_token_pool.available_size.return_value = 1 queue.req_to_metadata_buffer_idx_allocator = MagicMock() diff --git a/test/registered/unit/mem_cache/test_decode_radix_lock_ref.py b/test/registered/unit/mem_cache/test_decode_radix_lock_ref.py index 73528e65bf91..7bb7453f1a31 100644 --- a/test/registered/unit/mem_cache/test_decode_radix_lock_ref.py +++ b/test/registered/unit/mem_cache/test_decode_radix_lock_ref.py @@ -393,6 +393,7 @@ def test_pop_preallocated_rechecks_budget_after_lock(self): queue._update_handshake_waiters = MagicMock() queue._uses_swa_tail_prealloc = MagicMock(return_value=True) queue._swa_tail_len = MagicMock(return_value=8) + queue._prealloc_required_tokens = MagicMock(return_value=(8, 8)) queue._swa_aware_allocatable_token_budgets = MagicMock(return_value=(8, 8)) queue._swa_tail_allocatable_token_budget = MagicMock(return_value=8) queue._match_prefix_and_lock = MagicMock( @@ -449,7 +450,7 @@ def test_pop_preallocated_rechecks_budget_after_lock(self): skip_swa=True, ) self.assertFalse(req.swa_prefix_lock_released) - queue._swa_tail_len.assert_called_once_with(8) + queue._swa_tail_len.assert_called_with(8) queue._allocatable_token_budgets.assert_called_once() def test_repeated_incremental_no_leak(self): diff --git a/test/registered/unit/mem_cache/test_unified_mha_views.py b/test/registered/unit/mem_cache/test_unified_mha_views.py index b0be816ce612..a8deff504933 100644 --- a/test/registered/unit/mem_cache/test_unified_mha_views.py +++ b/test/registered/unit/mem_cache/test_unified_mha_views.py @@ -457,17 +457,9 @@ def test_move_kv_cache_relocates_whole_envelopes(self): "envelope move did not relocate exactly the named pages", ) - def test_transfer_entry_points_fail_loud(self): - """PD / CPU-copy entry points assume per-layer buffers indexed by TOKEN - id; against the row space they would silently mis-index (or hit a - missing-attr AttributeError). Every one of them must raise.""" + def test_prefix_valid_entry_point_fails_loud(self): + """Prefix-valid writes still assume token-major buffer indexing.""" _, pool = _make_pool_and_kv(1) - with self.assertRaises(NotImplementedError): - pool.get_contiguous_buf_infos() - with self.assertRaises(NotImplementedError): - pool.get_cpu_copy(torch.tensor([1])) - with self.assertRaises(NotImplementedError): - pool.load_cpu_copy(None, torch.tensor([1])) with self.assertRaises(NotImplementedError): pool.set_kv_buffer_prefix_valid() From 08a06eb0b1c7b9e830d09ce60768dea90f78f7c2 Mon Sep 17 00:00:00 2001 From: Yonghao Zhuang Date: Wed, 2 Sep 2026 00:43:20 -0700 Subject: [PATCH 18/55] Prefer same-side eviction for unified SWA --- python/sglang/srt/mem_cache/common.py | 22 ++++++++++++---------- 1 file changed, 12 insertions(+), 10 deletions(-) diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index e4dac917811c..55ecc27e46f6 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/python/sglang/srt/mem_cache/common.py @@ -182,7 +182,7 @@ def evict_from_tree_cache(tree_cache: BasePrefixCache | None, num_tokens: int): isinstance(allocator, UnifiedSWATokenToKVPoolAllocator) and allocator.supports_asymmetric_reservation ): - if allocator.can_reserve(num_tokens, num_tokens): + if allocator.ensure_capacity(num_tokens, num_tokens): return page_size = allocator.page_size @@ -206,31 +206,33 @@ def minimum_reclaim(max_tokens: int, fits) -> int: lo = mid + 1 return lo * page_size - full_reclaim = minimum_reclaim( - full_evictable, + swa_reclaim = minimum_reclaim( + swa_evictable, lambda value: allocator.can_reserve( num_tokens, num_tokens, - full_evictable_tokens=value, - swa_evictable_tokens=swa_evictable, + full_evictable_tokens=full_evictable, + swa_evictable_tokens=value, ), ) - swa_reclaim = minimum_reclaim( - swa_evictable, + full_reclaim = minimum_reclaim( + full_evictable, lambda value: allocator.can_reserve( num_tokens, num_tokens, - full_evictable_tokens=full_reclaim, - swa_evictable_tokens=value, + full_evictable_tokens=value, + swa_evictable_tokens=swa_reclaim, ), ) + # FULL is the primary allocation. Reclaim SWA only when all evictable + # FULL plus compaction still cannot satisfy the shared byte demand. evicted = tree_cache.evict_for_alloc( EvictParams( num_tokens=full_reclaim, swa_num_tokens=swa_reclaim, ) ) - if not allocator.can_reserve(num_tokens, num_tokens): + if not allocator.ensure_capacity(num_tokens, num_tokens): logger.warning( "Unified SWA eviction did not satisfy the planned reservation: " f"requested=({num_tokens}, {num_tokens}), " From fe5787a193c87935f15fc6a77fc34e0af19db0ad Mon Sep 17 00:00:00 2001 From: yhzhuang Date: Wed, 2 Sep 2026 11:13:58 -0700 Subject: [PATCH 19/55] Prefer SWA tombstones before unified FULL eviction --- python/sglang/srt/mem_cache/common.py | 74 +++++++-------- .../srt/mem_cache/multi_ended_allocator.py | 89 +++++++++++++++++-- .../srt/mem_cache/unified_radix_cache.py | 73 ++++++++++++++- 3 files changed, 194 insertions(+), 42 deletions(-) diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index 55ecc27e46f6..536a311be0b3 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/python/sglang/srt/mem_cache/common.py @@ -195,48 +195,50 @@ def evict_from_tree_cache(tree_cache: BasePrefixCache | None, num_tokens: int): swa_evictable_tokens=swa_evictable, ): return - - def minimum_reclaim(max_tokens: int, fits) -> int: - lo, hi = 0, max_tokens // page_size - while lo < hi: - mid = (lo + hi) // 2 - if fits(mid * page_size): - hi = mid - else: - lo = mid + 1 - return lo * page_size - - swa_reclaim = minimum_reclaim( - swa_evictable, - lambda value: allocator.can_reserve( - num_tokens, - num_tokens, - full_evictable_tokens=full_evictable, - swa_evictable_tokens=value, - ), - ) - full_reclaim = minimum_reclaim( - full_evictable, - lambda value: allocator.can_reserve( - num_tokens, - num_tokens, - full_evictable_tokens=value, - swa_evictable_tokens=swa_reclaim, - ), + full_min_tokens, swa_first_tokens, reclaim_bytes = allocator.reclaim_budget( + num_tokens, num_tokens ) - # FULL is the primary allocation. Reclaim SWA only when all evictable - # FULL plus compaction still cannot satisfy the shared byte demand. - evicted = tree_cache.evict_for_alloc( - EvictParams( - num_tokens=full_reclaim, - swa_num_tokens=swa_reclaim, + swa_first_tokens = min(swa_first_tokens, swa_evictable) + full_token_bytes = allocator.full_attn_allocator.entry_bytes + swa_token_bytes = allocator.swa_attn_allocator.entry_bytes + # A FULL victim may retain all, some, or none of its SWA pages. Measure + # allocator live deltas instead of assuming a fixed FULL:SWA ratio. + full_tokens_before = allocator.full_attn_allocator.allocated_count() + swa_tokens_before = allocator.swa_attn_allocator.allocated_count() + + from sglang.srt.mem_cache.unified_radix_cache import UnifiedRadixCache + + assert isinstance(tree_cache, UnifiedRadixCache) + + def enough() -> bool: + full_evicted = max( + 0, + full_tokens_before - allocator.full_attn_allocator.allocated_count(), + ) + swa_evicted = max( + 0, + swa_tokens_before - allocator.swa_attn_allocator.allocated_count(), + ) + if full_evicted < full_min_tokens: + return False + actual_reclaimed_bytes = ( + full_evicted * full_token_bytes + swa_evicted * swa_token_bytes ) + if actual_reclaimed_bytes < reclaim_bytes: + return False + return allocator.ensure_capacity(num_tokens, num_tokens) + + evicted = tree_cache.evict_swa_tombstones_then_full( + swa_num_tokens=swa_first_tokens, + full_num_tokens=full_evictable, + stop_full=enough, ) if not allocator.ensure_capacity(num_tokens, num_tokens): logger.warning( - "Unified SWA eviction did not satisfy the planned reservation: " + "Unified SWA eviction did not satisfy the reservation: " f"requested=({num_tokens}, {num_tokens}), " - f"planned=({full_reclaim}, {swa_reclaim}), " + f"budget=(full_min={full_min_tokens}, " + f"swa_first={swa_first_tokens}, bytes={reclaim_bytes}), " f"evicted=({evicted.num_tokens_evicted}, " f"{evicted.swa_num_tokens_evicted})" ) diff --git a/python/sglang/srt/mem_cache/multi_ended_allocator.py b/python/sglang/srt/mem_cache/multi_ended_allocator.py index 479babc6366e..7814828cdd93 100644 --- a/python/sglang/srt/mem_cache/multi_ended_allocator.py +++ b/python/sglang/srt/mem_cache/multi_ended_allocator.py @@ -3487,6 +3487,89 @@ def _compaction_allowed(self) -> bool: for allocator in (self.full_attn_allocator, self.swa_attn_allocator) ) + def reclaim_budget(self, full_tokens: int, swa_tokens: int) -> tuple[int, int, int]: + """Return ``(full_min_tokens, swa_first_tokens, shared_bytes)``. + + ``swa_first_tokens`` is the useful amount to request from a + tombstone-only SWA pass. ``full_min_tokens`` avoids repeated capacity + checks before a mandatory FULL-side reclaim is met. ``shared_bytes`` + is the byte shortfall which the actual eviction deltas must cover. + """ + page_size = self.page_size + full_pages = (max(0, int(full_tokens)) + page_size - 1) // page_size + swa_pages = (max(0, int(swa_tokens)) + page_size - 1) // page_size + fa, sa = self.full_attn_allocator, self.swa_attn_allocator + full_live_pages = fa.allocated_count() // page_size + swa_live_pages = sa.allocated_count() // page_size + + compacted = not self.lazy_compaction or self._compaction_allowed() + if compacted: + virtual_page_capacity = fa.num_virtual_ids - fa.min_page_index + full_page_capacity = min( + virtual_page_capacity, + fa.num_pages - fa.min_page_index, + ) + swa_page_capacity = min( + virtual_page_capacity, + sa.num_pages - sa.min_page_index, + ) + full_min_pages = max(0, full_live_pages + full_pages - full_page_capacity) + swa_min_pages = max(0, swa_live_pages + swa_pages - swa_page_capacity) + shared_bytes = max( + 0, + (full_live_pages + full_pages) * fa.entry_bytes_per_page + + (swa_live_pages + swa_pages) * sa.entry_bytes_per_page + - self._empty_shared_gap_bytes, + ) + shared_bytes_after_full = max( + 0, shared_bytes - full_min_pages * fa.entry_bytes_per_page + ) + swa_first_pages = max( + swa_min_pages, + -(-shared_bytes_after_full // sa.entry_bytes_per_page), + ) + else: + # With the PD move gate closed, only current holes and frontier gap + # are usable; pending pages and peer holes cannot be compacted. + full_holes = len(fa._free_phys_pages) + swa_holes = len(sa._free_phys_pages) + full_index_room = fa.num_pages - fa.min_page_index - fa._allocated_pages() + swa_index_room = sa.num_pages - sa.min_page_index - sa._allocated_pages() + full_min_pages = max( + 0, + full_pages - len(fa.free_virtual_ids), + full_pages - full_holes - full_index_room, + ) + swa_min_pages = max( + 0, + swa_pages - swa_holes - swa_index_room, + ) + full_extensions = max(0, full_pages - full_holes) + swa_extensions = max(0, swa_pages - swa_holes) + shared_bytes = max( + 0, + full_extensions * fa.entry_bytes_per_page + + swa_extensions * sa.entry_bytes_per_page + - fa._current_gap_bytes(), + ) + shared_bytes_after_full = max( + 0, + shared_bytes + - min(full_min_pages, full_extensions) * fa.entry_bytes_per_page, + ) + # Extra SWA holes cannot donate bytes to FULL while compaction is + # blocked, so stop once SWA's own extension demand is covered. + swa_first_pages = max( + swa_min_pages, + min( + swa_extensions, + -(-shared_bytes_after_full // sa.entry_bytes_per_page), + ), + ) + + swa_first_pages = min(swa_live_pages, swa_first_pages) + return full_min_pages * page_size, swa_first_pages * page_size, shared_bytes + def _fits_page_demand( self, num_full_pages: int, @@ -3524,11 +3607,7 @@ def _fits_page_demand( virtual_page_capacity, sa.num_pages - sa.min_page_index, ) - if ( - full_total_pages > full_page_capacity - or swa_total_pages > swa_page_capacity - or swa_total_pages > full_total_pages - ): + if full_total_pages > full_page_capacity or swa_total_pages > swa_page_capacity: return False if compacted: diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index 5cb53fac7ff2..5157ca9d7cdd 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -6,7 +6,15 @@ import time from dataclasses import replace from queue import Queue -from typing import TYPE_CHECKING, Iterator, NamedTuple, Optional, Sequence, TypeVar +from typing import ( + TYPE_CHECKING, + Callable, + Iterator, + NamedTuple, + Optional, + Sequence, + TypeVar, +) import torch @@ -624,6 +632,46 @@ def _evict( available_size_targets=available_size_targets, ) + return self._finish_eviction(tracker, start_time) + + def evict_swa_tombstones_then_full( + self, + *, + swa_num_tokens: int, + full_num_tokens: int, + stop_full: Callable[[], bool], + ) -> EvictResult: + """Evict SWA-only tombstones before falling back to FULL eviction. + + The SWA cursor tombstones independent candidates inline and returns + coupled leaves to its caller. Skipping those returned leaves makes the + first pass incapable of cascading into FULL. The second pass uses the + normal FULL eviction policy; ``stop_full`` observes allocator state + after each victim, including partially present SWA freed with a FULL node. + """ + if self.disable: + return EvictResult() + start_time = time.perf_counter() + tracker = {ct: 0 for ct in self.tree_components} + + if swa_num_tokens > 0: + self._evict_swa_tombstones(swa_num_tokens, tracker) + if not stop_full(): + self._evict_components( + self._evict_request_by_type(EvictParams(num_tokens=full_num_tokens)), + tracker, + stop=stop_full, + ) + + return self._finish_eviction(tracker, start_time) + + def _finish_eviction( + self, + tracker: dict[ComponentType, int], + start_time: float, + ) -> EvictResult: + """Finish write-back bookkeeping and build the public result.""" + if ( self.cache_controller is not None and self.cache_controller.write_policy == "write_back" @@ -698,12 +746,15 @@ def _evict_components( request_by_type: dict[ComponentType, int], tracker: dict[ComponentType, int], available_size_targets: Optional[dict[ComponentType, int]] = None, + stop: Optional[Callable[[], bool]] = 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 stop is not None and stop(): + return True if available_size_targets is None: return False target = available_size_targets.get(component_type) @@ -758,6 +809,26 @@ def target_reached(component_type: ComponentType) -> bool: finally: self.tree_core.evict_device_end(ct) + def _evict_swa_tombstones( + self, + request_cnt: int, + tracker: dict[ComponentType, int], + ) -> None: + """Evict SWA-only candidates, skipping leaves coupled to FULL.""" + self.tree_core.evict_device_start(ComponentType.SWA, request_cnt) + try: + while tracker[ComponentType.SWA] < request_cnt: + node_id, made_progress = self._evict_device_next_node( + ComponentType.SWA, tracker + ) + if node_id is not None: + # The cursor advanced but the coupled leaf is still intact. + continue + if not made_progress: + break + finally: + self.tree_core.evict_device_end(ComponentType.SWA) + def _record_dropped_tokens( self, tracker: dict[ComponentType, int], From e6e491bc3627c2f9ecf239ff196f1cee0ea4fc40 Mon Sep 17 00:00:00 2001 From: yhzhuang Date: Wed, 2 Sep 2026 13:27:43 -0700 Subject: [PATCH 20/55] Restore upstream unified cache eviction policy --- python/sglang/srt/mem_cache/common.py | 70 +--------------- .../srt/mem_cache/multi_ended_allocator.py | 83 ------------------- .../srt/mem_cache/unified_radix_cache.py | 64 +------------- 3 files changed, 2 insertions(+), 215 deletions(-) diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index 536a311be0b3..b1497c6f5c93 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/python/sglang/srt/mem_cache/common.py @@ -174,75 +174,7 @@ def evict_from_tree_cache(tree_cache: BasePrefixCache | None, num_tokens: int): allocator = tree_cache.token_to_kv_pool_allocator - from sglang.srt.mem_cache.multi_ended_allocator import ( - UnifiedSWATokenToKVPoolAllocator, - ) - - if ( - isinstance(allocator, UnifiedSWATokenToKVPoolAllocator) - and allocator.supports_asymmetric_reservation - ): - if allocator.ensure_capacity(num_tokens, num_tokens): - return - - page_size = allocator.page_size - full_evictable = tree_cache.full_evictable_size() // page_size * page_size - swa_evictable = tree_cache.swa_evictable_size() // page_size * page_size - if not allocator.can_reserve( - num_tokens, - num_tokens, - full_evictable_tokens=full_evictable, - swa_evictable_tokens=swa_evictable, - ): - return - full_min_tokens, swa_first_tokens, reclaim_bytes = allocator.reclaim_budget( - num_tokens, num_tokens - ) - swa_first_tokens = min(swa_first_tokens, swa_evictable) - full_token_bytes = allocator.full_attn_allocator.entry_bytes - swa_token_bytes = allocator.swa_attn_allocator.entry_bytes - # A FULL victim may retain all, some, or none of its SWA pages. Measure - # allocator live deltas instead of assuming a fixed FULL:SWA ratio. - full_tokens_before = allocator.full_attn_allocator.allocated_count() - swa_tokens_before = allocator.swa_attn_allocator.allocated_count() - - from sglang.srt.mem_cache.unified_radix_cache import UnifiedRadixCache - - assert isinstance(tree_cache, UnifiedRadixCache) - - def enough() -> bool: - full_evicted = max( - 0, - full_tokens_before - allocator.full_attn_allocator.allocated_count(), - ) - swa_evicted = max( - 0, - swa_tokens_before - allocator.swa_attn_allocator.allocated_count(), - ) - if full_evicted < full_min_tokens: - return False - actual_reclaimed_bytes = ( - full_evicted * full_token_bytes + swa_evicted * swa_token_bytes - ) - if actual_reclaimed_bytes < reclaim_bytes: - return False - return allocator.ensure_capacity(num_tokens, num_tokens) - - evicted = tree_cache.evict_swa_tombstones_then_full( - swa_num_tokens=swa_first_tokens, - full_num_tokens=full_evictable, - stop_full=enough, - ) - if not allocator.ensure_capacity(num_tokens, num_tokens): - logger.warning( - "Unified SWA eviction did not satisfy the reservation: " - f"requested=({num_tokens}, {num_tokens}), " - f"budget=(full_min={full_min_tokens}, " - f"swa_first={swa_first_tokens}, bytes={reclaim_bytes}), " - f"evicted=({evicted.num_tokens_evicted}, " - f"{evicted.swa_num_tokens_evicted})" - ) - elif isinstance(allocator, SWATokenToKVPoolAllocator): + if isinstance(allocator, SWATokenToKVPoolAllocator): # Hybrid allocator full_available_size = allocator.full_available_size() swa_available_size = allocator.swa_available_size() diff --git a/python/sglang/srt/mem_cache/multi_ended_allocator.py b/python/sglang/srt/mem_cache/multi_ended_allocator.py index 7814828cdd93..c3c1ae1ab871 100644 --- a/python/sglang/srt/mem_cache/multi_ended_allocator.py +++ b/python/sglang/srt/mem_cache/multi_ended_allocator.py @@ -3487,89 +3487,6 @@ def _compaction_allowed(self) -> bool: for allocator in (self.full_attn_allocator, self.swa_attn_allocator) ) - def reclaim_budget(self, full_tokens: int, swa_tokens: int) -> tuple[int, int, int]: - """Return ``(full_min_tokens, swa_first_tokens, shared_bytes)``. - - ``swa_first_tokens`` is the useful amount to request from a - tombstone-only SWA pass. ``full_min_tokens`` avoids repeated capacity - checks before a mandatory FULL-side reclaim is met. ``shared_bytes`` - is the byte shortfall which the actual eviction deltas must cover. - """ - page_size = self.page_size - full_pages = (max(0, int(full_tokens)) + page_size - 1) // page_size - swa_pages = (max(0, int(swa_tokens)) + page_size - 1) // page_size - fa, sa = self.full_attn_allocator, self.swa_attn_allocator - full_live_pages = fa.allocated_count() // page_size - swa_live_pages = sa.allocated_count() // page_size - - compacted = not self.lazy_compaction or self._compaction_allowed() - if compacted: - virtual_page_capacity = fa.num_virtual_ids - fa.min_page_index - full_page_capacity = min( - virtual_page_capacity, - fa.num_pages - fa.min_page_index, - ) - swa_page_capacity = min( - virtual_page_capacity, - sa.num_pages - sa.min_page_index, - ) - full_min_pages = max(0, full_live_pages + full_pages - full_page_capacity) - swa_min_pages = max(0, swa_live_pages + swa_pages - swa_page_capacity) - shared_bytes = max( - 0, - (full_live_pages + full_pages) * fa.entry_bytes_per_page - + (swa_live_pages + swa_pages) * sa.entry_bytes_per_page - - self._empty_shared_gap_bytes, - ) - shared_bytes_after_full = max( - 0, shared_bytes - full_min_pages * fa.entry_bytes_per_page - ) - swa_first_pages = max( - swa_min_pages, - -(-shared_bytes_after_full // sa.entry_bytes_per_page), - ) - else: - # With the PD move gate closed, only current holes and frontier gap - # are usable; pending pages and peer holes cannot be compacted. - full_holes = len(fa._free_phys_pages) - swa_holes = len(sa._free_phys_pages) - full_index_room = fa.num_pages - fa.min_page_index - fa._allocated_pages() - swa_index_room = sa.num_pages - sa.min_page_index - sa._allocated_pages() - full_min_pages = max( - 0, - full_pages - len(fa.free_virtual_ids), - full_pages - full_holes - full_index_room, - ) - swa_min_pages = max( - 0, - swa_pages - swa_holes - swa_index_room, - ) - full_extensions = max(0, full_pages - full_holes) - swa_extensions = max(0, swa_pages - swa_holes) - shared_bytes = max( - 0, - full_extensions * fa.entry_bytes_per_page - + swa_extensions * sa.entry_bytes_per_page - - fa._current_gap_bytes(), - ) - shared_bytes_after_full = max( - 0, - shared_bytes - - min(full_min_pages, full_extensions) * fa.entry_bytes_per_page, - ) - # Extra SWA holes cannot donate bytes to FULL while compaction is - # blocked, so stop once SWA's own extension demand is covered. - swa_first_pages = max( - swa_min_pages, - min( - swa_extensions, - -(-shared_bytes_after_full // sa.entry_bytes_per_page), - ), - ) - - swa_first_pages = min(swa_live_pages, swa_first_pages) - return full_min_pages * page_size, swa_first_pages * page_size, shared_bytes - def _fits_page_demand( self, num_full_pages: int, diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index 5157ca9d7cdd..9f7208de411d 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -6,15 +6,7 @@ import time from dataclasses import replace from queue import Queue -from typing import ( - TYPE_CHECKING, - Callable, - Iterator, - NamedTuple, - Optional, - Sequence, - TypeVar, -) +from typing import TYPE_CHECKING, Iterator, NamedTuple, Optional, Sequence, TypeVar import torch @@ -634,37 +626,6 @@ def _evict( return self._finish_eviction(tracker, start_time) - def evict_swa_tombstones_then_full( - self, - *, - swa_num_tokens: int, - full_num_tokens: int, - stop_full: Callable[[], bool], - ) -> EvictResult: - """Evict SWA-only tombstones before falling back to FULL eviction. - - The SWA cursor tombstones independent candidates inline and returns - coupled leaves to its caller. Skipping those returned leaves makes the - first pass incapable of cascading into FULL. The second pass uses the - normal FULL eviction policy; ``stop_full`` observes allocator state - after each victim, including partially present SWA freed with a FULL node. - """ - if self.disable: - return EvictResult() - start_time = time.perf_counter() - tracker = {ct: 0 for ct in self.tree_components} - - if swa_num_tokens > 0: - self._evict_swa_tombstones(swa_num_tokens, tracker) - if not stop_full(): - self._evict_components( - self._evict_request_by_type(EvictParams(num_tokens=full_num_tokens)), - tracker, - stop=stop_full, - ) - - return self._finish_eviction(tracker, start_time) - def _finish_eviction( self, tracker: dict[ComponentType, int], @@ -746,15 +707,12 @@ def _evict_components( request_by_type: dict[ComponentType, int], tracker: dict[ComponentType, int], available_size_targets: Optional[dict[ComponentType, int]] = None, - stop: Optional[Callable[[], bool]] = 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 stop is not None and stop(): - return True if available_size_targets is None: return False target = available_size_targets.get(component_type) @@ -809,26 +767,6 @@ def target_reached(component_type: ComponentType) -> bool: finally: self.tree_core.evict_device_end(ct) - def _evict_swa_tombstones( - self, - request_cnt: int, - tracker: dict[ComponentType, int], - ) -> None: - """Evict SWA-only candidates, skipping leaves coupled to FULL.""" - self.tree_core.evict_device_start(ComponentType.SWA, request_cnt) - try: - while tracker[ComponentType.SWA] < request_cnt: - node_id, made_progress = self._evict_device_next_node( - ComponentType.SWA, tracker - ) - if node_id is not None: - # The cursor advanced but the coupled leaf is still intact. - continue - if not made_progress: - break - finally: - self.tree_core.evict_device_end(ComponentType.SWA) - def _record_dropped_tokens( self, tracker: dict[ComponentType, int], From 66fed935c73cd434caae859fd9b05cc92c8411b2 Mon Sep 17 00:00:00 2001 From: yhzhuang Date: Wed, 2 Sep 2026 14:10:48 -0700 Subject: [PATCH 21/55] Use joint capacity accounting for unified SWA --- python/sglang/srt/mem_cache/common.py | 20 ++++-- .../srt/mem_cache/multi_ended_allocator.py | 67 +++++++++++++++---- 2 files changed, 70 insertions(+), 17 deletions(-) diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index b1497c6f5c93..d6acb09554f7 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/python/sglang/srt/mem_cache/common.py @@ -176,12 +176,22 @@ def evict_from_tree_cache(tree_cache: BasePrefixCache | None, num_tokens: int): if isinstance(allocator, SWATokenToKVPoolAllocator): # Hybrid allocator - full_available_size = allocator.full_available_size() - swa_available_size = allocator.swa_available_size() + from sglang.srt.mem_cache.multi_ended_allocator import ( + UnifiedSWATokenToKVPoolAllocator, + ) + + if ( + isinstance(allocator, UnifiedSWATokenToKVPoolAllocator) + and allocator.supports_asymmetric_reservation + ): + full_num_tokens, swa_num_tokens = allocator.allocation_shortfalls( + num_tokens, num_tokens + ) + else: + full_num_tokens = max(0, num_tokens - allocator.full_available_size()) + swa_num_tokens = max(0, num_tokens - allocator.swa_available_size()) - 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) + if full_num_tokens or swa_num_tokens: tree_cache.evict_for_alloc( EvictParams(num_tokens=full_num_tokens, swa_num_tokens=swa_num_tokens) ) diff --git a/python/sglang/srt/mem_cache/multi_ended_allocator.py b/python/sglang/srt/mem_cache/multi_ended_allocator.py index c3c1ae1ab871..f3e76b7c7e0e 100644 --- a/python/sglang/srt/mem_cache/multi_ended_allocator.py +++ b/python/sglang/srt/mem_cache/multi_ended_allocator.py @@ -3434,6 +3434,33 @@ def current_full_capacity(self) -> int: def current_swa_capacity(self) -> int: return self.swa_available_size() + self.swa_attn_allocator.allocated_count() + def allocation_shortfalls( + self, full_tokens: int | float, swa_tokens: int | float + ) -> tuple[int, int]: + """Return the shared-byte and SWA-slot shortfalls in token units. + + SWA bytes are charged to the FULL budget so the shared gap is counted + once. SWA-only eviction backs the slot budget, not the byte budget. + """ + if full_tokens < 0 or swa_tokens < 0: + raise ValueError("allocation sizes must be non-negative") + + page_size = self.page_size + full_pages = (math.ceil(full_tokens) + page_size - 1) // page_size + swa_pages = (math.ceil(swa_tokens) + page_size - 1) // page_size + full_page_bytes = self.full_attn_allocator.entry_bytes_per_page + swa_page_bytes = self.swa_attn_allocator.entry_bytes_per_page + swa_full_pages = ( + swa_pages * swa_page_bytes + full_page_bytes - 1 + ) // full_page_bytes + + joint_full_tokens = (full_pages + swa_full_pages) * page_size + aligned_swa_tokens = swa_pages * page_size + return ( + max(0, joint_full_tokens - self.full_available_size()), + max(0, aligned_swa_tokens - self.swa_available_size()), + ) + def can_reserve( self, full_tokens: int | float, @@ -3447,8 +3474,8 @@ def can_reserve( """Check pending FULL/SWA demand against the shared byte envelope. Scheduler admission keeps the historical one-token strict slack at an - empty-pool boundary. Actual allocation and eviction planning may consume - the exact physical capacity, so they leave ``require_token_slack`` off. + empty-pool boundary. For a live pool, shared bytes are backed only by + FULL eviction; SWA eviction backs its own slot shortfall. """ if not self.supports_asymmetric_reservation: if ( @@ -3466,19 +3493,35 @@ def can_reserve( ): return False + if empty_pool: + page_size = self.page_size + full_pages = (math.ceil(full_tokens) + page_size - 1) // page_size + swa_pages = (math.ceil(swa_tokens) + page_size - 1) // page_size + return self._fits_page_demand( + full_pages, + swa_pages, + compacted=True, + empty_pool=True, + ) + + full_shortfall, swa_shortfall = self.allocation_shortfalls( + full_tokens, swa_tokens + ) page_size = self.page_size - full_pages = (math.ceil(full_tokens) + page_size - 1) // page_size - swa_pages = (math.ceil(swa_tokens) + page_size - 1) // page_size - compacted = empty_pool or not self.lazy_compaction - if not compacted: - compacted = self._compaction_allowed() + if ( + full_shortfall > max(0, int(full_evictable_tokens)) // page_size * page_size + or swa_shortfall + > max(0, int(swa_evictable_tokens)) // page_size * page_size + ): + return False + + compacted = not self.lazy_compaction or self._compaction_allowed() return self._fits_page_demand( - full_pages, - swa_pages, - full_reclaim_pages=max(0, int(full_evictable_tokens)) // page_size, - swa_reclaim_pages=max(0, int(swa_evictable_tokens)) // page_size, + (math.ceil(full_tokens) + page_size - 1) // page_size, + (math.ceil(swa_tokens) + page_size - 1) // page_size, + full_reclaim_pages=full_shortfall // page_size, + swa_reclaim_pages=swa_shortfall // page_size, compacted=compacted, - empty_pool=empty_pool, ) def _compaction_allowed(self) -> bool: From c41b06843cc1368daafac286188b6821abe1c4a2 Mon Sep 17 00:00:00 2001 From: yhzhuang Date: Fri, 4 Sep 2026 11:28:48 -0700 Subject: [PATCH 22/55] Add unified capacity regression coverage --- python/sglang/srt/mem_cache/allocator/swa.py | 10 -- .../mem_cache/test_multi_ended_allocator.py | 96 +++++++++++++++++++ .../model_executor/test_pool_configurator.py | 62 ++++++++++++ 3 files changed, 158 insertions(+), 10 deletions(-) diff --git a/python/sglang/srt/mem_cache/allocator/swa.py b/python/sglang/srt/mem_cache/allocator/swa.py index d2c8bed5eb2c..785f9b83a30c 100644 --- a/python/sglang/srt/mem_cache/allocator/swa.py +++ b/python/sglang/srt/mem_cache/allocator/swa.py @@ -343,16 +343,6 @@ def set_full_to_swa_mapping( swa_indices = swa_indices.to(self.full_to_swa_index_mapping.dtype) self.full_to_swa_index_mapping[full_indices] = swa_indices - def recover_swa_with_locked_full( - self, kept_full: torch.Tensor, incoming_full: torch.Tensor - ) -> torch.Tensor: - """Keep locked FULL slots and transfer the incoming SWA slots to them.""" - swa_value = self.translate_loc_from_full_to_swa(incoming_full) - self.set_full_to_swa_mapping(kept_full, swa_value) - self.clear_full_to_swa_mapping(incoming_full) - self.full_attn_allocator.free(incoming_full) - return swa_value - def clear_full_to_swa_mapping(self, full_indices: torch.Tensor) -> None: if full_indices.numel() == 0: return diff --git a/test/registered/unit/mem_cache/test_multi_ended_allocator.py b/test/registered/unit/mem_cache/test_multi_ended_allocator.py index e24050807bd1..a9c364fe9e26 100644 --- a/test/registered/unit/mem_cache/test_multi_ended_allocator.py +++ b/test/registered/unit/mem_cache/test_multi_ended_allocator.py @@ -30,6 +30,7 @@ import torch +from sglang.srt.managers.schedule_policy import estimate_swa_kv_tokens from sglang.srt.mem_cache.multi_ended_allocator import ( FloatMultiEndedAllocator, MultiEndedAllocator, @@ -620,6 +621,7 @@ def _build( swa_layer_num=2, head_num=2, head_dim=4, + page_size=1, ): full_spec = MHASubPoolSpec( name="full", @@ -646,6 +648,7 @@ def _build( sub_pool_specs=[full_spec, swa_spec], device=_DEV, enable_memory_saver=False, + page_size=page_size, ) kvcache = _FakeUnifiedSWAKVPool(pool) allocator = UnifiedSWATokenToKVPoolAllocator( @@ -654,11 +657,104 @@ def _build( device=_DEV, full_max_total_num_tokens=n_full_slots, swa_max_total_num_tokens=n_swa_slots, + page_size=page_size, need_sort=False, forward_stream=None, ) return pool, allocator, kvcache + def test_empty_pool_reservation_matches_packed_byte_boundary(self): + page_size = 4 + _, allocator, _ = self._build( + n_full_slots=40, + n_swa_slots=24, + full_layer_num=4, + swa_layer_num=2, + page_size=page_size, + ) + full_allocator = allocator.full_attn_allocator + swa_allocator = allocator.swa_attn_allocator + swa_pages = 2 + full_pages = ( + allocator._empty_shared_gap_bytes + - swa_pages * swa_allocator.entry_bytes_per_page + ) // full_allocator.entry_bytes_per_page + packed_bytes = ( + full_pages * full_allocator.entry_bytes_per_page + + swa_pages * swa_allocator.entry_bytes_per_page + ) + + self.assertLessEqual( + full_pages + 1, full_allocator.num_pages - full_allocator.min_page_index + ) + self.assertLessEqual( + swa_pages, swa_allocator.num_pages - swa_allocator.min_page_index + ) + self.assertLessEqual(packed_bytes, allocator._empty_shared_gap_bytes) + self.assertGreater( + packed_bytes + full_allocator.entry_bytes_per_page, + allocator._empty_shared_gap_bytes, + ) + self.assertTrue( + allocator.can_reserve( + full_pages * page_size, + swa_pages * page_size, + empty_pool=True, + ) + ) + self.assertFalse( + allocator.can_reserve( + full_pages * page_size + 1, + swa_pages * page_size, + empty_pool=True, + ) + ) + + extend_tokens = 32 + max_new_tokens = 0 + reservation_full_tokens = extend_tokens + max_new_tokens + page_size + reservation_swa_tokens = estimate_swa_kv_tokens( + extend_tokens, + max_new_tokens, + sliding_window_size=16, + page_size=page_size, + include_allocated_tail=False, + ) + reservation_swa_with_tail = estimate_swa_kv_tokens( + extend_tokens, + max_new_tokens, + sliding_window_size=16, + page_size=page_size, + ) + reservation_bytes = ( + reservation_full_tokens // page_size + ) * full_allocator.entry_bytes_per_page + ( + reservation_swa_tokens // page_size + ) * swa_allocator.entry_bytes_per_page + reservation_bytes_with_tail = ( + reservation_full_tokens // page_size + ) * full_allocator.entry_bytes_per_page + ( + reservation_swa_with_tail // page_size + ) * swa_allocator.entry_bytes_per_page + self.assertLessEqual(reservation_bytes, allocator._empty_shared_gap_bytes) + self.assertGreater( + reservation_bytes_with_tail, allocator._empty_shared_gap_bytes + ) + self.assertTrue( + allocator.can_reserve( + reservation_full_tokens, + reservation_swa_tokens, + empty_pool=True, + ) + ) + self.assertFalse( + allocator.can_reserve( + reservation_full_tokens, + reservation_swa_with_tail, + empty_pool=True, + ) + ) + def _alloc(self, allocator, kvcache, n): """Allocate N virtual ids; stamp the data marker on both sub-pools.""" v = allocator.alloc(n) diff --git a/test/registered/unit/model_executor/test_pool_configurator.py b/test/registered/unit/model_executor/test_pool_configurator.py index 3e8cf29b283c..1883d5c68e1f 100644 --- a/test/registered/unit/model_executor/test_pool_configurator.py +++ b/test/registered/unit/model_executor/test_pool_configurator.py @@ -337,6 +337,68 @@ def test_memory_utilization(self): self.assertLessEqual(used, available) self.assertGreater(used, available * 0.99) + def test_unified_capacity_is_maximal_with_draft_pool(self): + page_size = 8 + full_layers = 2 + swa_layers = 1 + draft_layers = 2 + draft_swa_layers = 1 + ratio = 0.5 + mr = _make_model_runner( + self, + is_hybrid_swa=True, + full_attention_layer_ids=list(range(full_layers)), + swa_attention_layer_ids=list(range(full_layers, full_layers + swa_layers)), + swa_num_kv_heads=4, + swa_full_tokens_ratio=ratio, + page_size=page_size, + enable_unified_memory=True, + speculative_algorithm="EAGLE", + ) + mr.spec_algorithm.is_eagle.return_value = True + mr.spec_algorithm.is_none.return_value = False + mr.spec_aux_config.eagle_draft_num_layers = draft_layers + mr.spec_aux_config.eagle_draft_swa_num_layers = draft_swa_layers + + full_bytes_per_token = _full_per_token(mr) + swa_bytes_per_token = _swa_per_token(mr) + target_full_bytes_per_token = full_bytes_per_token * full_layers + draft_bytes_per_token = ( + full_bytes_per_token * (draft_layers - draft_swa_layers) + + swa_bytes_per_token * draft_swa_layers + ) + + def allocation_bytes(full_tokens, *, include_reserved_draft_page=True): + swa_tokens = int(full_tokens * ratio) // page_size * page_size + target_bytes = ( + full_tokens * target_full_bytes_per_token + + swa_tokens * swa_bytes_per_token * swa_layers + ) + virtual_span = max(target_bytes // target_full_bytes_per_token - 1, 0) + draft_tokens = (virtual_span + page_size - 1) // page_size * page_size + if include_reserved_draft_page: + draft_tokens += page_size + return target_bytes + draft_tokens * draft_bytes_per_token + + expected_full_tokens = 10 * page_size + available = allocation_bytes( + expected_full_tokens + page_size, + include_reserved_draft_page=False, + ) + with mock_cpu_env(): + from sglang.srt.model_executor.pool_configurator import ( + create_memory_pool_configurator, + ) + + cfg = create_memory_pool_configurator(mr) + config = cfg.calculate_pool_sizes(available, page_size) + + full_tokens = config.full_max_total_num_tokens + self.assertEqual(full_tokens % page_size, 0) + self.assertEqual(full_tokens, expected_full_tokens) + self.assertLessEqual(allocation_bytes(full_tokens), available) + self.assertGreater(allocation_bytes(full_tokens + page_size), available) + @patch( "sglang.srt.mem_cache.kv_cache_configurator.calculate_mla_kv_cache_dim", return_value=576, From 43ecd0688544acea1528069a83c1aeb3b20485e7 Mon Sep 17 00:00:00 2001 From: yhzhuang Date: Fri, 4 Sep 2026 11:29:04 -0700 Subject: [PATCH 23/55] Use physical SWA indices for PD transfer --- python/sglang/srt/disaggregation/decode.py | 2 +- .../srt/disaggregation/mooncake/conn.py | 7 ++++++ python/sglang/srt/disaggregation/prefill.py | 2 +- .../srt/mem_cache/allocator/hisparse.py | 7 ++++++ python/sglang/srt/mem_cache/allocator/swa.py | 6 +++++ .../srt/mem_cache/multi_ended_allocator.py | 6 +++++ .../test_pp_hybrid_kv_transfer.py | 20 ++++++++++++++++ .../mem_cache/test_multi_ended_allocator.py | 24 +++++++++++++++++++ 8 files changed, 72 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index e34679e08fc6..94c106796ca4 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -1540,7 +1540,7 @@ def _swa_payload(): decode_req.req.kv.req_pool_idx, window_start:seq_len ] window_kv_indices_swa = ( - self.token_to_kv_pool_allocator.translate_loc_from_full_to_swa( + self.token_to_kv_pool_allocator.translate_swa_indices_for_transfer( window_kv_indices_full ) ) diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py index afcc0ac4799c..2f5a8e0e6c5e 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/python/sglang/srt/disaggregation/mooncake/conn.py @@ -685,11 +685,18 @@ def _send_kvcache_generic( # Published layer IDs give exact pairing; plain-MHA peers publish none # and keep positional slicing. has_layer_ids = bool(src_layer_ids or dst_layer_ids) + # Unified SWA publishes one page-envelope region even on an MHA backend. + is_single_region_swa = ( + state_type == StateType.SWA + and len(src_data_ptrs) == 1 + and len(dst_data_ptrs) == 1 + ) if ( self.is_mla_backend or self.is_hybrid_mla_backend or force_flat or has_layer_ids + or is_single_region_swa ): # Layer IDs map PP-local buffers to global decode entries. # Registrations without them retain the existing PP mapping. diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index 8d08befecf0a..0266dd93ebad 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -1230,7 +1230,7 @@ def _swa_payload(): req.kv.req_pool_idx, window_start:seq_len ] window_kv_indices_swa = ( - self.token_to_kv_pool_allocator.translate_loc_from_full_to_swa( + self.token_to_kv_pool_allocator.translate_swa_indices_for_transfer( window_kv_indices_full ) ) diff --git a/python/sglang/srt/mem_cache/allocator/hisparse.py b/python/sglang/srt/mem_cache/allocator/hisparse.py index 4a10891f8514..bce9434cf5ce 100644 --- a/python/sglang/srt/mem_cache/allocator/hisparse.py +++ b/python/sglang/srt/mem_cache/allocator/hisparse.py @@ -353,6 +353,13 @@ def get_kvcache(self): def translate_loc_from_full_to_swa(self, kv_indices: torch.Tensor): return self.logical_attn_allocator.translate_loc_from_full_to_swa(kv_indices) + def translate_swa_indices_for_transfer( + self, kv_indices: torch.Tensor + ) -> torch.Tensor: + return self.logical_attn_allocator.translate_swa_indices_for_transfer( + kv_indices + ) + def full_available_size(self): return min( self.logical_attn_allocator.full_available_size(), diff --git a/python/sglang/srt/mem_cache/allocator/swa.py b/python/sglang/srt/mem_cache/allocator/swa.py index d2c8bed5eb2c..2945e7e9bfc0 100644 --- a/python/sglang/srt/mem_cache/allocator/swa.py +++ b/python/sglang/srt/mem_cache/allocator/swa.py @@ -149,6 +149,12 @@ def translate_loc_from_full_to_swa(self, kv_indices: torch.Tensor): assert self._kvcache.full_to_swa_index_mapping is not None return self._kvcache.translate_loc_from_full_to_swa(kv_indices) + def translate_swa_indices_for_transfer( + self, kv_indices: torch.Tensor + ) -> torch.Tensor: + """Map full-pool token ids to SWA-buffer token ids for PD transfer.""" + return self.translate_loc_from_full_to_swa(kv_indices) + def alloc(self, need_size: int): assert self.page_size == 1 if need_size > self.full_attn_allocator.available_size(): diff --git a/python/sglang/srt/mem_cache/multi_ended_allocator.py b/python/sglang/srt/mem_cache/multi_ended_allocator.py index dd84f10f1989..9852f6827caf 100644 --- a/python/sglang/srt/mem_cache/multi_ended_allocator.py +++ b/python/sglang/srt/mem_cache/multi_ended_allocator.py @@ -3670,6 +3670,12 @@ def translate_kv_indices_for_transfer( """Translate virtual token ids to full-pool physical ids for PD.""" return self.full_attn_allocator.translate_kv_loc(kv_indices.to(torch.int64)) + def translate_swa_indices_for_transfer( + self, kv_indices: torch.Tensor + ) -> torch.Tensor: + """Virtual token ids -> physical, not kernel-facing, SWA ids for PD.""" + return self.swa_attn_allocator.translate_kv_loc(kv_indices.to(torch.int64)) + def set_disagg_move_gate(self, gate: Callable[[], bool]) -> None: """Install the PD compaction gate on both physical sub-allocators.""" assert self.lazy_compaction, ( diff --git a/test/registered/unit/disaggregation/test_pp_hybrid_kv_transfer.py b/test/registered/unit/disaggregation/test_pp_hybrid_kv_transfer.py index e3a8bc7f4cdb..bc7007d026f2 100644 --- a/test/registered/unit/disaggregation/test_pp_hybrid_kv_transfer.py +++ b/test/registered/unit/disaggregation/test_pp_hybrid_kv_transfer.py @@ -6,6 +6,7 @@ import numpy as np +from sglang.srt.disaggregation.base.conn import StateType from sglang.srt.disaggregation.common.conn import CommonKVManager from sglang.srt.disaggregation.mooncake.conn import MooncakeKVManager from sglang.srt.disaggregation.prefill import _transfer_start_layer @@ -70,6 +71,7 @@ def test_non_hybrid_pool_keeps_global_start_layer(self): class _RecordingKVManager: get_mha_kv_ptrs_with_pp = CommonKVManager.get_mha_kv_ptrs_with_pp + get_mla_kv_ptrs_with_pp = CommonKVManager.get_mla_kv_ptrs_with_pp def __init__(self, *, prefill_start_layer: int, pp_size: int): self.is_mla_backend = False @@ -141,6 +143,24 @@ def test_f5_of_n12(self): self._run_case(model_full_ids=ids, stage_full_ids=ids[:5], start_offset=0) +class TestSingleRegionSWATransfer(CustomTestCase): + def test_one_region_swa_generates_transfer_block(self): + manager = _RecordingKVManager(prefill_start_layer=0, pp_size=1) + rc = MooncakeKVManager._send_kvcache_generic( + manager, + mooncake_session_id="session", + src_data_ptrs=[1000], + dst_data_ptrs=[2000], + item_lens=[64], + prefill_data_indices=np.array([3, 4], dtype=np.int32), + dst_data_indices=np.array([7, 8], dtype=np.int32), + executor=None, + state_type=StateType.SWA, + ) + self.assertEqual(rc, 0) + self.assertEqual(manager.blocks, [(1000 + 3 * 64, 2000 + 7 * 64, 2 * 64)]) + + class TestGetMhaKvPtrsWithPp(CustomTestCase): """Derived property: the modulo heuristic in get_mha_kv_ptrs_with_pp exists for the decode-has-draft-KV layout [K_main, V_main, draft_K, draft_V]. Pin diff --git a/test/registered/unit/mem_cache/test_multi_ended_allocator.py b/test/registered/unit/mem_cache/test_multi_ended_allocator.py index e24050807bd1..9ac83b912d34 100644 --- a/test/registered/unit/mem_cache/test_multi_ended_allocator.py +++ b/test/registered/unit/mem_cache/test_multi_ended_allocator.py @@ -30,6 +30,7 @@ import torch +from sglang.srt.mem_cache.common import kv_to_page_indices from sglang.srt.mem_cache.multi_ended_allocator import ( FloatMultiEndedAllocator, MultiEndedAllocator, @@ -2735,6 +2736,29 @@ def test_swa_translate_scales_page_stride(self): expected = v2p_swa[v // self.PS] * (self.PS * mult) + v % self.PS self.assertTrue(torch.equal(a.translate_loc_from_full_to_swa(v), expected)) + def test_swa_transfer_page_is_physical_not_kernel_scaled(self): + mult = 2 * self.SWA_L + a = self._build() + v = a.alloc(3 * self.PS) + self.assertIsNotNone(v) + + physical_pages = a.swa_attn_allocator.virtual_to_physical[ + v[:: self.PS] // self.PS + ] + physical_tokens = a.swa_attn_allocator.translate_kv_loc(v) + transfer_tokens = a.translate_swa_indices_for_transfer(v) + self.assertTrue(torch.equal(transfer_tokens, physical_tokens)) + self.assertEqual( + kv_to_page_indices(transfer_tokens, self.PS).tolist(), + physical_pages.tolist(), + ) + + kernel_tokens = a.translate_loc_from_full_to_swa(v) + self.assertEqual( + kv_to_page_indices(kernel_tokens, self.PS).tolist(), + (physical_pages * mult).tolist(), + ) + def test_swa_kernel_tombstone_still_lands_on_sink(self): """The scaled stride must not break the tombstone clamp: a tombstoned page's ids (v2p == -1 -> -stride + offset, negative for every in-page From 940d41b2928bbd934b57e76a3fe8baacd8a3df61 Mon Sep 17 00:00:00 2001 From: yhzhuang Date: Fri, 4 Sep 2026 13:56:56 -0700 Subject: [PATCH 24/55] Update PD priority test fixture for draft indices --- .../unit/managers/test_priority_scheduling_disaggregation.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/test/registered/unit/managers/test_priority_scheduling_disaggregation.py b/test/registered/unit/managers/test_priority_scheduling_disaggregation.py index 397e27666d9c..fac4f433e561 100644 --- a/test/registered/unit/managers/test_priority_scheduling_disaggregation.py +++ b/test/registered/unit/managers/test_priority_scheduling_disaggregation.py @@ -147,7 +147,10 @@ def pre_alloc_mock(req, prefix_indices=None, prefix_len=0, total_prefix_len=0): queue.token_to_kv_pool_allocator.available_size.return_value = 1000 queue.token_to_kv_pool = MagicMock() queue.transfer_queue = SimpleNamespace(queue=[], enable_staging=False) - queue.kv_manager = SimpleNamespace(kv_args=SimpleNamespace(state_types=[])) + queue.kv_manager = SimpleNamespace( + kv_args=SimpleNamespace(state_types=[]), + uses_separate_draft_kv_indices=False, + ) queue.tree_cache = MagicMock() scheduler = MagicMock() From d01513db8edf88c4caf5c1c974138fda13050f7b Mon Sep 17 00:00:00 2001 From: yhzhuang Date: Fri, 4 Sep 2026 13:56:56 -0700 Subject: [PATCH 25/55] Update HiSparse test fixture for unified memory --- .../unit/mem_cache/test_hisparse_max_token_pool_size.py | 1 + 1 file changed, 1 insertion(+) diff --git a/test/registered/unit/mem_cache/test_hisparse_max_token_pool_size.py b/test/registered/unit/mem_cache/test_hisparse_max_token_pool_size.py index eb76a1e56ad1..af2d35182b9a 100644 --- a/test/registered/unit/mem_cache/test_hisparse_max_token_pool_size.py +++ b/test/registered/unit/mem_cache/test_hisparse_max_token_pool_size.py @@ -79,6 +79,7 @@ def test_non_hisparse_uses_effective_max_total_num_tokens(self): def test_non_hisparse_hybrid_swa_prefers_full_max(self): instance = _make_model_runner( enable_hisparse=False, + server_args=SimpleNamespace(enable_unified_memory=False), token_to_kv_pool_allocator=SimpleNamespace(), is_hybrid_swa=True, max_total_num_tokens=1024, From 8bc15863561c3dfe11075232f666ca2f82e998be Mon Sep 17 00:00:00 2001 From: yhzhuang Date: Tue, 8 Sep 2026 14:38:27 -0700 Subject: [PATCH 26/55] Plan unified FULL/SWA reclaim against physical capacity Account for reclaimable pages on both sides of the arena and realize compaction before allocation. Keep tri-pool conservation statistics and reject unsupported mixed-chunk and pipeline-parallel layouts. --- python/sglang/srt/arg_groups/kv_cache_hook.py | 11 +- .../pool_stats_observer.py | 13 +- python/sglang/srt/mem_cache/common.py | 43 +++-- .../srt/mem_cache/multi_ended_allocator.py | 163 ++++++++++-------- .../srt/model_executor/pool_configurator.py | 21 +-- 5 files changed, 149 insertions(+), 102 deletions(-) diff --git a/python/sglang/srt/arg_groups/kv_cache_hook.py b/python/sglang/srt/arg_groups/kv_cache_hook.py index 26f788c14052..d01814ea8a90 100644 --- a/python/sglang/srt/arg_groups/kv_cache_hook.py +++ b/python/sglang/srt/arg_groups/kv_cache_hook.py @@ -198,11 +198,14 @@ def handle_unified_memory_pool(server_args: Any) -> None: cfg = resolving_view(server_args) if not cfg.enable_unified_memory: return + assert cfg.pp_size == 1, ( + "--enable-unified-memory does not support " + "pipeline parallelism (--pp-size > 1)." + ) + assert ( + not cfg.enable_mixed_chunk + ), "--enable-unified-memory does not support --enable-mixed-chunk." if cfg.disaggregation_mode != "null": - assert cfg.pp_size == 1, ( - "--enable-unified-memory with PD disaggregation does not support " - "pipeline parallelism (--pp-size > 1)." - ) # Constraints of the whole-envelope transfer; see # UnifiedMLATokenToKVPool.get_contiguous_buf_infos. assert cfg.disaggregation_transfer_backend == "mooncake", ( diff --git a/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py b/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py index 1d6fc836bbee..a3f66e16fb0c 100644 --- a/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py +++ b/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py @@ -11,6 +11,10 @@ Tuple, ) +from sglang.srt.mem_cache.multi_ended_allocator import ( + UnifiedMambaSWATokenToKVPoolAllocator, +) + if TYPE_CHECKING: from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache @@ -287,8 +291,13 @@ def _get_swa_token_info(self) -> PoolStats: allocator = self.token_to_kv_pool_allocator full_capacity = allocator.current_full_capacity swa_capacity = allocator.current_swa_capacity - full_available_size = allocator.full_available_size() - swa_available_size = allocator.swa_available_size() + if isinstance(allocator, UnifiedMambaSWATokenToKVPoolAllocator): + # The tri-pool reports static capacities paired with conserve views. + full_available_size = allocator.conserve_full_available_size() + swa_available_size = allocator.conserve_swa_available_size() + else: + full_available_size = allocator.full_available_size() + swa_available_size = allocator.swa_available_size() full_evictable_size = self.tree_cache.full_evictable_size() swa_evictable_size = self.tree_cache.swa_evictable_size() full_num_used = full_capacity - (full_available_size + full_evictable_size) diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index e02d22ec901d..ebf437dc1b2a 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/python/sglang/srt/mem_cache/common.py @@ -66,9 +66,9 @@ def free_swa_out_of_window_slots( return # 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.kv.cache_protected_len % page_size == 0, ( - "cache_protected_len must be page aligned" - ) + assert ( + req.kv.cache_protected_len % page_size == 0 + ), "cache_protected_len must be page aligned" req.kv.swa_evicted_seqlen = max( req.kv.swa_evicted_seqlen, req.kv.swa_dead_lo(page_size) ) @@ -176,13 +176,20 @@ def evict_from_tree_cache(tree_cache: BasePrefixCache | None, num_tokens: int): UnifiedSWATokenToKVPoolAllocator, ) - if ( + is_unified_allocator = ( isinstance(allocator, UnifiedSWATokenToKVPoolAllocator) and allocator.supports_asymmetric_reservation - ): - full_num_tokens, swa_num_tokens = allocator.allocation_shortfalls( - num_tokens, num_tokens + ) + if is_unified_allocator: + reclaim_plan = allocator.reclaim_plan( + num_tokens, + num_tokens, + full_evictable_tokens=tree_cache.full_evictable_size(), + swa_evictable_tokens=tree_cache.swa_evictable_size(), ) + if reclaim_plan is None: + return + full_num_tokens, swa_num_tokens = reclaim_plan else: full_num_tokens = max(0, num_tokens - allocator.full_available_size()) swa_num_tokens = max(0, num_tokens - allocator.swa_available_size()) @@ -191,6 +198,10 @@ def evict_from_tree_cache(tree_cache: BasePrefixCache | None, num_tokens: int): tree_cache.evict_for_alloc( EvictParams(num_tokens=full_num_tokens, swa_num_tokens=swa_num_tokens) ) + if is_unified_allocator: + # The reclaim plan may rely on compacting cross-side holes. Make + # that projected state real before callers inspect availability. + allocator.ensure_capacity(num_tokens, num_tokens) else: # Standard allocator: evict only the shortfall (mirrors the SWA arm) available_size = allocator.available_size() @@ -261,9 +272,9 @@ def release_kv_cache(req: Req, tree_cache: BasePrefixCache, is_insert: bool = Tr assert (not req.kv.holds_kv) == req.kv.is_kv_released # MambaRadixCache may alloc mamba state before alloc KV cache if not req.kv.holds_kv: - assert tree_cache.supports_mamba(), ( - "Only MambaRadixCache allow freeing before alloc" - ) + assert ( + tree_cache.supports_mamba() + ), "Only MambaRadixCache allow freeing before alloc" # TODO (csy, hanming): clean up this early allocation logic if req.kv.holds_mamba: tree_cache.req_to_token_pool.mamba_allocator.free( @@ -292,9 +303,9 @@ def release_kv_cache(req: Req, tree_cache: BasePrefixCache, is_insert: bool = Tr if isinstance(tree_cache.req_to_token_pool, HybridReqToTokenPool) and ( not tree_cache.supports_mamba() ): - assert req.kv.holds_mamba, ( - "mamba state is freed while the tree cache does not manage mamba states" - ) + assert ( + req.kv.holds_mamba + ), "mamba state is freed while the tree cache does not manage mamba states" tree_cache.req_to_token_pool.free_mamba_cache(req) # The DSV4-NPU ReqToTokenPool subclass's free() additionally releases the # c4/c128 state pages; other ReqToTokenPool subclasses are a no-op here. @@ -312,9 +323,9 @@ def _release_overallocated_kv_indices( # strip_thinking_cache intentionally reports output tokens as overallocated # so they fall into the free path below (#22373). if spec_algo is None and not get_serving().strip_thinking_cache: - assert start_p == end_p, ( - f"Unexpected overallocated KV cache, {req.kv.kv_committed_len=}, {req.kv.kv_allocated_len=}" - ) + assert ( + start_p == end_p + ), f"Unexpected overallocated KV cache, {req.kv.kv_committed_len=}, {req.kv.kv_allocated_len=}" if page_size > 1: start_p = ceil_align(start_p, page_size) diff --git a/python/sglang/srt/mem_cache/multi_ended_allocator.py b/python/sglang/srt/mem_cache/multi_ended_allocator.py index e9740fee8063..1419bb7af136 100644 --- a/python/sglang/srt/mem_cache/multi_ended_allocator.py +++ b/python/sglang/srt/mem_cache/multi_ended_allocator.py @@ -329,9 +329,9 @@ def __init__( self.num_pages = max_slots // self.pool_page_size # `min_page_index` = ceil(min_slot_index / pool_page_size), keeping the # reserved-sink invariant (min_page_index * entry_bytes_per_page >= entry_max). - assert virtual_num_pages is None or not is_id_owner, ( - "only a non-owner allocator may use another pool's virtual-id space" - ) + assert ( + virtual_num_pages is None or not is_id_owner + ), "only a non-owner allocator may use another pool's virtual-id space" # An ID owner has one virtual page per physical page. A non-owner may # deliberately have a different physical capacity while indexing the # owner's virtual page space (the unified SWA allocator does this). @@ -369,9 +369,9 @@ def __init__( self.high_peer: Optional[MultiEndedAllocator] = None # Inverse history of relocations (spec rollback), at PAGE granularity. - self._inverse_history: List[ - Tuple[torch.Tensor, torch.Tensor, torch.Tensor] - ] = [] + self._inverse_history: List[Tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = ( + [] + ) # --- Lazy compaction state (all unused when lazy_compaction=False) --- # `_free_phys_pages`: GPU free list of physical PAGE ids, sorted at `_flush`. @@ -1233,9 +1233,9 @@ def alloc_extend( v2p stays -1 and translation yields negative ids → CUDA OOB. """ with record_function("MultiEndedAlloc.alloc_extend"): - assert self.is_id_owner, ( - f"alloc_extend on a non-id-owner allocator ({self.sub_pool_name!r})" - ) + assert ( + self.is_id_owner + ), f"alloc_extend on a non-id-owner allocator ({self.sub_pool_name!r})" if num_new_pages is None: num_new_pages = get_num_new_pages( seq_lens=seq_lens_cpu, @@ -1302,9 +1302,9 @@ def alloc_decode( virtual page on THIS sub-allocator (else v2p stays -1 → CUDA OOB). """ with record_function("MultiEndedAlloc.alloc_decode"): - assert self.is_id_owner, ( - f"alloc_decode on a non-id-owner allocator ({self.sub_pool_name!r})" - ) + assert ( + self.is_id_owner + ), f"alloc_decode on a non-id-owner allocator ({self.sub_pool_name!r})" bs = len(seq_lens) # CPU-only count BEFORE the kernel, to snapshot the exact slice the # kernel will consume. @@ -3429,31 +3429,77 @@ def current_full_capacity(self) -> int: def current_swa_capacity(self) -> int: return self.swa_available_size() + self.swa_attn_allocator.allocated_count() - def allocation_shortfalls( - self, full_tokens: int | float, swa_tokens: int | float - ) -> tuple[int, int]: - """Return the shared-byte and SWA-slot shortfalls in token units. + def reclaim_plan( + self, + full_tokens: int | float, + swa_tokens: int | float, + *, + full_evictable_tokens: int = 0, + swa_evictable_tokens: int = 0, + empty_pool: bool = False, + ) -> Optional[Tuple[int, int]]: + """Return cumulative FULL/SWA eviction targets, or None if impossible. - SWA bytes are charged to the FULL budget so the shared gap is counted - once. SWA-only eviction backs the slot budget, not the byte budget. + The tree evicts FULL before SWA. Minimize required SWA reclaim with all + evictable FULL available, then trim excess FULL reclaim. FULL eviction's + actual SWA cascade is counted by the tree's shared eviction tracker. """ if full_tokens < 0 or swa_tokens < 0: - raise ValueError("allocation sizes must be non-negative") + return None page_size = self.page_size full_pages = (math.ceil(full_tokens) + page_size - 1) // page_size swa_pages = (math.ceil(swa_tokens) + page_size - 1) // page_size - full_page_bytes = self.full_attn_allocator.entry_bytes_per_page - swa_page_bytes = self.swa_attn_allocator.entry_bytes_per_page - swa_full_pages = ( - swa_pages * swa_page_bytes + full_page_bytes - 1 - ) // full_page_bytes - - joint_full_tokens = (full_pages + swa_full_pages) * page_size - aligned_swa_tokens = swa_pages * page_size + if swa_pages > full_pages: + return None + + compacted = empty_pool or not self.lazy_compaction or self._compaction_allowed() + + def fits(full_reclaim_pages: int, swa_reclaim_pages: int) -> bool: + return self._fits_page_demand( + full_pages, + swa_pages, + full_reclaim_pages=full_reclaim_pages, + swa_reclaim_pages=swa_reclaim_pages, + compacted=compacted, + empty_pool=empty_pool, + ) + + if empty_pool: + return (0, 0) if fits(0, 0) else None + if fits(0, 0): + return (0, 0) + + max_full_pages = min( + self.full_attn_allocator.allocated_count() // page_size, + max(0, int(full_evictable_tokens)) // page_size, + ) + max_swa_pages = min( + self.swa_attn_allocator.allocated_count() // page_size, + max(0, int(swa_evictable_tokens)) // page_size, + ) + if not fits(max_full_pages, max_swa_pages): + return None + + def first_fit(high: int, predicate: Callable[[int], bool]) -> int: + low = 0 + while low < high: + mid = (low + high) // 2 + if predicate(mid): + high = mid + else: + low = mid + 1 + return low + + swa_reclaim_pages = first_fit( + max_swa_pages, lambda value: fits(max_full_pages, value) + ) + full_reclaim_pages = first_fit( + max_full_pages, lambda value: fits(value, swa_reclaim_pages) + ) return ( - max(0, joint_full_tokens - self.full_available_size()), - max(0, aligned_swa_tokens - self.swa_available_size()), + full_reclaim_pages * page_size, + swa_reclaim_pages * page_size, ) def can_reserve( @@ -3469,9 +3515,11 @@ def can_reserve( """Check pending FULL/SWA demand against the shared byte envelope. Scheduler admission keeps the historical one-token strict slack at an - empty-pool boundary. For a live pool, shared bytes are backed only by - FULL eviction; SWA eviction backs its own slot shortfall. + empty-pool boundary. Live admission checks the state reachable after + reclaiming the currently evictable pages from both sides. """ + if full_tokens < 0 or swa_tokens < 0: + return False if not self.supports_asymmetric_reservation: if ( full_tokens != swa_tokens @@ -3481,42 +3529,20 @@ def can_reserve( ): return False return full_tokens <= self.available_size() - if full_tokens < 0 or swa_tokens < 0: - return False if require_token_slack and ( full_tokens >= self.size_full or swa_tokens >= self.size_swa ): return False - if empty_pool: - page_size = self.page_size - full_pages = (math.ceil(full_tokens) + page_size - 1) // page_size - swa_pages = (math.ceil(swa_tokens) + page_size - 1) // page_size - return self._fits_page_demand( - full_pages, - swa_pages, - compacted=True, - empty_pool=True, + return ( + self.reclaim_plan( + full_tokens, + swa_tokens, + full_evictable_tokens=full_evictable_tokens, + swa_evictable_tokens=swa_evictable_tokens, + empty_pool=empty_pool, ) - - full_shortfall, swa_shortfall = self.allocation_shortfalls( - full_tokens, swa_tokens - ) - page_size = self.page_size - if ( - full_shortfall > max(0, int(full_evictable_tokens)) // page_size * page_size - or swa_shortfall - > max(0, int(swa_evictable_tokens)) // page_size * page_size - ): - return False - - compacted = not self.lazy_compaction or self._compaction_allowed() - return self._fits_page_demand( - (math.ceil(full_tokens) + page_size - 1) // page_size, - (math.ceil(swa_tokens) + page_size - 1) // page_size, - full_reclaim_pages=full_shortfall // page_size, - swa_reclaim_pages=swa_shortfall // page_size, - compacted=compacted, + is not None ) def _compaction_allowed(self) -> bool: @@ -3719,9 +3745,8 @@ def swa_v2p_page_table(self) -> torch.Tensor: def alloc(self, need_size: int) -> Optional[torch.Tensor]: with record_function("UnifiedSWAAlloc.alloc"): - if need_size > self.available_size(): - if not _relieve_for_alloc(self, need_size): - return None + if not self.ensure_capacity(need_size, need_size): + return None # Snapshot the virtual PAGES full will consume, to bind them on swa too. num_pages = need_size // self.page_size fa = self.full_attn_allocator @@ -3757,9 +3782,8 @@ def alloc_extend( prefix_lens=prefix_lens_cpu, ) need_tokens = num_new_pages * self.page_size - if need_tokens > self.available_size(): - if not _relieve_for_alloc(self, need_tokens): - return None + if not self.ensure_capacity(need_tokens, need_tokens): + return None # Snapshot the virtual PAGES the kernel will consume; clone so swa keeps # its view after the slice is consumed. @@ -3796,9 +3820,8 @@ def alloc_decode( seq_lens=seq_lens_cpu, page_size=self.page_size, decode=True ) need_tokens = num_new_pages * self.page_size - if need_tokens > self.available_size(): - if not _relieve_for_alloc(self, need_tokens): - return None + if not self.ensure_capacity(need_tokens, need_tokens): + return None fa = self.full_attn_allocator new_virtual_pages = fa.free_virtual_ids[:num_new_pages].clone() diff --git a/python/sglang/srt/model_executor/pool_configurator.py b/python/sglang/srt/model_executor/pool_configurator.py index 7c7767ccd68c..637809c954ea 100644 --- a/python/sglang/srt/model_executor/pool_configurator.py +++ b/python/sglang/srt/model_executor/pool_configurator.py @@ -464,9 +464,9 @@ def __init__(self, kvc: KVCacheConfigurator): self._full_layers_num = len(model_config.full_attention_layer_ids) self._swa_layers_num = len(model_config.swa_attention_layer_ids) - assert self._swa_layers_num > 0, ( - "Hybrid SWA model must have at least one SWA layer" - ) + assert ( + self._swa_layers_num > 0 + ), "Hybrid SWA model must have at least one SWA layer" self._swa_full_tokens_ratio = get_schedule().swa_full_tokens_ratio self._sliding_window_size = kvc.sliding_window_size @@ -662,7 +662,8 @@ def align_page_size(x: int) -> int: swa_tokens = align_page_size(int(full_tokens * self._swa_full_tokens_ratio)) if ( - self._sliding_window_size is not None + not self._enable_unified_memory + and self._sliding_window_size is not None and self._sliding_window_size + self._page_size >= swa_tokens ): raise ValueError( @@ -1094,9 +1095,9 @@ def finalize_with_max_running_requests( def calculate_pool_sizes( self, available_bytes: int, page_size: int ) -> MemoryPoolConfig: - assert page_size % 128 == 0, ( - "page_size must be multiple of 128 for compressed attention" - ) + assert ( + page_size % 128 == 0 + ), "page_size must be multiple of 128 for compressed attention" if self.requested_max_running_requests_per_worker is not None: c128_state_fixed_bytes = self._get_c128_state_fixed_bytes( @@ -1124,9 +1125,9 @@ def calculate_pool_sizes( def calculate_pool_sizes_from_max_tokens( self, max_total_num_tokens: int, page_size: int ) -> MemoryPoolConfig: - assert page_size % 128 == 0, ( - "page_size must be multiple of 128 for compressed attention" - ) + assert ( + page_size % 128 == 0 + ), "page_size must be multiple of 128 for compressed attention" sizes = self._compute_dsv4_sizes(max_total_num_tokens, page_size) return self._to_config(sizes) From 653c010f0d9b16837a7aa311868f92e01b05f679 Mon Sep 17 00:00:00 2001 From: yhzhuang Date: Tue, 8 Sep 2026 16:30:59 -0700 Subject: [PATCH 27/55] Restore unified pool cleanup and scoped compatibility gates --- python/sglang/srt/arg_groups/kv_cache_hook.py | 11 ++++----- python/sglang/srt/mem_cache/common.py | 24 +++++++++---------- .../srt/mem_cache/multi_ended_allocator.py | 24 +++++++++---------- .../srt/mem_cache/unified_radix_cache.py | 9 ------- .../srt/model_executor/pool_configurator.py | 18 +++++++------- 5 files changed, 37 insertions(+), 49 deletions(-) diff --git a/python/sglang/srt/arg_groups/kv_cache_hook.py b/python/sglang/srt/arg_groups/kv_cache_hook.py index d01814ea8a90..26f788c14052 100644 --- a/python/sglang/srt/arg_groups/kv_cache_hook.py +++ b/python/sglang/srt/arg_groups/kv_cache_hook.py @@ -198,14 +198,11 @@ def handle_unified_memory_pool(server_args: Any) -> None: cfg = resolving_view(server_args) if not cfg.enable_unified_memory: return - assert cfg.pp_size == 1, ( - "--enable-unified-memory does not support " - "pipeline parallelism (--pp-size > 1)." - ) - assert ( - not cfg.enable_mixed_chunk - ), "--enable-unified-memory does not support --enable-mixed-chunk." if cfg.disaggregation_mode != "null": + assert cfg.pp_size == 1, ( + "--enable-unified-memory with PD disaggregation does not support " + "pipeline parallelism (--pp-size > 1)." + ) # Constraints of the whole-envelope transfer; see # UnifiedMLATokenToKVPool.get_contiguous_buf_infos. assert cfg.disaggregation_transfer_backend == "mooncake", ( diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index ebf437dc1b2a..b3f7e0b0a692 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/python/sglang/srt/mem_cache/common.py @@ -66,9 +66,9 @@ def free_swa_out_of_window_slots( return # 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.kv.cache_protected_len % page_size == 0 - ), "cache_protected_len must be page aligned" + assert req.kv.cache_protected_len % page_size == 0, ( + "cache_protected_len must be page aligned" + ) req.kv.swa_evicted_seqlen = max( req.kv.swa_evicted_seqlen, req.kv.swa_dead_lo(page_size) ) @@ -272,9 +272,9 @@ def release_kv_cache(req: Req, tree_cache: BasePrefixCache, is_insert: bool = Tr assert (not req.kv.holds_kv) == req.kv.is_kv_released # MambaRadixCache may alloc mamba state before alloc KV cache if not req.kv.holds_kv: - assert ( - tree_cache.supports_mamba() - ), "Only MambaRadixCache allow freeing before alloc" + assert tree_cache.supports_mamba(), ( + "Only MambaRadixCache allow freeing before alloc" + ) # TODO (csy, hanming): clean up this early allocation logic if req.kv.holds_mamba: tree_cache.req_to_token_pool.mamba_allocator.free( @@ -303,9 +303,9 @@ def release_kv_cache(req: Req, tree_cache: BasePrefixCache, is_insert: bool = Tr if isinstance(tree_cache.req_to_token_pool, HybridReqToTokenPool) and ( not tree_cache.supports_mamba() ): - assert ( - req.kv.holds_mamba - ), "mamba state is freed while the tree cache does not manage mamba states" + assert req.kv.holds_mamba, ( + "mamba state is freed while the tree cache does not manage mamba states" + ) tree_cache.req_to_token_pool.free_mamba_cache(req) # The DSV4-NPU ReqToTokenPool subclass's free() additionally releases the # c4/c128 state pages; other ReqToTokenPool subclasses are a no-op here. @@ -323,9 +323,9 @@ def _release_overallocated_kv_indices( # strip_thinking_cache intentionally reports output tokens as overallocated # so they fall into the free path below (#22373). if spec_algo is None and not get_serving().strip_thinking_cache: - assert ( - start_p == end_p - ), f"Unexpected overallocated KV cache, {req.kv.kv_committed_len=}, {req.kv.kv_allocated_len=}" + assert start_p == end_p, ( + f"Unexpected overallocated KV cache, {req.kv.kv_committed_len=}, {req.kv.kv_allocated_len=}" + ) if page_size > 1: start_p = ceil_align(start_p, page_size) diff --git a/python/sglang/srt/mem_cache/multi_ended_allocator.py b/python/sglang/srt/mem_cache/multi_ended_allocator.py index 1419bb7af136..e82a7c7e37b1 100644 --- a/python/sglang/srt/mem_cache/multi_ended_allocator.py +++ b/python/sglang/srt/mem_cache/multi_ended_allocator.py @@ -329,9 +329,9 @@ def __init__( self.num_pages = max_slots // self.pool_page_size # `min_page_index` = ceil(min_slot_index / pool_page_size), keeping the # reserved-sink invariant (min_page_index * entry_bytes_per_page >= entry_max). - assert ( - virtual_num_pages is None or not is_id_owner - ), "only a non-owner allocator may use another pool's virtual-id space" + assert virtual_num_pages is None or not is_id_owner, ( + "only a non-owner allocator may use another pool's virtual-id space" + ) # An ID owner has one virtual page per physical page. A non-owner may # deliberately have a different physical capacity while indexing the # owner's virtual page space (the unified SWA allocator does this). @@ -369,9 +369,9 @@ def __init__( self.high_peer: Optional[MultiEndedAllocator] = None # Inverse history of relocations (spec rollback), at PAGE granularity. - self._inverse_history: List[Tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = ( - [] - ) + self._inverse_history: List[ + Tuple[torch.Tensor, torch.Tensor, torch.Tensor] + ] = [] # --- Lazy compaction state (all unused when lazy_compaction=False) --- # `_free_phys_pages`: GPU free list of physical PAGE ids, sorted at `_flush`. @@ -1233,9 +1233,9 @@ def alloc_extend( v2p stays -1 and translation yields negative ids → CUDA OOB. """ with record_function("MultiEndedAlloc.alloc_extend"): - assert ( - self.is_id_owner - ), f"alloc_extend on a non-id-owner allocator ({self.sub_pool_name!r})" + assert self.is_id_owner, ( + f"alloc_extend on a non-id-owner allocator ({self.sub_pool_name!r})" + ) if num_new_pages is None: num_new_pages = get_num_new_pages( seq_lens=seq_lens_cpu, @@ -1302,9 +1302,9 @@ def alloc_decode( virtual page on THIS sub-allocator (else v2p stays -1 → CUDA OOB). """ with record_function("MultiEndedAlloc.alloc_decode"): - assert ( - self.is_id_owner - ), f"alloc_decode on a non-id-owner allocator ({self.sub_pool_name!r})" + assert self.is_id_owner, ( + f"alloc_decode on a non-id-owner allocator ({self.sub_pool_name!r})" + ) bs = len(seq_lens) # CPU-only count BEFORE the kernel, to snapshot the exact slice the # kernel will consume. diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index 8fd73e11812c..6c13b7cafc72 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -625,15 +625,6 @@ def _evict( available_size_targets=available_size_targets, ) - return self._finish_eviction(tracker, start_time) - - def _finish_eviction( - self, - tracker: dict[ComponentType, int], - start_time: float, - ) -> EvictResult: - """Finish write-back bookkeeping and build the public result.""" - if ( self.cache_controller is not None and self.cache_controller.write_policy == "write_back" diff --git a/python/sglang/srt/model_executor/pool_configurator.py b/python/sglang/srt/model_executor/pool_configurator.py index 637809c954ea..679e7249f942 100644 --- a/python/sglang/srt/model_executor/pool_configurator.py +++ b/python/sglang/srt/model_executor/pool_configurator.py @@ -464,9 +464,9 @@ def __init__(self, kvc: KVCacheConfigurator): self._full_layers_num = len(model_config.full_attention_layer_ids) self._swa_layers_num = len(model_config.swa_attention_layer_ids) - assert ( - self._swa_layers_num > 0 - ), "Hybrid SWA model must have at least one SWA layer" + assert self._swa_layers_num > 0, ( + "Hybrid SWA model must have at least one SWA layer" + ) self._swa_full_tokens_ratio = get_schedule().swa_full_tokens_ratio self._sliding_window_size = kvc.sliding_window_size @@ -1095,9 +1095,9 @@ def finalize_with_max_running_requests( def calculate_pool_sizes( self, available_bytes: int, page_size: int ) -> MemoryPoolConfig: - assert ( - page_size % 128 == 0 - ), "page_size must be multiple of 128 for compressed attention" + assert page_size % 128 == 0, ( + "page_size must be multiple of 128 for compressed attention" + ) if self.requested_max_running_requests_per_worker is not None: c128_state_fixed_bytes = self._get_c128_state_fixed_bytes( @@ -1125,9 +1125,9 @@ def calculate_pool_sizes( def calculate_pool_sizes_from_max_tokens( self, max_total_num_tokens: int, page_size: int ) -> MemoryPoolConfig: - assert ( - page_size % 128 == 0 - ), "page_size must be multiple of 128 for compressed attention" + assert page_size % 128 == 0, ( + "page_size must be multiple of 128 for compressed attention" + ) sizes = self._compute_dsv4_sizes(max_total_num_tokens, page_size) return self._to_config(sizes) From f604efb38fe5ac6e9b94ce37bc44a1d5125418fb Mon Sep 17 00:00:00 2001 From: yhzhuang Date: Thu, 10 Sep 2026 15:47:12 -0700 Subject: [PATCH 28/55] Keep HiSparse fixture compatible with Python 3.10 --- .../unit/mem_cache/test_hisparse_max_token_pool_size.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/test/registered/unit/mem_cache/test_hisparse_max_token_pool_size.py b/test/registered/unit/mem_cache/test_hisparse_max_token_pool_size.py index dd4f1726d4aa..4c590ce02bf1 100644 --- a/test/registered/unit/mem_cache/test_hisparse_max_token_pool_size.py +++ b/test/registered/unit/mem_cache/test_hisparse_max_token_pool_size.py @@ -78,9 +78,6 @@ def test_non_hisparse_uses_effective_max_total_num_tokens(self): self.assertEqual(instance.max_token_pool_size, 1024) def test_non_hisparse_hybrid_swa_prefers_full_max(self): - self.enterContext( - get_context().override_server_args(enable_unified_memory=False) - ) instance = _make_model_runner( enable_hisparse=False, token_to_kv_pool_allocator=SimpleNamespace(), @@ -89,8 +86,9 @@ def test_non_hisparse_hybrid_swa_prefers_full_max(self): full_max_total_num_tokens=3000, swa_max_total_num_tokens=500, ) - self.assertEqual(instance.max_token_pool_size, 3000) - self.assertEqual(instance.effective_max_total_num_tokens, 3000) + with get_context().override_server_args(enable_unified_memory=False): + self.assertEqual(instance.max_token_pool_size, 3000) + self.assertEqual(instance.effective_max_total_num_tokens, 3000) def _make_prealloc_queue( From 70312b502dffbc6eda79a9187b0e120e7e940906 Mon Sep 17 00:00:00 2001 From: yhzhuang Date: Thu, 10 Sep 2026 16:04:23 -0700 Subject: [PATCH 29/55] Remove unrelated UMP diff churn --- python/sglang/srt/arg_groups/kv_cache_hook.py | 9 +++-- python/sglang/srt/disaggregation/decode.py | 2 + python/sglang/srt/managers/schedule_policy.py | 6 +-- .../scheduler_components/invariant_checker.py | 38 +++++++++---------- python/sglang/srt/mem_cache/allocator/base.py | 10 ----- .../mem_cache/allocator/unified_sub_pool.py | 24 +++++------- 6 files changed, 36 insertions(+), 53 deletions(-) diff --git a/python/sglang/srt/arg_groups/kv_cache_hook.py b/python/sglang/srt/arg_groups/kv_cache_hook.py index dde637e03805..d872310fc9d7 100644 --- a/python/sglang/srt/arg_groups/kv_cache_hook.py +++ b/python/sglang/srt/arg_groups/kv_cache_hook.py @@ -200,10 +200,6 @@ def handle_unified_memory_pool(server_args: Any) -> None: if not cfg.enable_unified_memory: return if cfg.disaggregation_mode != "null": - assert cfg.pp_size == 1, ( - "--enable-unified-memory with PD disaggregation does not support " - "pipeline parallelism (--pp-size > 1)." - ) # Constraints of the whole-envelope transfer; see # UnifiedMLATokenToKVPool.get_contiguous_buf_infos. assert cfg.disaggregation_transfer_backend == "mooncake", ( @@ -211,6 +207,11 @@ def handle_unified_memory_pool(server_args: Any) -> None: "the mooncake transfer backend; got " f"{cfg.disaggregation_transfer_backend!r}." ) + assert cfg.pp_size == 1, ( + "--enable-unified-memory with PD disaggregation does not support " + "pipeline parallelism (whole-envelope transfer has no per-layer " + "entries to subset)." + ) assert not envs.SGLANG_DISABLE_LAZY_COMPACTION.get(), ( "--enable-unified-memory with PD disaggregation requires lazy " "compaction; unset SGLANG_DISABLE_LAZY_COMPACTION." diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index f3f55314498c..9ac2121e3759 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -445,6 +445,7 @@ def _reclaim_swa_tail_capacity( EvictParams(swa_num_tokens=required - available) ) available = self.token_to_kv_pool_allocator.swa_available_size() + if available < required: return ( f"SWA eviction insufficient: needed={required}, " @@ -845,6 +846,7 @@ def resume_retracted_reqs( break if uses_swa_tail_prealloc and swa_required > swa_allocatable_tokens: break + resumed_reqs.append(req) indices_to_remove.add(i) req.is_retracted = False diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index 7399f57bd8c1..2a59e58e3636 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -903,8 +903,7 @@ def _swa_chunk_cap( ) if cap <= 0: return 0 - cap = cap // self.page_size * self.page_size - return cap + return cap // self.page_size * self.page_size assert max_chunk_tokens is not None cap = max_chunk_tokens // self.page_size * self.page_size @@ -1481,6 +1480,7 @@ def add_one_req( # this returns 0, so the debit sites below reuse the value. mamba_gap_reserve = self._mamba_gap_budget_for_req(req) total_tokens += mamba_gap_reserve + # adjusting the input_tokens based on host_hit_length and page_size swa_extend_input_len = cand_extend_input_len - req.host_hit_length real_input_tokens = self.ceil_paged_tokens(swa_extend_input_len) @@ -1510,7 +1510,7 @@ def add_one_req( return AddReqResult.OTHER with self._lock_node(req.last_node): - # Capacity may decrease after the lock acquisition. + # self.rem_total_tokens may decrease after the lock acquisition if not self.is_unified_swa and total_tokens >= self.rem_total_tokens: return AddReqResult.NO_TOKEN diff --git a/python/sglang/srt/managers/scheduler_components/invariant_checker.py b/python/sglang/srt/managers/scheduler_components/invariant_checker.py index d232b6d9f0fb..a3a95d0128d8 100644 --- a/python/sglang/srt/managers/scheduler_components/invariant_checker.py +++ b/python/sglang/srt/managers/scheduler_components/invariant_checker.py @@ -96,9 +96,6 @@ def _check_pool_invariant( def _check_full_pool(self, ps: PoolStats, uncached: int = 0) -> Tuple[bool, str]: allocator = self.token_to_kv_pool_allocator - is_unified_mamba_swa = isinstance( - allocator, UnifiedMambaSWATokenToKVPoolAllocator - ) full_capacity = ( allocator.current_full_capacity if isinstance(allocator, UnifiedSWATokenToKVPoolAllocator) @@ -139,14 +136,14 @@ def _check_full_pool(self, ps: PoolStats, uncached: int = 0) -> Tuple[bool, str] // allocator.page_size * allocator.page_size ) - full_available_size = ( - allocator.conserve_full_available_size() - if is_unified_mamba_swa - else ps.full_available_size - ) + full_available = ps.full_available_size + if isinstance(allocator, UnifiedMambaSWATokenToKVPoolAllocator): + # Pair the static per-layer total with the conserve view, never the + # byte-coordinated one -- see `conserve_full_available_size`. + full_available = allocator.conserve_full_available_size() leak, msg = self._check_pool_invariant( "full", - full_available_size, + full_available, full_evictable_size, protected, session_held, @@ -172,24 +169,23 @@ def _check_swa_pool(self, ps: PoolStats, uncached: int = 0) -> Tuple[bool, str]: f"evictable={ps.swa_evictable_size}, " f"total={self.swa_tokens_per_layer}" ) + swa_available = ps.swa_available_size if isinstance(allocator, UnifiedMambaSWATokenToKVPoolAllocator): - available = allocator.conserve_swa_available_size() - total = self.swa_tokens_per_layer - else: - available = ps.swa_available_size - total = ( - allocator.current_swa_capacity - if isinstance(allocator, UnifiedSWATokenToKVPoolAllocator) - and allocator.supports_asymmetric_reservation - else self.swa_tokens_per_layer - ) + # Tri-pool: same floating-boundary phantom as the full pool -- use the + # slot-conservation view, not the byte-coordinated min (see _check_full_pool). + swa_available = allocator.conserve_swa_available_size() return self._check_pool_invariant( "swa", - available, + swa_available, ps.swa_evictable_size, self.tree_cache.swa_protected_size(), self.pool_stats_observer.session_held_swa_tokens(), - total, + ( + allocator.current_swa_capacity + if isinstance(allocator, UnifiedSWATokenToKVPoolAllocator) + and allocator.supports_asymmetric_reservation + else self.swa_tokens_per_layer + ), uncached, ) diff --git a/python/sglang/srt/mem_cache/allocator/base.py b/python/sglang/srt/mem_cache/allocator/base.py index 42871b243f9f..fa640b650672 100644 --- a/python/sglang/srt/mem_cache/allocator/base.py +++ b/python/sglang/srt/mem_cache/allocator/base.py @@ -88,16 +88,6 @@ def verify_byte_accounting(self) -> list: violation strings, empty when healthy. Static pools have no byte model.""" return [] - @property - def current_full_capacity(self) -> int: - """Capacity represented by full_available_size plus live allocations.""" - return self.size_full - - @property - def current_swa_capacity(self) -> int: - """SWA counterpart of current_full_capacity for hybrid allocators.""" - return self.size_swa - def mamba_full_cache_donor(self) -> MambaFullCacheDonor | None: """Return the shared-pool donor capability, if this allocator has one.""" return None diff --git a/python/sglang/srt/mem_cache/allocator/unified_sub_pool.py b/python/sglang/srt/mem_cache/allocator/unified_sub_pool.py index c70f697c149f..6e2e0ed61375 100644 --- a/python/sglang/srt/mem_cache/allocator/unified_sub_pool.py +++ b/python/sglang/srt/mem_cache/allocator/unified_sub_pool.py @@ -312,27 +312,21 @@ def __init__( self.page_size = page_size * dcp_size self.num_pages = max_slots // self.pool_page_size # `min_page_index` = ceil(min_slot_index / pool_page_size), keeping the - # reserved-sink invariant (min_page_index * entry_bytes_per_page >= entry_max). + # reserved sink floor covered (see `_reserved_floor_bytes`). + self.min_page_index = ( + self.min_slot_index + self.pool_page_size - 1 + ) // self.pool_page_size + self.entry_bytes_per_page = self.entry_bytes * self.pool_page_size + + # v2p is indexed by VIRTUAL page id, p2v by PHYSICAL page id. A non-owner + # consumes the owner's ids, so the two counts are unrelated. assert virtual_num_pages is None or not is_id_owner, ( "only a non-owner allocator may use another pool's virtual-id space" ) - # An ID owner has one virtual page per physical page. A non-owner may - # deliberately have a different physical capacity while indexing the - # owner's virtual page space (the unified SWA allocator does this). self.num_virtual_ids = ( - virtual_num_pages if virtual_num_pages is not None else self.num_pages + self.num_pages if virtual_num_pages is None else virtual_num_pages ) - if is_id_owner: - assert self.num_virtual_ids == self.num_pages assert self.num_virtual_ids > 0, "virtual page count must be positive" - self.min_page_index = ( - self.min_slot_index + self.pool_page_size - 1 - ) // self.pool_page_size - self.entry_bytes_per_page = self.entry_bytes * self.pool_page_size - - # v2p is indexed by VIRTUAL page id, p2v by PHYSICAL page id. A - # non-owner consumes the owner's ids, so its v2p spans the owner's - # count; the two are unrelated and either can be the larger. # Page 0 is the padding anchor; the trailing row is the -1 sentinel. self.virtual_to_physical = torch.full( (self.num_virtual_ids + 1,), From 4e0e1c35dcd86d1b2e1ab12f15f1feab1854fbcb Mon Sep 17 00:00:00 2001 From: yhzhuang Date: Sat, 12 Sep 2026 09:10:24 -0700 Subject: [PATCH 30/55] Fix unified pool admission rounding --- python/sglang/srt/managers/schedule_policy.py | 6 +-- python/sglang/srt/managers/scheduler.py | 2 +- .../test_scheduler_init_req_max_new_tokens.py | 34 ++++++++++++ .../mem_cache/test_multi_ended_allocator.py | 52 +++++++++++++++++++ 4 files changed, 88 insertions(+), 6 deletions(-) diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index 2a59e58e3636..8e68faedbdc7 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -982,11 +982,7 @@ def _check_swa_admission( full_ever_tokens = unified_swa_tokens = None if self.is_unified_swa: full_input_tokens = len(req.full_untruncated_fill_ids) - full_ever_tokens = ( - self.ceil_paged_tokens(full_input_tokens) - + max_new_tokens - + self.page_size - ) + full_ever_tokens = full_input_tokens + max_new_tokens + self.page_size unified_swa_tokens = self._swa_budget_for_req( full_input_tokens, max_new_tokens ) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index dc35a96ad948..7b7b8d375f52 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -2528,7 +2528,7 @@ def init_req_max_new_tokens(self, req): ): def fits_shared_pool(candidate: int) -> bool: - full_tokens = paged_input_len + candidate + self.page_size + full_tokens = input_len + candidate + self.page_size swa_tokens = estimate_swa_kv_tokens( input_len, candidate, diff --git a/test/registered/unit/managers/test_scheduler_init_req_max_new_tokens.py b/test/registered/unit/managers/test_scheduler_init_req_max_new_tokens.py index 431dd85b744e..43f182778e26 100644 --- a/test/registered/unit/managers/test_scheduler_init_req_max_new_tokens.py +++ b/test/registered/unit/managers/test_scheduler_init_req_max_new_tokens.py @@ -2,8 +2,11 @@ import unittest from types import SimpleNamespace +import torch + from sglang.srt.environ import envs from sglang.srt.managers.scheduler import Scheduler +from sglang.srt.mem_cache.unified_memory_pool import init_unified_swa_pools from sglang.srt.runtime_context import get_parallel from sglang.test.ci.ci_register import register_cpu_ci @@ -181,6 +184,37 @@ def test_admission_rules_sweep(self): ) self._init_and_check(scheduler, req) + def test_unified_budget_rounds_prompt_and_decode_together(self): + bundle = init_unified_swa_pools( + device="cpu", + kv_cache_dtype=torch.float16, + head_num=1, + head_dim=4, + v_head_dim=4, + swa_head_num=1, + swa_head_dim=4, + swa_v_head_dim=4, + page_size=4, + start_layer=0, + end_layer=2, + swa_attention_layer_ids=[1], + full_attention_layer_ids=[0], + total_bytes=384, + enable_memory_saver=False, + need_sort=False, + lazy_compaction=True, + ) + scheduler = self._new_scheduler(page_size=4) + scheduler.token_to_kv_pool_allocator = bundle.token_to_kv_pool_allocator + scheduler.sliding_window_size = 4 + scheduler.chunked_prefill_size = 4 + scheduler.max_new_tokens_limit = None + for prompt_len in (4, 5, 6, 7): + with self.subTest(prompt_len=prompt_len): + req = self._new_req(max_new_tokens=1, input_len=prompt_len) + scheduler.init_req_max_new_tokens(req) + self.assertEqual(req.sampling_params.max_new_tokens, 1) + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/mem_cache/test_multi_ended_allocator.py b/test/registered/unit/mem_cache/test_multi_ended_allocator.py index 2baf5cceb07c..bac0687f9f1d 100644 --- a/test/registered/unit/mem_cache/test_multi_ended_allocator.py +++ b/test/registered/unit/mem_cache/test_multi_ended_allocator.py @@ -601,6 +601,58 @@ def _build( ) return pool, allocator, kvcache + def test_reclaim_plan_matches_exhaustive_page_targets(self): + page_size = 4 + _, allocator, _ = self._build( + n_full_slots=40, n_swa_slots=24, page_size=page_size + ) + allocator.lazy_compaction = True + for sub_pool in (allocator.full_attn_allocator, allocator.swa_attn_allocator): + sub_pool.lazy_compaction = True + sub_pool.disagg_move_gate = lambda: False + live = allocator.alloc(16) + self.assertIsNotNone(live) + allocator.free(live[4:8]) + allocator.free_swa(live[8:12]) + + for compacted in (False, True): + for sub_pool in ( + allocator.full_attn_allocator, + allocator.swa_attn_allocator, + ): + sub_pool.disagg_move_gate = lambda: compacted + for full_evictable, swa_evictable in ((0, 0), (7, 5), (12, 8), (100, 100)): + max_full = min(12, full_evictable) // page_size + max_swa = min(8, swa_evictable) // page_size + for full_pages in range(9): + for swa_pages in range(full_pages + 1): + feasible = [ + (full * page_size, swa * page_size) + for swa in range(max_swa + 1) + for full in range(max_full + 1) + if allocator._fits_page_demand( + full_pages, + swa_pages, + full_reclaim_pages=full, + swa_reclaim_pages=swa, + compacted=compacted, + ) + ] + with self.subTest( + compacted=compacted, + evictable=(full_evictable, swa_evictable), + pages=(full_pages, swa_pages), + ): + self.assertEqual( + allocator.reclaim_plan( + full_pages * page_size, + swa_pages * page_size, + full_evictable_tokens=full_evictable, + swa_evictable_tokens=swa_evictable, + ), + feasible[0] if feasible else None, + ) + def test_empty_pool_reservation_matches_packed_byte_boundary(self): page_size = 4 _, allocator, _ = self._build( From 70671d52eafe13751974bc74630efa74e37542bb Mon Sep 17 00:00:00 2001 From: yhzhuang Date: Sat, 12 Sep 2026 21:38:57 -0700 Subject: [PATCH 31/55] Refactor unified SWA budget handling --- python/sglang/srt/managers/schedule_policy.py | 258 ++++++++---------- python/sglang/srt/managers/scheduler.py | 78 +++--- .../scheduler_components/invariant_checker.py | 35 +-- .../pool_stats_observer.py | 15 +- .../mem_cache/allocator/unified_hybrid_swa.py | 61 +++-- python/sglang/srt/mem_cache/common.py | 27 +- .../srt/model_executor/pool_configurator.py | 27 +- 7 files changed, 237 insertions(+), 264 deletions(-) diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index 8e68faedbdc7..de727a9f1b73 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -54,7 +54,7 @@ ) from sglang.srt.mem_cache.allocator.unified_hybrid_swa import ( UnifiedMambaSWATokenToKVPoolAllocator, - UnifiedSWATokenToKVPoolAllocator, + supports_swa_byte_budget, ) from sglang.srt.mem_cache.allocator.unified_mamba import ( UnifiedMambaTokenToKVPoolAllocator, @@ -640,12 +640,7 @@ def __init__( self.token_to_kv_pool_allocator, PureSWATokenToKVPoolAllocator ) self.is_hybrid_ssm_cache = self.tree_cache.supports_mamba() - self.is_unified_swa = ( - isinstance( - self.token_to_kv_pool_allocator, UnifiedSWATokenToKVPoolAllocator - ) - and self.token_to_kv_pool_allocator.supports_asymmetric_reservation - ) + self.is_unified_swa = supports_swa_byte_budget(self.token_to_kv_pool_allocator) self.rem_swa_token_offset = ( num_mixed_decode_tokens if self.is_unified_swa else 0 @@ -856,8 +851,6 @@ def _unified_swa_reservation_fits( *, empty_pool: bool = False, ) -> bool: - if not self.is_unified_swa: - return True allocator = self.token_to_kv_pool_allocator if empty_pool: full_evictable_tokens = 0 @@ -876,36 +869,14 @@ def _unified_swa_reservation_fits( require_token_slack=empty_pool, ) - def _swa_chunk_cap( + def _unified_swa_chunk_cap( self, max_new_tokens: int, swa_host_hit_length: int = 0, *, - max_chunk_tokens: Optional[int] = None, + max_chunk_tokens: int, ) -> int: - """Largest page-aligned extend chunk the SWA pool can admit right now. - - Split pools reserve headroom below rem_swa_tokens; unified pools use - the joint byte budget. Returns 0 if not even one page fits. Only valid - when is_hybrid_swa is True. - - Escape hatch for a request whose budget can never pass the - _swa_budget_for_req gate (extend near/above the pool size, or a large - load-back charge): without shrinking its chunk it would be rejected - forever (head-of-line livelock). Shrinking is sound because past a - chunk boundary only the sliding window stays locked — the rest turns - evictable — so each pass's transient footprint fits the pool.""" - if not self.is_unified_swa: - # extend_input_len=0: this solves for the extend chunk itself, so - # the reserved headroom is the post-chunk decode window only. - cap = int(self.rem_swa_tokens) - self._swa_reserved_tokens( - 0, max_new_tokens, swa_host_hit_length - ) - if cap <= 0: - return 0 - return cap // self.page_size * self.page_size - - assert max_chunk_tokens is not None + """Largest page-aligned chunk that fits the joint FULL/SWA byte budget.""" cap = max_chunk_tokens // self.page_size * self.page_size lo, hi = 0, cap // self.page_size while lo < hi: @@ -921,14 +892,18 @@ def _swa_chunk_cap( hi = mid - 1 return lo * self.page_size + def _swa_chunk_cap(self, max_new_tokens: int, swa_host_hit_length: int = 0) -> int: + """Largest chunk fitting the fixed SWA partition after decode headroom.""" + cap = int(self.rem_swa_tokens) - self._swa_reserved_tokens( + 0, max_new_tokens, swa_host_hit_length + ) + return max(0, cap // self.page_size * self.page_size) + def _swa_req_never_fits( self, extend_input_len: int, max_new_tokens: int, swa_host_hit_length: int = 0, - *, - full_tokens: Optional[int] = None, - unified_swa_tokens: Optional[int] = None, ) -> bool: """True when a request's SWA budget exceeds the *entire* SWA pool, so it can never be admitted whole no matter how far the pool drains. @@ -943,72 +918,83 @@ def _swa_req_never_fits( swa_tokens = self._swa_budget_for_req( extend_input_len, max_new_tokens, swa_host_hit_length ) - if self.is_unified_swa: - assert full_tokens is not None - assert unified_swa_tokens is not None - return not self._unified_swa_reservation_fits( - full_tokens, unified_swa_tokens, empty_pool=True - ) if self._swa_req_ring: return swa_tokens > self.token_to_kv_pool_allocator.size_swa return swa_tokens >= self.token_to_kv_pool_allocator.size_swa - def _check_swa_admission( + def _check_unified_swa_admission( self, req: Req, *, extend_input_len: int, total_tokens: int, ) -> tuple[bool, Optional[int]]: - """Check SWA pressure and return the usable chunk limit.""" - if not self.is_unified_swa: - extend_input_len = self.ceil_paged_tokens(extend_input_len) max_new_tokens = self._swa_new_tokens(req) swa_needed = self._swa_budget_for_req( extend_input_len, max_new_tokens, swa_host_hit_length=req.swa_host_hit_length, ) - reservation_fits = ( - self._unified_swa_reservation_fits(total_tokens, swa_needed) - if self.is_unified_swa - else swa_needed <= self.rem_swa_tokens - if self._swa_req_ring - else swa_needed < self.rem_swa_tokens - ) - if reservation_fits: + if self._unified_swa_reservation_fits(total_tokens, swa_needed): return True, self.rem_chunk_tokens - full_ever_tokens = unified_swa_tokens = None - if self.is_unified_swa: - full_input_tokens = len(req.full_untruncated_fill_ids) - full_ever_tokens = full_input_tokens + max_new_tokens + self.page_size - unified_swa_tokens = self._swa_budget_for_req( - full_input_tokens, max_new_tokens - ) - if not self._swa_req_never_fits( - extend_input_len, - max_new_tokens, - req.swa_host_hit_length, - full_tokens=full_ever_tokens, - unified_swa_tokens=unified_swa_tokens, + # Only shrink requests that cannot fit even after the pool drains; + # transient pressure must wait instead of consuming decode headroom. + full_input_tokens = len(req.full_untruncated_fill_ids) + if self._unified_swa_reservation_fits( + full_input_tokens + max_new_tokens + self.page_size, + self._swa_budget_for_req(full_input_tokens, max_new_tokens), + empty_pool=True, ): return False, None - - max_chunk_tokens = self.rem_chunk_tokens - chunk_max_new = max_new_tokens - if self.is_unified_swa: - max_chunk_tokens = min(max_chunk_tokens or 0, max(0, extend_input_len - 1)) - chunk_max_new = 0 - swa_cap = self._swa_chunk_cap( - chunk_max_new, + swa_cap = self._unified_swa_chunk_cap( + 0, req.swa_host_hit_length, - max_chunk_tokens=max_chunk_tokens, + max_chunk_tokens=min( + self.rem_chunk_tokens or 0, max(0, extend_input_len - 1) + ), ) if self.rem_chunk_tokens is None or swa_cap <= 0: return False, None return True, min(self.rem_chunk_tokens, swa_cap) + def _check_swa_admission( + self, req: Req, *, extend_input_len: int + ) -> tuple[bool, Optional[int]]: + extend_input_len = self.ceil_paged_tokens(extend_input_len) + max_new_tokens = self._swa_new_tokens(req) + swa_needed = self._swa_budget_for_req( + extend_input_len, max_new_tokens, req.swa_host_hit_length + ) + # Ring-slot capacity is exact; the fixed token partition keeps slack. + if ( + swa_needed <= self.rem_swa_tokens + if self._swa_req_ring + else swa_needed < self.rem_swa_tokens + ): + return True, self.rem_chunk_tokens + if not self._swa_req_never_fits( + extend_input_len, max_new_tokens, req.swa_host_hit_length + ): + return False, None + swa_cap = self._swa_chunk_cap(max_new_tokens, req.swa_host_hit_length) + if self.rem_chunk_tokens is None or swa_cap <= 0: + return False, None + return True, min(self.rem_chunk_tokens, swa_cap) + + def _check_prefill_budget( + self, req: Req, *, extend_input_len: int, total_tokens: int + ) -> tuple[bool, Optional[int]]: + if self.is_unified_swa: + return self._check_unified_swa_admission( + req, extend_input_len=extend_input_len, total_tokens=total_tokens + ) + if total_tokens >= self.rem_total_tokens: + return False, None + if self.is_hybrid_swa: + return self._check_swa_admission(req, extend_input_len=extend_input_len) + return True, self.rem_chunk_tokens + def _mamba_gap_budget_for_req(self, req: Req) -> int: """Shared-gap reservation (full-token-equivalents) for a request's new mamba state. Charged only on the SHARED Mamba pool (`_mamba_slot_cost > 0`) @@ -1215,6 +1201,25 @@ def add_dllm_staging_req(self, req: Req): else AddReqResult.CONTINUE ) + def _unified_swa_chunk_limit( + self, req: Req, extend_input_len: int, max_chunk_tokens: int + ) -> Optional[int]: + candidate_tokens = min(extend_input_len, max_chunk_tokens) + candidate_finishes = candidate_tokens >= extend_input_len + max_new_tokens = self._swa_new_tokens(req) if candidate_finishes else 0 + if self._unified_swa_reservation_fits( + self.ceil_paged_tokens(candidate_tokens) + max_new_tokens + self.page_size, + self._swa_budget_for_req(candidate_tokens, max_new_tokens), + ): + return max_chunk_tokens + max_partial = ( + max(0, candidate_tokens - 1) if candidate_finishes else candidate_tokens + ) + joint_cap = self._unified_swa_chunk_cap(0, max_chunk_tokens=max_partial) + if joint_cap <= 0: + return None + return min(max_chunk_tokens, joint_cap) + def add_chunked_req(self, req: Req): if self.dllm_config is not None: _rem_tokens = self._get_dllm_remain_tokens() @@ -1223,17 +1228,11 @@ def add_chunked_req(self, req: Req): _rem_tokens = self.rem_chunk_tokens else: _rem_tokens = min(self.rem_chunk_tokens, int(self.rem_total_tokens)) - if ( - self.is_hybrid_swa - and not self.is_unified_swa - and not self._swa_req_ring - ): - # alloc_extend needs extend_num_tokens + page_size per request, - # so reserve one page here to avoid OOM. - # Ring mode skips it: rem_swa_tokens counts slots, not chunk tokens. - _rem_tokens = min( - _rem_tokens, int(self.rem_swa_tokens) - self.page_size - ) + if self.is_hybrid_swa and not self._swa_req_ring: + # alloc_extend needs one extra page; ring mode counts slots. + _rem_tokens = min( + _rem_tokens, int(self.rem_swa_tokens) - self.page_size + ) # The chunked_req must be added to the list; otherwise, it will cause a memory leak. # Therefore, in certain cases where _rem_tokens <= 0, it should be replaced with rem_chunk_tokens. if _rem_tokens <= 0: @@ -1256,22 +1255,11 @@ def add_chunked_req(self, req: Req): req.prefix_indices ) if self.is_unified_swa: - candidate_tokens = min(cand_extend_input_len, _rem_tokens) - candidate_finishes = candidate_tokens >= cand_extend_input_len - candidate_max_new = self._swa_new_tokens(req) if candidate_finishes else 0 - if not self._unified_swa_reservation_fits( - self.ceil_paged_tokens(candidate_tokens) - + candidate_max_new - + self.page_size, - self._swa_budget_for_req(candidate_tokens, candidate_max_new), - ): - max_partial = candidate_tokens - if candidate_finishes: - max_partial = max(0, candidate_tokens - 1) - joint_cap = self._swa_chunk_cap(0, max_chunk_tokens=max_partial) - if joint_cap <= 0: - return req - _rem_tokens = min(_rem_tokens, joint_cap) + _rem_tokens = self._unified_swa_chunk_limit( + req, cand_extend_input_len, _rem_tokens + ) + if _rem_tokens is None: + return req truncated = cand_extend_input_len > _rem_tokens new_len = min(cand_extend_input_len, _rem_tokens) req.set_extend_range(len(req.prefix_indices), len(req.prefix_indices) + new_len) @@ -1316,22 +1304,25 @@ def add_one_req_ignore_eos(self, req: Req): # Shared Mamba pool: fold the new mamba state's shared-gap cost into the # budget gate so admission can't over-commit (0 for baseline / non-Mamba). paged_input += self._mamba_gap_budget_for_req(req) - if self.is_hybrid_swa: + if self.is_unified_swa: max_new_tokens = self._swa_new_tokens(req) swa_needed = self._swa_budget_for_req(cand_extend_input_len, max_new_tokens) - if self.is_unified_swa: - if not self._unified_swa_reservation_fits( - paged_input + max_new_tokens + self.page_size, - swa_needed, - ): - return AddReqResult.NO_TOKEN - elif ( - paged_input > min(self.cur_rem_tokens, self.rem_total_tokens) - or swa_needed > self.rem_swa_tokens + if not self._unified_swa_reservation_fits( + paged_input + max_new_tokens + self.page_size, + swa_needed, + ): + return AddReqResult.NO_TOKEN + else: + if paged_input > min(self.cur_rem_tokens, self.rem_total_tokens): + return AddReqResult.NO_TOKEN + if ( + self.is_hybrid_swa + and self._swa_budget_for_req( + cand_extend_input_len, self._swa_new_tokens(req) + ) + > self.rem_swa_tokens ): return AddReqResult.NO_TOKEN - elif paged_input > min(self.cur_rem_tokens, self.rem_total_tokens): - return AddReqResult.NO_TOKEN def add_req_state(r, insert_sort=False): new_token_ratio = ( @@ -1482,19 +1473,14 @@ def add_one_req( real_input_tokens = self.ceil_paged_tokens(swa_extend_input_len) prefix_len = len(req.prefix_indices) - if not self.is_unified_swa and total_tokens >= self.rem_total_tokens: + can_admit, chunk_tokens_limit = self._check_prefill_budget( + req, + extend_input_len=swa_extend_input_len, + total_tokens=total_tokens, + ) + if not can_admit: return AddReqResult.NO_TOKEN - chunk_tokens_limit = self.rem_chunk_tokens - if self.is_hybrid_swa: - can_admit, chunk_tokens_limit = self._check_swa_admission( - req, - extend_input_len=swa_extend_input_len, - total_tokens=total_tokens, - ) - if not can_admit: - return AddReqResult.NO_TOKEN - if ( self.rem_chunk_tokens is None and len(self.can_run_list) != 0 @@ -1506,19 +1492,15 @@ def add_one_req( return AddReqResult.OTHER with self._lock_node(req.last_node): - # self.rem_total_tokens may decrease after the lock acquisition - if not self.is_unified_swa and total_tokens >= self.rem_total_tokens: + # Locking removes evictable capacity, so repeat the budget check. + can_admit, chunk_tokens_limit = self._check_prefill_budget( + req, + extend_input_len=swa_extend_input_len, + total_tokens=total_tokens, + ) + if not can_admit: return AddReqResult.NO_TOKEN - if self.is_hybrid_swa: - can_admit, chunk_tokens_limit = self._check_swa_admission( - req, - extend_input_len=swa_extend_input_len, - total_tokens=total_tokens, - ) - if not can_admit: - return AddReqResult.NO_TOKEN - # Negotiate only after every KV-budget gate (a NO_TOKEN rank must # report not-prefillable via finalize()) and before init_load_back # (a delay verdict must not start KV load-back). diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 7b7b8d375f52..e50ddf012669 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -281,7 +281,7 @@ ) from sglang.srt.mem_cache import kv_cache_builder from sglang.srt.mem_cache.allocator.unified_hybrid_swa import ( - UnifiedSWATokenToKVPoolAllocator, + supports_swa_byte_budget, ) from sglang.srt.mem_cache.common import ( maybe_cache_unfinished_req, @@ -2513,7 +2513,6 @@ def init_req_max_new_tokens(self, req): max_new_tokens = min(max_new_tokens, self.max_new_tokens_limit) # Keep this bound consistent with PrefillAdder's admission budget. - paged_input_len = -(-input_len // self.page_size) * self.page_size max_new_tokens = max( 0, min( @@ -2521,45 +2520,11 @@ def init_req_max_new_tokens(self, req): self.max_req_len - input_len - 1, ), ) - allocator = self.token_to_kv_pool_allocator - if ( - isinstance(allocator, UnifiedSWATokenToKVPoolAllocator) - and allocator.supports_asymmetric_reservation - ): - - def fits_shared_pool(candidate: int) -> bool: - full_tokens = input_len + candidate + self.page_size - swa_tokens = estimate_swa_kv_tokens( - input_len, - candidate, - sliding_window_size=self.sliding_window_size, - page_size=self.page_size, - allocation_limit=self.chunked_prefill_size, - ) - return allocator.can_reserve( - full_tokens, - swa_tokens, - empty_pool=True, - require_token_slack=True, - ) - - if not fits_shared_pool(0): - req.set_finish_with_abort( - "Request prompt exceeds the unified FULL/SWA KV byte budget: " - f"input_len={input_len}." - ) - max_new_tokens = 0 - elif not fits_shared_pool(max_new_tokens): - lo, hi = 0, max_new_tokens - while lo < hi: - mid = (lo + hi + 1) // 2 - if fits_shared_pool(mid): - lo = mid - else: - hi = mid - 1 - max_new_tokens = lo + if supports_swa_byte_budget(self.token_to_kv_pool_allocator): + max_new_tokens = self._clip_unified_swa_max_new_tokens(req, max_new_tokens) else: # Non-unified pools keep their fixed token partition. + paged_input_len = -(-input_len // self.page_size) * self.page_size max_new_tokens = min( max_new_tokens, self.max_total_num_tokens * get_parallel().attn_dcp_size @@ -2574,6 +2539,41 @@ def fits_shared_pool(candidate: int) -> bool: if req.sampling_params.min_new_tokens > req.sampling_params.max_new_tokens: req.sampling_params.min_new_tokens = req.sampling_params.max_new_tokens + def _clip_unified_swa_max_new_tokens(self, req: Req, max_new_tokens: int) -> int: + input_len = len(req.origin_input_ids) + + def fits_shared_pool(candidate: int) -> bool: + swa_tokens = estimate_swa_kv_tokens( + input_len, + candidate, + sliding_window_size=self.sliding_window_size, + page_size=self.page_size, + allocation_limit=self.chunked_prefill_size, + ) + return self.token_to_kv_pool_allocator.can_reserve( + input_len + candidate + self.page_size, + swa_tokens, + empty_pool=True, + require_token_slack=True, + ) + + if not fits_shared_pool(0): + req.set_finish_with_abort( + "Request prompt exceeds the unified FULL/SWA KV byte budget: " + f"input_len={input_len}." + ) + return 0 + if fits_shared_pool(max_new_tokens): + return max_new_tokens + lo, hi = 0, max_new_tokens + while lo < hi: + mid = (lo + hi + 1) // 2 + if fits_shared_pool(mid): + lo = mid + else: + hi = mid - 1 + return lo + def _process_and_broadcast_mm_inputs( self, raw_mm_inputs, diff --git a/python/sglang/srt/managers/scheduler_components/invariant_checker.py b/python/sglang/srt/managers/scheduler_components/invariant_checker.py index a3a95d0128d8..30742013ad54 100644 --- a/python/sglang/srt/managers/scheduler_components/invariant_checker.py +++ b/python/sglang/srt/managers/scheduler_components/invariant_checker.py @@ -22,10 +22,6 @@ ) from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator from sglang.srt.mem_cache.allocator.swa import is_swa_req_ring -from sglang.srt.mem_cache.allocator.unified_hybrid_swa import ( - UnifiedMambaSWATokenToKVPoolAllocator, - UnifiedSWATokenToKVPoolAllocator, -) from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache from sglang.srt.mem_cache.memory_pool import ReqToTokenPool from sglang.srt.observability.scheduler_stage_metrics import ( @@ -96,18 +92,12 @@ def _check_pool_invariant( def _check_full_pool(self, ps: PoolStats, uncached: int = 0) -> Tuple[bool, str]: allocator = self.token_to_kv_pool_allocator - full_capacity = ( - allocator.current_full_capacity - if isinstance(allocator, UnifiedSWATokenToKVPoolAllocator) - and allocator.supports_asymmetric_reservation - else self.full_tokens_per_layer - ) - if self.is_hybrid_swa and not full_capacity: + if self.is_hybrid_swa and not ps.full_capacity: return False, "" if self.is_hybrid_swa: protected = self.tree_cache.full_protected_size() session_held = self.pool_stats_observer.session_held_full_tokens() - total = full_capacity + total = ps.full_capacity elif self.is_hybrid_ssm: # Branch on cache type for the protected accessor (MambaRadixCache # splits full/mamba; ChunkCache only has the single protected_size). @@ -136,14 +126,9 @@ def _check_full_pool(self, ps: PoolStats, uncached: int = 0) -> Tuple[bool, str] // allocator.page_size * allocator.page_size ) - full_available = ps.full_available_size - if isinstance(allocator, UnifiedMambaSWATokenToKVPoolAllocator): - # Pair the static per-layer total with the conserve view, never the - # byte-coordinated one -- see `conserve_full_available_size`. - full_available = allocator.conserve_full_available_size() leak, msg = self._check_pool_invariant( "full", - full_available, + ps.full_available_size, full_evictable_size, protected, session_held, @@ -169,23 +154,13 @@ def _check_swa_pool(self, ps: PoolStats, uncached: int = 0) -> Tuple[bool, str]: f"evictable={ps.swa_evictable_size}, " f"total={self.swa_tokens_per_layer}" ) - swa_available = ps.swa_available_size - if isinstance(allocator, UnifiedMambaSWATokenToKVPoolAllocator): - # Tri-pool: same floating-boundary phantom as the full pool -- use the - # slot-conservation view, not the byte-coordinated min (see _check_full_pool). - swa_available = allocator.conserve_swa_available_size() return self._check_pool_invariant( "swa", - swa_available, + ps.swa_available_size, ps.swa_evictable_size, self.tree_cache.swa_protected_size(), self.pool_stats_observer.session_held_swa_tokens(), - ( - allocator.current_swa_capacity - if isinstance(allocator, UnifiedSWATokenToKVPoolAllocator) - and allocator.supports_asymmetric_reservation - else self.swa_tokens_per_layer - ), + ps.swa_capacity, uncached, ) diff --git a/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py b/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py index 1b8577b3b94a..38c6a5cc2797 100644 --- a/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py +++ b/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py @@ -14,7 +14,7 @@ from sglang.srt.mem_cache.allocator.swa import is_swa_req_ring from sglang.srt.mem_cache.allocator.unified_hybrid_swa import ( UnifiedMambaSWATokenToKVPoolAllocator, - UnifiedSWATokenToKVPoolAllocator, + supports_swa_byte_budget, ) if TYPE_CHECKING: @@ -39,6 +39,8 @@ class PoolStats: is_hisparse: bool = False # For hybrid-swa pools + full_capacity: Optional[int] = None + swa_capacity: Optional[int] = None swa_num_used: Optional[int] = None swa_token_usage: Optional[float] = None swa_available_size: Optional[int] = None @@ -293,13 +295,12 @@ def _get_swa_token_info(self) -> PoolStats: allocator = self.token_to_kv_pool_allocator full_capacity = self.full_tokens_per_layer swa_capacity = self.swa_tokens_per_layer - if ( - isinstance(allocator, UnifiedSWATokenToKVPoolAllocator) - and allocator.supports_asymmetric_reservation - ): + if supports_swa_byte_budget(allocator): full_capacity = allocator.current_full_capacity swa_capacity = allocator.current_swa_capacity - if isinstance(allocator, UnifiedMambaSWATokenToKVPoolAllocator): + full_available_size = allocator.full_available_size() + swa_available_size = allocator.swa_available_size() + elif isinstance(allocator, UnifiedMambaSWATokenToKVPoolAllocator): # The tri-pool reports static capacities paired with conserve views. full_available_size = allocator.conserve_full_available_size() swa_available_size = allocator.conserve_swa_available_size() @@ -336,6 +337,8 @@ def _get_swa_token_info(self) -> PoolStats: return PoolStats( is_hybrid_swa=True, + full_capacity=full_capacity, + swa_capacity=swa_capacity, full_num_used=full_num_used, full_token_usage=full_token_usage, full_available_size=full_available_size, diff --git a/python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py b/python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py index 20e85d94b5f7..092140575caf 100644 --- a/python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py +++ b/python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py @@ -18,7 +18,7 @@ import logging import math -from typing import Callable, List, Optional, Sequence, Tuple +from typing import Callable, List, Optional, Sequence, Tuple, TypeGuard import torch from torch.profiler import record_function @@ -42,6 +42,16 @@ logger = logging.getLogger(__name__) +def supports_swa_byte_budget( + allocator: BaseTokenToKVPoolAllocator | None, +) -> TypeGuard[UnifiedSWATokenToKVPoolAllocator]: + """Whether FULL/SWA demand can be checked against a two-pool byte budget.""" + return ( + isinstance(allocator, UnifiedSWATokenToKVPoolAllocator) + and allocator.supports_asymmetric_reservation + ) + + class UnifiedSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator): """Composite allocator for the hybrid SWA pair (full + swa MHA sub-pools). @@ -49,6 +59,8 @@ class UnifiedSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator): `available_size()` (joint bytes, in TOKENS) is the only safe alloc pre-check. """ + # Unequal FULL/SWA reservations share one byte budget. The tri-pool opts + # out because its Mamba state competes for those same bytes. supports_asymmetric_reservation = True # Parent's `size` property has no setter but base init does `self.size = size`; @@ -394,15 +406,6 @@ def can_reserve( """ if full_tokens < 0 or swa_tokens < 0: return False - if not self.supports_asymmetric_reservation: - if ( - full_tokens != swa_tokens - or full_evictable_tokens - or swa_evictable_tokens - or empty_pool - ): - return False - return full_tokens <= self.available_size() if require_token_slack and ( full_tokens >= self.size_full or swa_tokens >= self.size_swa ): @@ -503,13 +506,6 @@ def ensure_capacity(self, full_tokens: int, swa_tokens: int) -> bool: return False if full_tokens == 0 and swa_tokens == 0: return True - if not self.supports_asymmetric_reservation: - if full_tokens != swa_tokens: - return False - need_tokens = int(full_tokens) - if need_tokens <= self.available_size(): - return True - return _relieve_for_alloc(self, need_tokens) page_size = self.page_size num_full_pages = (int(full_tokens) + page_size - 1) // page_size num_swa_pages = (int(swa_tokens) + page_size - 1) // page_size @@ -1024,6 +1020,37 @@ def _wire_peers(self) -> None: # -- capacity -- + def can_reserve( + self, + full_tokens: int | float, + swa_tokens: int | float, + *, + full_evictable_tokens: int = 0, + swa_evictable_tokens: int = 0, + empty_pool: bool = False, + require_token_slack: bool = False, + ) -> bool: + if ( + full_tokens < 0 + or swa_tokens < 0 + or full_tokens != swa_tokens + or full_evictable_tokens + or swa_evictable_tokens + or empty_pool + ): + return False + return full_tokens <= self.available_size() + + def ensure_capacity(self, full_tokens: int, swa_tokens: int) -> bool: + if full_tokens < 0 or swa_tokens < 0 or full_tokens != swa_tokens: + return False + if full_tokens == 0: + return True + need_tokens = int(full_tokens) + if need_tokens <= self.available_size(): + return True + return _relieve_for_alloc(self, need_tokens) + def _compute_available_size(self) -> int: """Joint TOKENS for `alloc(N)`: N costs N full pages AND N swa pages, drawn from DIFFERENT bands -- full extends only into the high band, the float into diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index c0f96814bb1e..cf72e6c81865 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/python/sglang/srt/mem_cache/common.py @@ -161,14 +161,10 @@ def evict_from_tree_cache(tree_cache: BasePrefixCache | None, num_tokens: int): if isinstance(allocator, SWATokenToKVPoolAllocator): # Hybrid allocator from sglang.srt.mem_cache.allocator.unified_hybrid_swa import ( - UnifiedSWATokenToKVPoolAllocator, + supports_swa_byte_budget, ) - is_unified_allocator = ( - isinstance(allocator, UnifiedSWATokenToKVPoolAllocator) - and allocator.supports_asymmetric_reservation - ) - if is_unified_allocator: + if supports_swa_byte_budget(allocator): reclaim_plan = allocator.reclaim_plan( num_tokens, num_tokens, @@ -178,18 +174,23 @@ def evict_from_tree_cache(tree_cache: BasePrefixCache | None, num_tokens: int): if reclaim_plan is None: return full_num_tokens, swa_num_tokens = reclaim_plan - else: - full_num_tokens = max(0, num_tokens - allocator.full_available_size()) - swa_num_tokens = max(0, num_tokens - allocator.swa_available_size()) + if full_num_tokens or swa_num_tokens: + tree_cache.evict_for_alloc( + EvictParams( + num_tokens=full_num_tokens, swa_num_tokens=swa_num_tokens + ) + ) + # Realize any compaction assumed by the reclaim plan before callers + # inspect availability. + allocator.ensure_capacity(num_tokens, num_tokens) + return + full_num_tokens = max(0, num_tokens - allocator.full_available_size()) + swa_num_tokens = max(0, num_tokens - allocator.swa_available_size()) if full_num_tokens or swa_num_tokens: tree_cache.evict_for_alloc( EvictParams(num_tokens=full_num_tokens, swa_num_tokens=swa_num_tokens) ) - if is_unified_allocator: - # The reclaim plan may rely on compacting cross-side holes. Make - # that projected state real before callers inspect availability. - allocator.ensure_capacity(num_tokens, num_tokens) else: # Standard allocator: evict only the shortfall (mirrors the SWA arm) available_size = allocator.available_size() diff --git a/python/sglang/srt/model_executor/pool_configurator.py b/python/sglang/srt/model_executor/pool_configurator.py index 95d37a7ba8d9..49a449ebde80 100644 --- a/python/sglang/srt/model_executor/pool_configurator.py +++ b/python/sglang/srt/model_executor/pool_configurator.py @@ -781,6 +781,9 @@ def align_page_size(x: int) -> int: f"full_layer_tokens={full_tokens}, swa_layer_tokens={swa_tokens}" ) + return self._make_pool_config(full_tokens, swa_tokens) + + def _make_pool_config(self, full_tokens: int, swa_tokens: int) -> MemoryPoolConfig: return MemoryPoolConfig( max_total_num_tokens=full_tokens, full_max_total_num_tokens=full_tokens, @@ -907,16 +910,7 @@ def calculate_pool_sizes( f"Reduce --max-running-requests, lower SGLANG_SWA_EVICTION_INTERVAL, " f"or increase --mem-fraction-static." ) - return MemoryPoolConfig( - max_total_num_tokens=full_tokens, - full_max_total_num_tokens=full_tokens, - swa_max_total_num_tokens=swa_tokens, - unified_memory_pool_bytes=( - self._unified_pool_bytes(full_tokens, swa_tokens) - if self._enable_unified_memory - else None - ), - ) + return self._make_pool_config(full_tokens, swa_tokens) def calculate_pool_sizes_from_max_tokens( self, max_total_num_tokens: int, page_size: int @@ -924,17 +918,8 @@ def calculate_pool_sizes_from_max_tokens( # Constrained max_total goes to the full pool; SWA stays at its cap. swa_tokens = ceil_align(self._swa_cap, page_size) full_tokens = (max_total_num_tokens // page_size) * page_size - return MemoryPoolConfig( - max_total_num_tokens=full_tokens, - full_max_total_num_tokens=full_tokens, - swa_max_total_num_tokens=min(swa_tokens, max_total_num_tokens), - unified_memory_pool_bytes=( - self._unified_pool_bytes( - full_tokens, min(swa_tokens, max_total_num_tokens) - ) - if self._enable_unified_memory - else None - ), + return self._make_pool_config( + full_tokens, min(swa_tokens, max_total_num_tokens) ) From cb88f1f24c940a0436f418acafe71671a46a0c1a Mon Sep 17 00:00:00 2001 From: yhzhuang Date: Sat, 12 Sep 2026 21:51:14 -0700 Subject: [PATCH 32/55] Centralize decode SWA reservation routing --- python/sglang/srt/disaggregation/decode.py | 64 ++++++---------------- 1 file changed, 18 insertions(+), 46 deletions(-) diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index fe68d0692ba1..c3dfe287e527 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -76,7 +76,7 @@ from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator from sglang.srt.mem_cache.allocator.swa import is_swa_req_ring from sglang.srt.mem_cache.allocator.unified_hybrid_swa import ( - UnifiedSWATokenToKVPoolAllocator, + supports_swa_byte_budget, ) from sglang.srt.mem_cache.base_prefix_cache import ( BasePrefixCache, @@ -432,10 +432,11 @@ def _uses_swa_tail_prealloc(self) -> bool: ) def _supports_unified_swa_reservation(self) -> bool: - allocator = self.token_to_kv_pool_allocator - return ( - isinstance(allocator, UnifiedSWATokenToKVPoolAllocator) - and allocator.supports_asymmetric_reservation + return supports_swa_byte_budget(self.token_to_kv_pool_allocator) + + def _uses_swa_reservation(self) -> bool: + return self._uses_swa_tail_prealloc() or supports_swa_byte_budget( + self.scheduler.token_to_kv_pool_allocator ) def _unified_swa_reservation_fits( @@ -448,10 +449,6 @@ def _unified_swa_reservation_fits( empty_pool: bool = False, ) -> bool: allocator = self.token_to_kv_pool_allocator - if not self._supports_unified_swa_reservation(): - return True - assert isinstance(allocator, UnifiedSWATokenToKVPoolAllocator) - full_evictable = swa_evictable = 0 if not empty_pool: full_evictable = self._radix_full_evictable() @@ -484,7 +481,6 @@ def _prealloc_reservation_fits( *, full_allocatable_tokens: int, swa_allocatable_tokens: Optional[int], - uses_swa_tail_prealloc: bool, ) -> bool: if self._supports_unified_swa_reservation(): return self._unified_swa_reservation_fits( @@ -493,9 +489,8 @@ def _prealloc_reservation_fits( full_allocatable_tokens=full_allocatable_tokens, swa_allocatable_tokens=swa_allocatable_tokens, ) - if not uses_swa_tail_prealloc: + if swa_allocatable_tokens is None: return full_tokens <= full_allocatable_tokens - assert swa_allocatable_tokens is not None return ( full_tokens <= full_allocatable_tokens and swa_tokens <= swa_allocatable_tokens @@ -514,7 +509,6 @@ def _reclaim_swa_tail_capacity( allocator = self.token_to_kv_pool_allocator page_size = allocator.page_size if self._supports_unified_swa_reservation(): - assert isinstance(allocator, UnifiedSWATokenToKVPoolAllocator) full_required = ceil_align(full_len, page_size) swa_required = ceil_align(swa_tail_len, page_size) evict_from_tree_cache( @@ -931,16 +925,8 @@ def resume_retracted_reqs( # allocate memory resumed_reqs = [] indices_to_remove = set() - uses_swa_tail_prealloc = self._uses_swa_tail_prealloc() - uses_unified_swa = False - if not uses_swa_tail_prealloc: - allocator = self.scheduler.token_to_kv_pool_allocator - uses_unified_swa = ( - isinstance(allocator, UnifiedSWATokenToKVPoolAllocator) - and allocator.supports_asymmetric_reservation - ) - uses_swa_reservation = uses_swa_tail_prealloc or uses_unified_swa - if uses_swa_reservation: + swa_allocatable_tokens = None + if self._uses_swa_reservation(): full_allocatable_tokens, swa_allocatable_tokens = ( self._swa_aware_allocatable_token_budgets(count_retracted=False) ) @@ -961,10 +947,7 @@ def resume_retracted_reqs( full_required, swa_required, full_allocatable_tokens=full_allocatable_tokens, - swa_allocatable_tokens=( - swa_allocatable_tokens if uses_swa_reservation else None - ), - uses_swa_tail_prealloc=uses_swa_tail_prealloc, + swa_allocatable_tokens=swa_allocatable_tokens, ): break @@ -981,7 +964,7 @@ def resume_retracted_reqs( req.is_retracted = False self._pre_alloc(req) full_allocatable_tokens -= full_required - if uses_swa_reservation: + if swa_allocatable_tokens is not None: swa_allocatable_tokens = self._swa_tail_allocatable_token_budget( count_retracted=False, extra_reserved_reqs=len(resumed_reqs), @@ -1242,16 +1225,8 @@ def pop_preallocated( ) uses_swa_tail_prealloc = self._uses_swa_tail_prealloc() - uses_unified_swa = False - if not uses_swa_tail_prealloc: - allocator = self.scheduler.token_to_kv_pool_allocator - uses_unified_swa = ( - isinstance(allocator, UnifiedSWATokenToKVPoolAllocator) - and allocator.supports_asymmetric_reservation - ) - uses_swa_reservation = uses_swa_tail_prealloc or uses_unified_swa - swa_allocatable_tokens = 0 - if uses_swa_reservation: + swa_allocatable_tokens = None + if self._uses_swa_reservation(): retractable_swa_tokens = sum( self._swa_retractable_len(r) for r in self.scheduler.running_batch.reqs ) @@ -1435,23 +1410,20 @@ def pop_preallocated( swa_required, swa_len + max_new_tokens - retractable_swa_tokens, ) - elif uses_unified_swa: + elif swa_allocatable_tokens is not None: swa_required_for_admission = full_required_for_admission if not self._prealloc_reservation_fits( full_required_for_admission, swa_required_for_admission, full_allocatable_tokens=full_allocatable_tokens, - swa_allocatable_tokens=( - swa_allocatable_tokens if uses_swa_reservation else None - ), - uses_swa_tail_prealloc=uses_swa_tail_prealloc, + swa_allocatable_tokens=swa_allocatable_tokens, ): if prefix_match is not None and prefix_match.l1_prefix_len > 0: self._release_matched_prefix_lock(decode_req.req) break - if uses_swa_tail_prealloc or uses_unified_swa: + if swa_allocatable_tokens is not None: reclaim_error = self._reclaim_swa_tail_capacity( swa_len, decode_req.req.rid, @@ -1494,7 +1466,7 @@ def pop_preallocated( extra_reserved_reqs=len(preallocated_reqs) + 1, hicache_reserved_tokens=reserved_restore_tokens, ) - if uses_swa_reservation: + if swa_allocatable_tokens is not None: swa_allocatable_tokens = self._swa_tail_allocatable_token_budget( retractable_tokens=retractable_tokens, retractable_swa_tokens=retractable_swa_tokens, @@ -1812,7 +1784,7 @@ def _allocatable_token_budgets( # HiSparse pre-alloc only allocates logical indices, so the # logical pool is the binding constraint for admission control. available_size = logical_allocator.available_size() - elif self._uses_swa_tail_prealloc() or self._supports_unified_swa_reservation(): + elif self._uses_swa_reservation(): available_size = self.token_to_kv_pool_allocator.full_available_size() if get_disagg().disaggregation_decode_enable_radix_cache: available_size += self._radix_full_evictable() From 29545f8182389539c5f77fc356991dac834246b2 Mon Sep 17 00:00:00 2001 From: yhzhuang Date: Sat, 12 Sep 2026 22:51:55 -0700 Subject: [PATCH 33/55] Simplify shared SWA budget paths --- python/sglang/srt/managers/schedule_policy.py | 37 ++++++++----------- .../pool_stats_observer.py | 4 +- python/sglang/srt/mem_cache/common.py | 23 +++++------- .../sglang/srt/model_executor/model_runner.py | 15 +++----- .../unit/managers/test_prefill_adder.py | 2 +- 5 files changed, 33 insertions(+), 48 deletions(-) diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index de727a9f1b73..72f22d32b677 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -813,13 +813,13 @@ def _swa_budget_for_req( if self._swa_req_ring: # One ring slot per request, in the same unit as swa_available_size. return allocator.swa_ring_cost_tokens - if self.rem_chunk_tokens is not None: - alloc = min(extend_input_len, self.rem_chunk_tokens) - else: - alloc = extend_input_len - window = self.tree_cache.sliding_window_size - return max(alloc - window, 0) + self._swa_reserved_tokens( - extend_input_len, max_new_tokens, swa_host_hit_length + return estimate_swa_kv_tokens( + extend_input_len, + max_new_tokens, + sliding_window_size=self.tree_cache.sliding_window_size, + page_size=self.page_size, + allocation_limit=self.rem_chunk_tokens, + host_hit_length=swa_host_hit_length, ) def _swa_reserved_tokens( @@ -1304,25 +1304,20 @@ def add_one_req_ignore_eos(self, req: Req): # Shared Mamba pool: fold the new mamba state's shared-gap cost into the # budget gate so admission can't over-commit (0 for baseline / non-Mamba). paged_input += self._mamba_gap_budget_for_req(req) - if self.is_unified_swa: + if self.is_hybrid_swa: max_new_tokens = self._swa_new_tokens(req) swa_needed = self._swa_budget_for_req(cand_extend_input_len, max_new_tokens) - if not self._unified_swa_reservation_fits( + if self.is_unified_swa: + fits = self._unified_swa_reservation_fits( paged_input + max_new_tokens + self.page_size, swa_needed, - ): - return AddReqResult.NO_TOKEN + ) else: - if paged_input > min(self.cur_rem_tokens, self.rem_total_tokens): - return AddReqResult.NO_TOKEN - if ( - self.is_hybrid_swa - and self._swa_budget_for_req( - cand_extend_input_len, self._swa_new_tokens(req) - ) - > self.rem_swa_tokens - ): - return AddReqResult.NO_TOKEN + fits = paged_input <= min(self.cur_rem_tokens, self.rem_total_tokens) and ( + not self.is_hybrid_swa or swa_needed <= self.rem_swa_tokens + ) + if not fits: + return AddReqResult.NO_TOKEN def add_req_state(r, insert_sort=False): new_token_ratio = ( diff --git a/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py b/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py index 38c6a5cc2797..138663ef9127 100644 --- a/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py +++ b/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py @@ -298,9 +298,7 @@ def _get_swa_token_info(self) -> PoolStats: if supports_swa_byte_budget(allocator): full_capacity = allocator.current_full_capacity swa_capacity = allocator.current_swa_capacity - full_available_size = allocator.full_available_size() - swa_available_size = allocator.swa_available_size() - elif isinstance(allocator, UnifiedMambaSWATokenToKVPoolAllocator): + if isinstance(allocator, UnifiedMambaSWATokenToKVPoolAllocator): # The tri-pool reports static capacities paired with conserve views. full_available_size = allocator.conserve_full_available_size() swa_available_size = allocator.conserve_swa_available_size() diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index cf72e6c81865..807623207300 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/python/sglang/srt/mem_cache/common.py @@ -164,7 +164,8 @@ def evict_from_tree_cache(tree_cache: BasePrefixCache | None, num_tokens: int): supports_swa_byte_budget, ) - if supports_swa_byte_budget(allocator): + is_unified_swa = supports_swa_byte_budget(allocator) + if is_unified_swa: reclaim_plan = allocator.reclaim_plan( num_tokens, num_tokens, @@ -174,23 +175,17 @@ def evict_from_tree_cache(tree_cache: BasePrefixCache | None, num_tokens: int): if reclaim_plan is None: return full_num_tokens, swa_num_tokens = reclaim_plan - if full_num_tokens or swa_num_tokens: - tree_cache.evict_for_alloc( - EvictParams( - num_tokens=full_num_tokens, swa_num_tokens=swa_num_tokens - ) - ) - # Realize any compaction assumed by the reclaim plan before callers - # inspect availability. - allocator.ensure_capacity(num_tokens, num_tokens) - return - - full_num_tokens = max(0, num_tokens - allocator.full_available_size()) - swa_num_tokens = max(0, num_tokens - allocator.swa_available_size()) + else: + full_num_tokens = max(0, num_tokens - allocator.full_available_size()) + swa_num_tokens = max(0, num_tokens - allocator.swa_available_size()) if full_num_tokens or swa_num_tokens: tree_cache.evict_for_alloc( EvictParams(num_tokens=full_num_tokens, swa_num_tokens=swa_num_tokens) ) + if is_unified_swa: + # Even a zero-reclaim plan can depend on compaction before callers + # inspect availability. + allocator.ensure_capacity(num_tokens, num_tokens) else: # Standard allocator: evict only the shortfall (mirrors the SWA arm) available_size = allocator.available_size() diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 27fd702b4535..da1b5c16bd96 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -1369,15 +1369,12 @@ def unload_lora_adapter(self, lora_ref: LoRARef): @property def effective_max_total_num_tokens(self): """Return the max token pool size considering hybrid swa settings.""" - if self.is_hybrid_swa: - if get_memory().enable_unified_memory: - capacity = self.token_to_kv_pool_allocator.size_full - if get_schedule().max_total_tokens is not None: - capacity = min(capacity, get_schedule().max_total_tokens) - else: - capacity = ( - self.full_max_total_num_tokens or self.swa_max_total_num_tokens - ) + if self.is_hybrid_swa and get_memory().enable_unified_memory: + capacity = self.token_to_kv_pool_allocator.size_full + if get_schedule().max_total_tokens is not None: + capacity = min(capacity, get_schedule().max_total_tokens) + elif self.is_hybrid_swa: + capacity = self.full_max_total_num_tokens or self.swa_max_total_num_tokens else: capacity = self.max_total_num_tokens if (req_to_token_pool := getattr(self, "req_to_token_pool", None)) is not None: diff --git a/test/registered/unit/managers/test_prefill_adder.py b/test/registered/unit/managers/test_prefill_adder.py index c881f612818c..c3c1dd1cf4a0 100644 --- a/test/registered/unit/managers/test_prefill_adder.py +++ b/test/registered/unit/managers/test_prefill_adder.py @@ -636,7 +636,7 @@ def test_swa_admission_admits_short_cached_resume_at_two_window_pool(self): req.sampling_params = SimpleNamespace(max_new_tokens=40, ignore_eos=False) # Pre-fix: a constant sliding-window reservation rejects the resume. - with patch.object(adder, "_swa_reserved_tokens", return_value=WINDOW + PAGE): + with patch.object(adder, "_swa_budget_for_req", return_value=WINDOW + PAGE): self.assertIs( adder.add_one_req( req, has_chunked_req=False, truncation_align_size=None From 4d1fa0ff30f7f17c9b4003620af71187d03d48ad Mon Sep 17 00:00:00 2001 From: yhzhuang Date: Sun, 13 Sep 2026 00:57:32 -0700 Subject: [PATCH 34/55] Fix unified capturer indexing and asymmetric SWA reservation --- .../mem_cache/allocator/unified_hybrid_swa.py | 6 +- .../srt/mem_cache/kv_index_translator.py | 10 ++++ .../sglang/srt/model_executor/model_runner.py | 8 ++- python/sglang/srt/state_capturer/base.py | 12 +++- .../mem_cache/test_kv_index_translator.py | 56 ++++++++++++++++++- .../mem_cache/test_multi_ended_allocator.py | 21 ++++++- 6 files changed, 103 insertions(+), 10 deletions(-) diff --git a/python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py b/python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py index 092140575caf..b1218ac84814 100644 --- a/python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py +++ b/python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py @@ -336,8 +336,8 @@ def reclaim_plan( page_size = self.page_size full_pages = (math.ceil(full_tokens) + page_size - 1) // page_size swa_pages = (math.ceil(swa_tokens) + page_size - 1) // page_size - if swa_pages > full_pages: - return None + # Restoring host SWA for device-resident FULL can require more new + # SWA pages than new FULL pages; only the shared budget constrains it. compacted = empty_pool or not self.lazy_compaction or self._compaction_allowed() @@ -509,8 +509,6 @@ def ensure_capacity(self, full_tokens: int, swa_tokens: int) -> bool: page_size = self.page_size num_full_pages = (int(full_tokens) + page_size - 1) // page_size num_swa_pages = (int(swa_tokens) + page_size - 1) // page_size - if num_swa_pages > num_full_pages: - return False if self._fits_page_demand( num_full_pages, num_swa_pages, diff --git a/python/sglang/srt/mem_cache/kv_index_translator.py b/python/sglang/srt/mem_cache/kv_index_translator.py index 2e1b970680fa..1ccb5b35077c 100644 --- a/python/sglang/srt/mem_cache/kv_index_translator.py +++ b/python/sglang/srt/mem_cache/kv_index_translator.py @@ -171,6 +171,16 @@ def __init__( ) self._index_table_memo: Optional[Tuple[weakref.ref, KVIndexTable]] = None + def capture_token_capacity(self, max_token_pool_size: int) -> int: + """Host capture rows are indexed by request-token IDs, not kernel IDs. + + Unified IDs span the whole virtual table even when admission is capped. + The logical page size also covers DCP-widened request-token IDs. + """ + if self.is_translating: + return self._full_v2p_table.numel() * self.page_size + return max_token_pool_size + self.page_size + # -- per-batch view -------------------------------------------------------- @property diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index da1b5c16bd96..aa4baf45025c 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -1114,7 +1114,9 @@ def init_routed_experts_capturer(self): RoutedExpertsCapturer.create( model=self.model, model_config=self.model_config, - num_tokens=self.max_token_pool_size + self.page_size, + num_tokens=self.kv_index_translator.capture_token_capacity( + self.max_token_pool_size + ), max_running_requests=self.max_running_requests, device=self.device, ) @@ -1124,7 +1126,9 @@ def init_indexer_capturer(self): set_global_indexer_capturer( create_indexer_capturer( model_config=self.model_config, - num_tokens=self.max_token_pool_size + self.page_size, + num_tokens=self.kv_index_translator.capture_token_capacity( + self.max_token_pool_size + ), max_running_requests=self.max_running_requests, device=self.device, ) diff --git a/python/sglang/srt/state_capturer/base.py b/python/sglang/srt/state_capturer/base.py index 6f6b8757841b..f451f7e6f5d2 100644 --- a/python/sglang/srt/state_capturer/base.py +++ b/python/sglang/srt/state_capturer/base.py @@ -201,13 +201,21 @@ def on_forward_end( slice_gpu = self._get_local_slice( forward_batch, can_run_graph, cuda_graph_batch ) + # get_topk reads req_to_token IDs; attention may have rebound the write + # loc to kernel-facing IDs, which are neither stable nor the same space. + out_cache_loc = forward_batch.out_cache_loc_virtual + if out_cache_loc is None: + out_cache_loc = forward_batch.out_cache_loc + else: + # Kernel batches may be padded; only real request tokens are stored. + slice_gpu = slice_gpu[: out_cache_loc.shape[0]] if no_copy_to_cpu: # Clone before the next overlapping forward reuses these buffers. return TopkCaptureOutput( - out_cache_loc=forward_batch.out_cache_loc.clone(), + out_cache_loc=out_cache_loc.clone(), topk=slice_gpu.clone(), host_cache=self.host_cache, ) - out_cache_loc_cpu = forward_batch.out_cache_loc.cpu() + out_cache_loc_cpu = out_cache_loc.cpu() self.host_cache.buffer[out_cache_loc_cpu] = slice_gpu.cpu() return None diff --git a/test/registered/unit/mem_cache/test_kv_index_translator.py b/test/registered/unit/mem_cache/test_kv_index_translator.py index a5e818e24143..fb497732be86 100644 --- a/test/registered/unit/mem_cache/test_kv_index_translator.py +++ b/test/registered/unit/mem_cache/test_kv_index_translator.py @@ -35,6 +35,8 @@ from sglang.srt.mem_cache.kv_index_translator import KVIndexTranslator, KVReadTables from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool from sglang.srt.mem_cache.unified_memory_pool import MHASubPoolSpec, UnifiedKVPool +from sglang.srt.state_capturer.base import BaseTopkCapturer +from sglang.test.test_utils import CustomTestCase _DEV = "cpu" _FULL_L = 2 @@ -135,6 +137,7 @@ def test_non_unified_returns_same_objects(self): device=_DEV, ) self.assertFalse(src.is_translating) + self.assertEqual(src.capture_token_capacity(17), 18) rows = torch.tensor([2, 0]) view = src.build_index_table( req_pool_indices=rows, seq_lens=torch.tensor([5, 3]) @@ -594,7 +597,7 @@ def test_dead_batch_never_matches(self): self.assertEqual(v2.ids.shape[0], 1) -class TestWriteLoc(unittest.TestCase): +class TestWriteLoc(CustomTestCase): """The two-phase write contract: `rebind_write_loc` rebinds the full side once at ForwardBatch construction, and the sliding-window write loc derives POINTWISE from the full-side values -- pads, slices, and fresh copies @@ -630,6 +633,57 @@ def test_rebind_translates_full_side_only(self): self.assertTrue(torch.equal(fb.out_cache_loc, want_full)) self.assertTrue(torch.equal(virt, keep)) + def test_topk_capture_round_trips_request_token_ids(self): + for ps in (1, 4, 64): + for translating in (False, True): + for overlap in (False, True): + with self.subTest( + page_size=ps, translating=translating, overlap=overlap + ): + src, _, _, _, virt, _, _ = self._built(ps=ps, n=3 * ps) + req_pool = SimpleNamespace(req_to_token=virt.clone()[None, :]) + fb = _FakeForwardBatch(out_cache_loc=virt.clone()) + if translating: + src.rebind_write_loc(fb) + fb.out_cache_loc = torch.cat( + [fb.out_cache_loc, virt.new_zeros(2)] + ) + else: + fb.out_cache_loc_virtual = None + + # Admission may be capped below IDs issued after reuse. + capacity = src.capture_token_capacity(ps) + self.assertGreater(capacity, int(virt.max())) + expected = ( + torch.arange(len(virt) * 4, dtype=torch.int32).reshape( + -1, 2, 2 + ) + + 1 + ) + cap = object.__new__(BaseTopkCapturer) + cap.topk_size = 2 + cap.device_cache = SimpleNamespace( + buffer=torch.cat([expected, expected.new_zeros(2, 2, 2)]) + ) + cap.host_cache = SimpleNamespace( + buffer=torch.zeros(capacity, 2, 2, dtype=torch.int32) + ) + result = cap.on_forward_end( + fb, False, None, no_copy_to_cpu=overlap + ) + if overlap: + fb.out_cache_loc.zero_() + if translating: + fb.out_cache_loc_virtual.zero_() + cap.device_cache.buffer.zero_() + result.map_device_tensors(lambda value: value.cpu()) + result.finalize() + self.assertTrue( + torch.equal( + cap.get_topk(0, len(virt) + 1, req_pool), expected + ) + ) + def test_swa_write_loc_round_trips_from_full_side(self): """Derived property: `field(full(t)) == swa(t)` for any virtual run t, across page sizes and multipliers.""" diff --git a/test/registered/unit/mem_cache/test_multi_ended_allocator.py b/test/registered/unit/mem_cache/test_multi_ended_allocator.py index bac0687f9f1d..6f677c94f64e 100644 --- a/test/registered/unit/mem_cache/test_multi_ended_allocator.py +++ b/test/registered/unit/mem_cache/test_multi_ended_allocator.py @@ -625,7 +625,7 @@ def test_reclaim_plan_matches_exhaustive_page_targets(self): max_full = min(12, full_evictable) // page_size max_swa = min(8, swa_evictable) // page_size for full_pages in range(9): - for swa_pages in range(full_pages + 1): + for swa_pages in range(9): feasible = [ (full * page_size, swa * page_size) for swa in range(max_swa + 1) @@ -653,6 +653,25 @@ def test_reclaim_plan_matches_exhaustive_page_targets(self): feasible[0] if feasible else None, ) + def test_restore_swa_without_allocating_more_full(self): + _, allocator, _ = self._build(page_size=4) + indices = allocator.alloc(8) + full_before = allocator.translate_kv_indices_for_transfer(indices).clone() + allocator.free_swa(indices) + + self.assertEqual(allocator.reclaim_plan(0, 8), (0, 0)) + self.assertTrue(allocator.can_reserve(0, 8)) + self.assertTrue(allocator.ensure_capacity(0, 8)) + allocator.swa_attn_allocator.alloc_with_virtual((indices // 4).unique()) + self.assertTrue( + torch.equal( + allocator.translate_kv_indices_for_transfer(indices), full_before + ) + ) + self.assertTrue( + bool((allocator.swa_attn_allocator.translate_kv_loc(indices) > 0).all()) + ) + def test_empty_pool_reservation_matches_packed_byte_boundary(self): page_size = 4 _, allocator, _ = self._build( From 567f4f5d0f1347f940365ae2e37b2b73af0434b4 Mon Sep 17 00:00:00 2001 From: yhzhuang Date: Sun, 13 Sep 2026 01:34:47 -0700 Subject: [PATCH 35/55] Consolidate unified allocation, admission and capacity policies --- python/sglang/srt/managers/schedule_policy.py | 115 ++++++++---------- .../pool_stats_observer.py | 19 ++- .../mem_cache/allocator/unified_hybrid_swa.py | 25 +++- python/sglang/srt/mem_cache/common.py | 24 +--- .../srt/mem_cache/kv_cache_configurator.py | 22 +++- .../srt/mem_cache/unified_memory_pool.py | 6 - .../sglang/srt/model_executor/model_runner.py | 12 +- .../test_hisparse_max_token_pool_size.py | 2 + .../mem_cache/test_multi_ended_allocator.py | 2 +- .../test_unified_byte_budget_sizing.py | 12 +- .../test_unified_radix_cache_unittest.py | 2 +- 11 files changed, 126 insertions(+), 115 deletions(-) diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index 72f22d32b677..24110fea767a 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -88,7 +88,6 @@ def estimate_swa_kv_tokens( page_size: int, allocation_limit: Optional[int] = None, host_hit_length: int = 0, - include_allocated_tail: bool = True, ) -> int: """Peak SWA reservation for one prefill/decode request.""" if sliding_window_size is None or sliding_window_size <= 0: @@ -99,9 +98,9 @@ def estimate_swa_kv_tokens( if allocation_limit is None else min(extend_input_len, allocation_limit) ) - allocated_tail = ( - max(allocated - sliding_window_size, 0) if include_allocated_tail else 0 - ) + allocated_tail = max(allocated - sliding_window_size, 0) + # A full window would double-charge short cached-prefix resumes; + # including extend keeps the reservation above the prefill allocation. reserved = ( allocated_tail + min(extend_input_len + max_new_tokens, sliding_window_size) @@ -822,18 +821,6 @@ def _swa_budget_for_req( host_hit_length=swa_host_hit_length, ) - def _swa_reserved_tokens( - self, extend_input_len: int, max_new_tokens: int, swa_host_hit_length: int = 0 - ) -> int: - return estimate_swa_kv_tokens( - extend_input_len, - max_new_tokens, - sliding_window_size=self.tree_cache.sliding_window_size, - page_size=self.page_size, - host_hit_length=swa_host_hit_length, - include_allocated_tail=False, - ) - def _swa_new_tokens(self, req: Req) -> int: """Tokens a request may still decode, for SWA headroom sizing. Mirrors the remaining-then-clip order of add_one_req's max_new: clip-then-subtract @@ -894,8 +881,14 @@ def _unified_swa_chunk_cap( def _swa_chunk_cap(self, max_new_tokens: int, swa_host_hit_length: int = 0) -> int: """Largest chunk fitting the fixed SWA partition after decode headroom.""" - cap = int(self.rem_swa_tokens) - self._swa_reserved_tokens( - 0, max_new_tokens, swa_host_hit_length + # Only the sliding window stays locked across chunks, so the next + # chunk can reuse the transient space outside its decode headroom. + cap = int(self.rem_swa_tokens) - estimate_swa_kv_tokens( + 0, + max_new_tokens, + sliding_window_size=self.tree_cache.sliding_window_size, + page_size=self.page_size, + host_hit_length=swa_host_hit_length, ) return max(0, cap // self.page_size * self.page_size) @@ -922,62 +915,52 @@ def _swa_req_never_fits( return swa_tokens > self.token_to_kv_pool_allocator.size_swa return swa_tokens >= self.token_to_kv_pool_allocator.size_swa - def _check_unified_swa_admission( + def _check_swa_admission( self, req: Req, *, extend_input_len: int, total_tokens: int, ) -> tuple[bool, Optional[int]]: + if not self.is_unified_swa: + extend_input_len = self.ceil_paged_tokens(extend_input_len) max_new_tokens = self._swa_new_tokens(req) swa_needed = self._swa_budget_for_req( extend_input_len, max_new_tokens, swa_host_hit_length=req.swa_host_hit_length, ) - if self._unified_swa_reservation_fits(total_tokens, swa_needed): - return True, self.rem_chunk_tokens - - # Only shrink requests that cannot fit even after the pool drains; - # transient pressure must wait instead of consuming decode headroom. - full_input_tokens = len(req.full_untruncated_fill_ids) - if self._unified_swa_reservation_fits( - full_input_tokens + max_new_tokens + self.page_size, - self._swa_budget_for_req(full_input_tokens, max_new_tokens), - empty_pool=True, - ): - return False, None - swa_cap = self._unified_swa_chunk_cap( - 0, - req.swa_host_hit_length, - max_chunk_tokens=min( - self.rem_chunk_tokens or 0, max(0, extend_input_len - 1) - ), - ) - if self.rem_chunk_tokens is None or swa_cap <= 0: - return False, None - return True, min(self.rem_chunk_tokens, swa_cap) - - def _check_swa_admission( - self, req: Req, *, extend_input_len: int - ) -> tuple[bool, Optional[int]]: - extend_input_len = self.ceil_paged_tokens(extend_input_len) - max_new_tokens = self._swa_new_tokens(req) - swa_needed = self._swa_budget_for_req( - extend_input_len, max_new_tokens, req.swa_host_hit_length - ) - # Ring-slot capacity is exact; the fixed token partition keeps slack. - if ( - swa_needed <= self.rem_swa_tokens - if self._swa_req_ring - else swa_needed < self.rem_swa_tokens - ): + swa_cap = 0 + # Only permanent shortfalls may shrink a chunk; transient pressure waits. + if self.is_unified_swa: + fits = self._unified_swa_reservation_fits(total_tokens, swa_needed) + if not fits: + full_input_tokens = len(req.full_untruncated_fill_ids) + if not self._unified_swa_reservation_fits( + full_input_tokens + max_new_tokens + self.page_size, + self._swa_budget_for_req(full_input_tokens, max_new_tokens), + empty_pool=True, + ): + swa_cap = self._unified_swa_chunk_cap( + 0, + req.swa_host_hit_length, + max_chunk_tokens=min( + self.rem_chunk_tokens or 0, max(0, extend_input_len - 1) + ), + ) + else: + # Ring-slot capacity is exact; the fixed token partition keeps slack. + fits = ( + swa_needed <= self.rem_swa_tokens + if self._swa_req_ring + else swa_needed < self.rem_swa_tokens + ) + if not fits and self._swa_req_never_fits( + extend_input_len, max_new_tokens, req.swa_host_hit_length + ): + swa_cap = self._swa_chunk_cap(max_new_tokens, req.swa_host_hit_length) + if fits: return True, self.rem_chunk_tokens - if not self._swa_req_never_fits( - extend_input_len, max_new_tokens, req.swa_host_hit_length - ): - return False, None - swa_cap = self._swa_chunk_cap(max_new_tokens, req.swa_host_hit_length) if self.rem_chunk_tokens is None or swa_cap <= 0: return False, None return True, min(self.rem_chunk_tokens, swa_cap) @@ -985,14 +968,12 @@ def _check_swa_admission( def _check_prefill_budget( self, req: Req, *, extend_input_len: int, total_tokens: int ) -> tuple[bool, Optional[int]]: - if self.is_unified_swa: - return self._check_unified_swa_admission( - req, extend_input_len=extend_input_len, total_tokens=total_tokens - ) - if total_tokens >= self.rem_total_tokens: + if not self.is_unified_swa and total_tokens >= self.rem_total_tokens: return False, None if self.is_hybrid_swa: - return self._check_swa_admission(req, extend_input_len=extend_input_len) + return self._check_swa_admission( + req, extend_input_len=extend_input_len, total_tokens=total_tokens + ) return True, self.rem_chunk_tokens def _mamba_gap_budget_for_req(self, req: Req) -> int: diff --git a/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py b/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py index 138663ef9127..823fe15b67cf 100644 --- a/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py +++ b/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py @@ -291,13 +291,13 @@ def _get_mamba_token_info(self): mamba_evictable_size=mamba_evictable_size, ) - def _get_swa_token_info(self) -> PoolStats: + def _swa_capacity_and_available(self) -> tuple[tuple[int, int], tuple[int, int]]: allocator = self.token_to_kv_pool_allocator - full_capacity = self.full_tokens_per_layer - swa_capacity = self.swa_tokens_per_layer if supports_swa_byte_budget(allocator): - full_capacity = allocator.current_full_capacity - swa_capacity = allocator.current_swa_capacity + return ( + (allocator.current_full_capacity, allocator.full_available_size()), + (allocator.current_swa_capacity, allocator.swa_available_size()), + ) if isinstance(allocator, UnifiedMambaSWATokenToKVPoolAllocator): # The tri-pool reports static capacities paired with conserve views. full_available_size = allocator.conserve_full_available_size() @@ -305,6 +305,15 @@ def _get_swa_token_info(self) -> PoolStats: else: full_available_size = allocator.full_available_size() swa_available_size = allocator.swa_available_size() + return ( + (self.full_tokens_per_layer, full_available_size), + (self.swa_tokens_per_layer, swa_available_size), + ) + + def _get_swa_token_info(self) -> PoolStats: + (full_capacity, full_available_size), (swa_capacity, swa_available_size) = ( + self._swa_capacity_and_available() + ) full_evictable_size = self.tree_cache.full_evictable_size() swa_evictable_size = self.tree_cache.swa_evictable_size() # Per-request SWA ring: released with the req slot, yet cached radix diff --git a/python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py b/python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py index b1218ac84814..3a7136d9da32 100644 --- a/python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py +++ b/python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py @@ -523,14 +523,35 @@ def ensure_capacity(self, full_tokens: int, swa_tokens: int) -> bool: compacted=True, ): return False - self.full_attn_allocator._flush(urgent=True) - self.swa_attn_allocator._flush(urgent=True) + self.full_attn_allocator.flush_for_allocation() + self.swa_attn_allocator.flush_for_allocation() return self._fits_page_demand( num_full_pages, num_swa_pages, compacted=False, ) + def evict_to_free_tokens(self, tree_cache, num_tokens: int) -> None: + from sglang.srt.mem_cache.base_prefix_cache import EvictParams + + if tree_cache is None or tree_cache.is_chunk_cache(): + return + reclaim_plan = self.reclaim_plan( + num_tokens, + num_tokens, + full_evictable_tokens=tree_cache.full_evictable_size(), + swa_evictable_tokens=tree_cache.swa_evictable_size(), + ) + if reclaim_plan is None: + return + full_reclaim, swa_reclaim = reclaim_plan + if full_reclaim or swa_reclaim: + tree_cache.evict_for_alloc( + EvictParams(num_tokens=full_reclaim, swa_num_tokens=swa_reclaim) + ) + # A zero-reclaim plan can still depend on compaction before allocation. + self.ensure_capacity(num_tokens, num_tokens) + @property def draft_virtual_id_space(self) -> int: return self.full_attn_allocator.max_slots - 1 diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index 807623207300..16f1e557c90e 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/python/sglang/srt/mem_cache/common.py @@ -164,28 +164,16 @@ def evict_from_tree_cache(tree_cache: BasePrefixCache | None, num_tokens: int): supports_swa_byte_budget, ) - is_unified_swa = supports_swa_byte_budget(allocator) - if is_unified_swa: - reclaim_plan = allocator.reclaim_plan( - num_tokens, - num_tokens, - full_evictable_tokens=tree_cache.full_evictable_size(), - swa_evictable_tokens=tree_cache.swa_evictable_size(), - ) - if reclaim_plan is None: - return - full_num_tokens, swa_num_tokens = reclaim_plan - else: - full_num_tokens = max(0, num_tokens - allocator.full_available_size()) - swa_num_tokens = max(0, num_tokens - allocator.swa_available_size()) + if supports_swa_byte_budget(allocator): + allocator.evict_to_free_tokens(tree_cache, num_tokens) + return + + full_num_tokens = max(0, num_tokens - allocator.full_available_size()) + swa_num_tokens = max(0, num_tokens - allocator.swa_available_size()) if full_num_tokens or swa_num_tokens: tree_cache.evict_for_alloc( EvictParams(num_tokens=full_num_tokens, swa_num_tokens=swa_num_tokens) ) - if is_unified_swa: - # Even a zero-reclaim plan can depend on compaction before callers - # inspect availability. - allocator.ensure_capacity(num_tokens, num_tokens) else: # Standard allocator: evict only the shortfall (mirrors the SWA arm) available_size = allocator.available_size() diff --git a/python/sglang/srt/mem_cache/kv_cache_configurator.py b/python/sglang/srt/mem_cache/kv_cache_configurator.py index f9fb8bbfbc2d..2ddb58efb21f 100644 --- a/python/sglang/srt/mem_cache/kv_cache_configurator.py +++ b/python/sglang/srt/mem_cache/kv_cache_configurator.py @@ -290,6 +290,23 @@ def __post_init__(self) -> None: self.draft_model_idx in self.model_config.swa_attention_layer_ids ) + def hybrid_swa_token_capacity( + self, + *, + allocator: BaseTokenToKVPoolAllocator, + full_capacity: Optional[int], + swa_capacity: Optional[int], + ) -> int: + if get_memory().enable_unified_memory: + capacity = allocator.size_full + max_total_tokens = get_schedule().max_total_tokens + return ( + min(capacity, max_total_tokens) + if max_total_tokens is not None + else capacity + ) + return full_capacity or swa_capacity + def _build_fp4_quant_method(self, *, num_layers: int): if not is_float4_e2m1fn_x2(self.kv_cache_dtype): return None @@ -2379,9 +2396,8 @@ def config_from_budget( f"{config.max_total_num_tokens}" ) if max_tokens != config.max_total_num_tokens: - # Token-capped re-derivation: the profiled budget no longer - # applies; the recalced config's unified_total_bytes stays None - # and the factories fall back to the token-count byte sum. + # Re-derive the capped budget: SWA carries unified_memory_pool_bytes; + # Mamba factories fall back to token-count sizing without unified_total_bytes. config = configurator.calculate_pool_sizes_from_max_tokens( max_tokens, get_schedule().page_size ) diff --git a/python/sglang/srt/mem_cache/unified_memory_pool.py b/python/sglang/srt/mem_cache/unified_memory_pool.py index 6c6a113cadf4..33f29a51bfc9 100644 --- a/python/sglang/srt/mem_cache/unified_memory_pool.py +++ b/python/sglang/srt/mem_cache/unified_memory_pool.py @@ -1711,7 +1711,6 @@ def init_unified_swa_pools( need_sort: bool, forward_stream: Optional[torch.cuda.Stream] = None, lazy_compaction: bool = False, - unified_total_bytes: Optional[int] = None, model_context_len: Optional[int] = None, sliding_window_size: Optional[int] = None, ) -> UnifiedSWAPoolBundle: @@ -1750,11 +1749,6 @@ def init_unified_swa_pools( store_dtype=store_dtype, grow_direction="up", ) - if total_bytes is not None and unified_total_bytes is not None: - raise ValueError("total_bytes and unified_total_bytes are mutually exclusive") - if total_bytes is None: - total_bytes = unified_total_bytes - legacy_allocator_capacities = {} if total_bytes is None: if full_max_total_num_tokens is None or swa_max_total_num_tokens is None: diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index aa4baf45025c..1f8798b29a02 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -1373,12 +1373,12 @@ def unload_lora_adapter(self, lora_ref: LoRARef): @property def effective_max_total_num_tokens(self): """Return the max token pool size considering hybrid swa settings.""" - if self.is_hybrid_swa and get_memory().enable_unified_memory: - capacity = self.token_to_kv_pool_allocator.size_full - if get_schedule().max_total_tokens is not None: - capacity = min(capacity, get_schedule().max_total_tokens) - elif self.is_hybrid_swa: - capacity = self.full_max_total_num_tokens or self.swa_max_total_num_tokens + if self.is_hybrid_swa: + capacity = self.kv_cache_configurator.hybrid_swa_token_capacity( + allocator=self.token_to_kv_pool_allocator, + full_capacity=self.full_max_total_num_tokens, + swa_capacity=self.swa_max_total_num_tokens, + ) else: capacity = self.max_total_num_tokens if (req_to_token_pool := getattr(self, "req_to_token_pool", None)) is not None: diff --git a/test/registered/unit/mem_cache/test_hisparse_max_token_pool_size.py b/test/registered/unit/mem_cache/test_hisparse_max_token_pool_size.py index 4c590ce02bf1..5a7222b3872a 100644 --- a/test/registered/unit/mem_cache/test_hisparse_max_token_pool_size.py +++ b/test/registered/unit/mem_cache/test_hisparse_max_token_pool_size.py @@ -14,6 +14,7 @@ from unittest.mock import MagicMock from sglang.srt.disaggregation.decode import DecodePreallocQueue +from sglang.srt.mem_cache.kv_cache_configurator import KVCacheConfigurator from sglang.srt.model_executor.model_runner import ModelRunner from sglang.srt.runtime_context import get_context from sglang.test.ci.ci_register import register_cpu_ci @@ -29,6 +30,7 @@ def _make_model_runner(**attrs): raise AttributeError on the internal `self.effective_max_total_num_tokens` read inside `max_token_pool_size`.""" instance = object.__new__(ModelRunner) + instance.kv_cache_configurator = object.__new__(KVCacheConfigurator) for name, value in attrs.items(): setattr(instance, name, value) return instance diff --git a/test/registered/unit/mem_cache/test_multi_ended_allocator.py b/test/registered/unit/mem_cache/test_multi_ended_allocator.py index 6f677c94f64e..aaed067bf844 100644 --- a/test/registered/unit/mem_cache/test_multi_ended_allocator.py +++ b/test/registered/unit/mem_cache/test_multi_ended_allocator.py @@ -727,7 +727,7 @@ def test_empty_pool_reservation_matches_packed_byte_boundary(self): max_new_tokens, sliding_window_size=16, page_size=page_size, - include_allocated_tail=False, + allocation_limit=16, ) reservation_swa_with_tail = estimate_swa_kv_tokens( extend_tokens, diff --git a/test/registered/unit/mem_cache/test_unified_byte_budget_sizing.py b/test/registered/unit/mem_cache/test_unified_byte_budget_sizing.py index f43ed5be8d77..d90ca5c45c17 100644 --- a/test/registered/unit/mem_cache/test_unified_byte_budget_sizing.py +++ b/test/registered/unit/mem_cache/test_unified_byte_budget_sizing.py @@ -13,9 +13,9 @@ # ============================================================================== """Byte-budget buffer sizing for the unified 2-pool factories. -With ``unified_total_bytes`` set the buffer is that many bytes exactly (the -mamba pair adds the state pool's bytes on top -- the budget is captured AFTER -the state carve-out); without it, sizing falls back to the token-count re-sum. +The SWA factory's ``total_bytes`` sets the exact buffer size. The Mamba pair's +``unified_total_bytes`` adds state bytes because its budget excludes that carve-out. +Without an explicit budget, sizing falls back to the token-count re-sum. Sizing from the ratio-derived token counts instead would re-introduce the configurator's rounding, which floors by the cell size and then page-aligns EACH side, losing up to about a page of tokens per side. @@ -91,7 +91,7 @@ def test_swa_factory_honors_the_budget_exactly(self): (64 + 32) * e + (e - 2), # almost one more entry ): with self.subTest(budget=budget): - bundle = _swa_factory(unified_total_bytes=budget) + bundle = _swa_factory(total_bytes=budget) self.assertEqual(bundle.unified_memory_pool.total_bytes, budget) def test_fallback_is_the_token_count_resum(self): @@ -166,7 +166,7 @@ def test_infeasible_budget_raises_before_construction(self): must fail loud instead of livelocking later.""" with self.assertRaises(RuntimeError) as ctx: _swa_factory( - unified_total_bytes=8 * _entry_bytes(), + total_bytes=8 * _entry_bytes(), model_context_len=4096, sliding_window_size=4096, ) @@ -184,7 +184,7 @@ def test_feasible_floor_inputs_are_not_rejected(self): ): with self.subTest(case=case): bundle = _swa_factory( - unified_total_bytes=200 * e, + total_bytes=200 * e, model_context_len=model_context_len, sliding_window_size=sliding_window_size, ) diff --git a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py index 71c8d7e47fdf..64ce90165534 100644 --- a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py +++ b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py @@ -4083,7 +4083,7 @@ def test_buffer_only_anchor_lock_cap_clamped_by_context_headroom(self): def test_buffer_load_back_swa_window_charged_at_admission(self): """Admission contract: a request the SWA budget gate accepts must be - allocatable at batch time (_swa_reserved_tokens: "an admitted request + allocatable at batch time (estimate_swa_kv_tokens: "an admitted request cannot OOM"). Regression: buffer mode surfaced a staged prefetch as host_hit_length only, so the gate never charged the SWA window that consumption (init_load_back -> cc.load) allocates and the request From 6c7be6252fd49a45c8e438372809a1d418d18e79 Mon Sep 17 00:00:00 2001 From: Cheng Wan Date: Sun, 13 Sep 2026 16:08:54 -0700 Subject: [PATCH 36/55] Refactor SWA admission into allocator-owned memory budgets --- python/sglang/srt/managers/schedule_policy.py | 393 ++--------------- python/sglang/srt/managers/scheduler.py | 61 +-- .../pool_stats_observer.py | 28 +- python/sglang/srt/mem_cache/allocator/base.py | 34 +- .../srt/mem_cache/allocator/hisparse.py | 14 + python/sglang/srt/mem_cache/allocator/swa.py | 35 ++ .../mem_cache/allocator/unified_hybrid_swa.py | 91 +++- python/sglang/srt/mem_cache/common.py | 36 +- python/sglang/srt/mem_cache/prefill_budget.py | 405 ++++++++++++++++++ .../unit/managers/test_prefill_adder.py | 176 ++++++-- .../test_scheduler_init_req_max_new_tokens.py | 12 +- .../mem_cache/test_multi_ended_allocator.py | 2 +- .../mem_cache/test_prefill_memory_budget.py | 224 ++++++++++ .../test_unified_radix_allocation_eviction.py | 6 +- 14 files changed, 1009 insertions(+), 508 deletions(-) create mode 100644 python/sglang/srt/mem_cache/prefill_budget.py create mode 100644 test/registered/unit/mem_cache/test_prefill_memory_budget.py diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index 24110fea767a..385f4870d689 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -48,13 +48,10 @@ DeepSeekV4HiSparseTokenToKVPoolAllocator, ) from sglang.srt.mem_cache.allocator.swa import ( - PureSWATokenToKVPoolAllocator, SWATokenToKVPoolAllocator, - is_swa_req_ring, ) from sglang.srt.mem_cache.allocator.unified_hybrid_swa import ( UnifiedMambaSWATokenToKVPoolAllocator, - supports_swa_byte_budget, ) from sglang.srt.mem_cache.allocator.unified_mamba import ( UnifiedMambaTokenToKVPoolAllocator, @@ -80,37 +77,6 @@ ) -def estimate_swa_kv_tokens( - extend_input_len: int, - max_new_tokens: int, - *, - sliding_window_size: Optional[int], - page_size: int, - allocation_limit: Optional[int] = None, - host_hit_length: int = 0, -) -> int: - """Peak SWA reservation for one prefill/decode request.""" - if sliding_window_size is None or sliding_window_size <= 0: - reserved = extend_input_len + max_new_tokens + page_size - else: - allocated = ( - extend_input_len - if allocation_limit is None - else min(extend_input_len, allocation_limit) - ) - allocated_tail = max(allocated - sliding_window_size, 0) - # A full window would double-charge short cached-prefix resumes; - # including extend keeps the reservation above the prefill allocation. - reserved = ( - allocated_tail - + min(extend_input_len + max_new_tokens, sliding_window_size) - + page_size - ) - if host_hit_length > 0: - reserved += -(-host_hit_length // page_size) * page_size - return reserved - - # Threshold for in-batch prefix cache. # If a request has a matched prefix length (against existing cache) less than this value, # the scheduler runs the in-batch prefix caching check for this request. @@ -591,8 +557,6 @@ def __init__( self.prefill_tile_block_m = prefill_tile_block_m self.tree_cache = tree_cache self.token_to_kv_pool_allocator = token_to_kv_pool_allocator - # Per-request SWA ring: one fixed slot per request, not a token budget. - self._swa_req_ring = is_swa_req_ring(token_to_kv_pool_allocator) self.running_batch = running_batch self.new_token_ratio = new_token_ratio self.rem_input_tokens = rem_input_tokens - num_mixed_decode_tokens @@ -604,8 +568,9 @@ def __init__( if self.rem_chunk_tokens is not None: self.rem_chunk_tokens -= num_mixed_decode_tokens - self.rem_total_token_offset = num_mixed_decode_tokens - self.cur_rem_token_offset = num_mixed_decode_tokens + self.memory_budget = token_to_kv_pool_allocator.create_prefill_budget( + tree_cache, num_mixed_decode_tokens=num_mixed_decode_tokens + ) self.req_states = None self.can_run_list = [] @@ -622,7 +587,7 @@ def __init__( if running_batch is not None: # Estimate the offset in the remaining token space - self.rem_total_token_offset += sum( + self.memory_budget.total_offset += sum( [ self._get_running_request_total_token_offset(r) for r in running_batch.reqs @@ -635,16 +600,7 @@ def __init__( self.token_to_kv_pool_allocator, (SWATokenToKVPoolAllocator, DeepSeekV4HiSparseTokenToKVPoolAllocator), ) - self.is_all_swa = isinstance( - self.token_to_kv_pool_allocator, PureSWATokenToKVPoolAllocator - ) self.is_hybrid_ssm_cache = self.tree_cache.supports_mamba() - self.is_unified_swa = supports_swa_byte_budget(self.token_to_kv_pool_allocator) - - self.rem_swa_token_offset = ( - num_mixed_decode_tokens if self.is_unified_swa else 0 - ) - # A new state slot eats shared-gap bytes that `rem_total_tokens` counts # as free, so reserve per slot or admission over-commits. Gate on the # ALLOCATOR, not `is_hybrid_ssm_cache`: that is False for `ChunkCache`, @@ -731,95 +687,11 @@ def _get_running_request_total_token_offset(self, req: Req) -> int: @property def rem_total_tokens(self): - if self.is_all_swa: - available_and_evictable = ( - self.token_to_kv_pool_allocator.swa_available_size() - + self.tree_cache.swa_evictable_size() - ) - elif self.is_hybrid_swa: - available_and_evictable = ( - self.token_to_kv_pool_allocator.full_available_size() - + self.tree_cache.full_evictable_size() - ) - elif self.is_hybrid_ssm_cache: - available_and_evictable = ( - self.token_to_kv_pool_allocator.available_size() - + self.tree_cache.full_evictable_size() - ) - else: - available_and_evictable = ( - self.token_to_kv_pool_allocator.available_size() - + self.tree_cache.evictable_size() - ) - return available_and_evictable - self.rem_total_token_offset - - @property - def rem_swa_tokens(self): - allocator = self.token_to_kv_pool_allocator - if self._swa_req_ring: - # swa_available_size() already reports ring capacity; tree - # swa_evictable is in linear token units and frees no ring space. - return allocator.swa_available_size() - self.rem_swa_token_offset - return ( - allocator.swa_available_size() - + self.tree_cache.swa_evictable_size() - - self.rem_swa_token_offset - ) + return self.memory_budget.remaining_total @property def cur_rem_tokens(self): - if self.is_all_swa: - available_and_evictable = ( - self.token_to_kv_pool_allocator.swa_available_size() - + self.tree_cache.swa_evictable_size() - ) - elif self.is_hybrid_swa: - available_and_evictable = ( - self.token_to_kv_pool_allocator.full_available_size() - + self.tree_cache.full_evictable_size() - ) - elif self.is_hybrid_ssm_cache: - available_and_evictable = ( - self.token_to_kv_pool_allocator.available_size() - + self.tree_cache.full_evictable_size() - ) - else: - available_and_evictable = ( - self.token_to_kv_pool_allocator.available_size() - + self.tree_cache.evictable_size() - ) - - return available_and_evictable - self.cur_rem_token_offset - - def _swa_budget_for_req( - self, extend_input_len: int, max_new_tokens: int, swa_host_hit_length: int = 0 - ) -> int: - """SWA pool budget per request. Only valid when is_hybrid_swa is True. - - With chunked prefill + overlap scheduler, the peak SWA occupancy is: - chunk N (running, not yet in tree) + sliding window (locked in tree) - + chunk N+1 (new allocation) - Since chunk N and locked tokens are already excluded from - swa_available + swa_evictable, the budget only needs to cover the - chunk N+1 allocation plus decode headroom: - - budget = max(alloc - window, 0) + min(extend + max_new_tokens, window) + page - - where alloc = min(extend, rem_chunk); the min() cap keeps the two terms - from double-counting extend, so budget <= extend + max_new_tokens + page. - """ - allocator = self.token_to_kv_pool_allocator - if self._swa_req_ring: - # One ring slot per request, in the same unit as swa_available_size. - return allocator.swa_ring_cost_tokens - return estimate_swa_kv_tokens( - extend_input_len, - max_new_tokens, - sliding_window_size=self.tree_cache.sliding_window_size, - page_size=self.page_size, - allocation_limit=self.rem_chunk_tokens, - host_hit_length=swa_host_hit_length, - ) + return self.memory_budget.remaining_current def _swa_new_tokens(self, req: Req) -> int: """Tokens a request may still decode, for SWA headroom sizing. Mirrors @@ -831,150 +703,17 @@ def _swa_new_tokens(self, req: Req) -> int: CLIP_MAX_NEW_TOKENS, ) - def _unified_swa_reservation_fits( - self, - full_tokens: int, - swa_tokens: int, - *, - empty_pool: bool = False, - ) -> bool: - allocator = self.token_to_kv_pool_allocator - if empty_pool: - full_evictable_tokens = 0 - swa_evictable_tokens = 0 - else: - full_tokens += self.rem_total_token_offset - swa_tokens += self.rem_swa_token_offset - full_evictable_tokens = self.tree_cache.full_evictable_size() - swa_evictable_tokens = self.tree_cache.swa_evictable_size() - return allocator.can_reserve( - full_tokens, - swa_tokens, - full_evictable_tokens=full_evictable_tokens, - swa_evictable_tokens=swa_evictable_tokens, - empty_pool=empty_pool, - require_token_slack=empty_pool, - ) - - def _unified_swa_chunk_cap( - self, - max_new_tokens: int, - swa_host_hit_length: int = 0, - *, - max_chunk_tokens: int, - ) -> int: - """Largest page-aligned chunk that fits the joint FULL/SWA byte budget.""" - cap = max_chunk_tokens // self.page_size * self.page_size - lo, hi = 0, cap // self.page_size - while lo < hi: - mid = (lo + hi + 1) // 2 - chunk_tokens = mid * self.page_size - full_tokens = self.ceil_paged_tokens(chunk_tokens) + self.page_size - swa_tokens = self._swa_budget_for_req( - chunk_tokens, max_new_tokens, swa_host_hit_length - ) - if self._unified_swa_reservation_fits(full_tokens, swa_tokens): - lo = mid - else: - hi = mid - 1 - return lo * self.page_size - - def _swa_chunk_cap(self, max_new_tokens: int, swa_host_hit_length: int = 0) -> int: - """Largest chunk fitting the fixed SWA partition after decode headroom.""" - # Only the sliding window stays locked across chunks, so the next - # chunk can reuse the transient space outside its decode headroom. - cap = int(self.rem_swa_tokens) - estimate_swa_kv_tokens( - 0, - max_new_tokens, - sliding_window_size=self.tree_cache.sliding_window_size, - page_size=self.page_size, - host_hit_length=swa_host_hit_length, - ) - return max(0, cap // self.page_size * self.page_size) - - def _swa_req_never_fits( - self, - extend_input_len: int, - max_new_tokens: int, - swa_host_hit_length: int = 0, - ) -> bool: - """True when a request's SWA budget exceeds the *entire* SWA pool, so it - can never be admitted whole no matter how far the pool drains. - - This is the head-of-line livelock the _swa_chunk_cap escape hatch exists - for; the hatch must fire only in this case. A request that merely - exceeds *current* rem_swa (transient pressure) would fit once running - decodes free their windows, so it must wait — admitting it into the - decode headroom collapses the SWA evictable cushion and forces running - requests to retract (observed as a severe retraction/re-prefill storm on - hybrid-SWA models at high concurrency).""" - swa_tokens = self._swa_budget_for_req( - extend_input_len, max_new_tokens, swa_host_hit_length - ) - if self._swa_req_ring: - return swa_tokens > self.token_to_kv_pool_allocator.size_swa - return swa_tokens >= self.token_to_kv_pool_allocator.size_swa - - def _check_swa_admission( - self, - req: Req, - *, - extend_input_len: int, - total_tokens: int, - ) -> tuple[bool, Optional[int]]: - if not self.is_unified_swa: - extend_input_len = self.ceil_paged_tokens(extend_input_len) - max_new_tokens = self._swa_new_tokens(req) - swa_needed = self._swa_budget_for_req( - extend_input_len, - max_new_tokens, - swa_host_hit_length=req.swa_host_hit_length, - ) - swa_cap = 0 - # Only permanent shortfalls may shrink a chunk; transient pressure waits. - if self.is_unified_swa: - fits = self._unified_swa_reservation_fits(total_tokens, swa_needed) - if not fits: - full_input_tokens = len(req.full_untruncated_fill_ids) - if not self._unified_swa_reservation_fits( - full_input_tokens + max_new_tokens + self.page_size, - self._swa_budget_for_req(full_input_tokens, max_new_tokens), - empty_pool=True, - ): - swa_cap = self._unified_swa_chunk_cap( - 0, - req.swa_host_hit_length, - max_chunk_tokens=min( - self.rem_chunk_tokens or 0, max(0, extend_input_len - 1) - ), - ) - else: - # Ring-slot capacity is exact; the fixed token partition keeps slack. - fits = ( - swa_needed <= self.rem_swa_tokens - if self._swa_req_ring - else swa_needed < self.rem_swa_tokens - ) - if not fits and self._swa_req_never_fits( - extend_input_len, max_new_tokens, req.swa_host_hit_length - ): - swa_cap = self._swa_chunk_cap(max_new_tokens, req.swa_host_hit_length) - if fits: - return True, self.rem_chunk_tokens - if self.rem_chunk_tokens is None or swa_cap <= 0: - return False, None - return True, min(self.rem_chunk_tokens, swa_cap) - def _check_prefill_budget( self, req: Req, *, extend_input_len: int, total_tokens: int ) -> tuple[bool, Optional[int]]: - if not self.is_unified_swa and total_tokens >= self.rem_total_tokens: - return False, None - if self.is_hybrid_swa: - return self._check_swa_admission( - req, extend_input_len=extend_input_len, total_tokens=total_tokens - ) - return True, self.rem_chunk_tokens + return self.memory_budget.check_prefill( + extend_input_len=extend_input_len, + total_tokens=total_tokens, + max_new_tokens=self._swa_new_tokens(req), + input_tokens=len(req.full_untruncated_fill_ids), + swa_host_hit_length=req.swa_host_hit_length, + chunk_limit=self.rem_chunk_tokens, + ) def _mamba_gap_budget_for_req(self, req: Req) -> int: """Shared-gap reservation (full-token-equivalents) for a request's new @@ -995,12 +734,7 @@ def ceil_paged_tokens(self, tokens: int) -> int: return -(-tokens // self.page_size) * self.page_size def budget_state(self): - if self.is_unified_swa: - no_token = not self._unified_swa_reservation_fits(0, 0) - else: - no_token = self.rem_total_tokens <= 0 or self.cur_rem_tokens <= 0 - if not no_token and self.is_hybrid_swa: - no_token = self.rem_swa_tokens <= 0 + no_token = not self.memory_budget.has_capacity() # Gate new mamba slots separately: rem_total_tokens' full_evictable can't # cover a mamba slot, which needs mamba-recoverable bytes (see __init__). if not no_token and self.rem_mamba_slots is not None: @@ -1032,17 +766,12 @@ def _update_prefill_budget( # TODO(lsyin): check this workaround logic, which only ensures the prefill will not out of memory, and may be too conservative extend_input_len = self.ceil_paged_tokens(extend_input_len) - # alloc_extend reserves an extra page_size per request to make sure the budget doesn't over-commit - page_overhead = self.page_size - # `mamba_gap_reserve` (shared Mamba pool only; 0 otherwise) charges the new - # mamba state's shared-gap cost to BOTH full budgets: the slot is allocated - # immediately (counts against `cur_rem`) and held for the request lifetime - # (counts against `rem_total`). See `_mamba_gap_budget_for_req`. - self.rem_total_token_offset += ( - extend_input_len + max_new_tokens + page_overhead + mamba_gap_reserve - ) - self.cur_rem_token_offset += ( - extend_input_len + page_overhead + mamba_gap_reserve + self.memory_budget.reserve( + extend_input_len, + max_new_tokens, + extra_tokens=mamba_gap_reserve, + chunk_limit=self.rem_chunk_tokens, + is_chunked_continuation=is_chunked_continuation, ) # The new mamba slot also consumes one mamba-recoverable slot (gated # separately so full_evictable can't cover it — see __init__). @@ -1050,14 +779,6 @@ def _update_prefill_budget( self.rem_mamba_slots -= 1 self.rem_input_tokens -= extend_input_len - if self.is_hybrid_swa: - # The ring slot is reserved once at first admission; charging it - # again on a continuation would double-count and over-throttle. - if not (self._swa_req_ring and is_chunked_continuation): - self.rem_swa_token_offset += self._swa_budget_for_req( - extend_input_len, max_new_tokens - ) - if self.dllm_config is not None: self.rem_dllm_tokens -= extend_input_len elif self.rem_chunk_tokens is not None: @@ -1182,44 +903,15 @@ def add_dllm_staging_req(self, req: Req): else AddReqResult.CONTINUE ) - def _unified_swa_chunk_limit( - self, req: Req, extend_input_len: int, max_chunk_tokens: int - ) -> Optional[int]: - candidate_tokens = min(extend_input_len, max_chunk_tokens) - candidate_finishes = candidate_tokens >= extend_input_len - max_new_tokens = self._swa_new_tokens(req) if candidate_finishes else 0 - if self._unified_swa_reservation_fits( - self.ceil_paged_tokens(candidate_tokens) + max_new_tokens + self.page_size, - self._swa_budget_for_req(candidate_tokens, max_new_tokens), - ): - return max_chunk_tokens - max_partial = ( - max(0, candidate_tokens - 1) if candidate_finishes else candidate_tokens - ) - joint_cap = self._unified_swa_chunk_cap(0, max_chunk_tokens=max_partial) - if joint_cap <= 0: - return None - return min(max_chunk_tokens, joint_cap) - def add_chunked_req(self, req: Req): if self.dllm_config is not None: _rem_tokens = self._get_dllm_remain_tokens() else: - if self.is_unified_swa: - _rem_tokens = self.rem_chunk_tokens - else: - _rem_tokens = min(self.rem_chunk_tokens, int(self.rem_total_tokens)) - if self.is_hybrid_swa and not self._swa_req_ring: - # alloc_extend needs one extra page; ring mode counts slots. - _rem_tokens = min( - _rem_tokens, int(self.rem_swa_tokens) - self.page_size - ) - # The chunked_req must be added to the list; otherwise, it will cause a memory leak. - # Therefore, in certain cases where _rem_tokens <= 0, it should be replaced with rem_chunk_tokens. - if _rem_tokens <= 0: - if self.is_hybrid_swa: - return req - _rem_tokens = self.rem_chunk_tokens + _rem_tokens = self.memory_budget.available_chunk_tokens( + self.rem_chunk_tokens + ) + if _rem_tokens is None: + return req # A mid-chunk rank prefills this pass regardless of the delayer # verdict, so report prefillable=True and ignore the result. @@ -1235,12 +927,13 @@ def add_chunked_req(self, req: Req): cand_extend_input_len = len(req.full_untruncated_fill_ids) - len( req.prefix_indices ) - if self.is_unified_swa: - _rem_tokens = self._unified_swa_chunk_limit( - req, cand_extend_input_len, _rem_tokens - ) - if _rem_tokens is None: - return req + _rem_tokens = self.memory_budget.fit_chunk( + extend_input_len=cand_extend_input_len, + max_new_tokens=self._swa_new_tokens(req), + chunk_limit=_rem_tokens, + ) + if _rem_tokens is None: + return req truncated = cand_extend_input_len > _rem_tokens new_len = min(cand_extend_input_len, _rem_tokens) req.set_extend_range(len(req.prefix_indices), len(req.prefix_indices) + new_len) @@ -1285,18 +978,12 @@ def add_one_req_ignore_eos(self, req: Req): # Shared Mamba pool: fold the new mamba state's shared-gap cost into the # budget gate so admission can't over-commit (0 for baseline / non-Mamba). paged_input += self._mamba_gap_budget_for_req(req) - if self.is_hybrid_swa: - max_new_tokens = self._swa_new_tokens(req) - swa_needed = self._swa_budget_for_req(cand_extend_input_len, max_new_tokens) - if self.is_unified_swa: - fits = self._unified_swa_reservation_fits( - paged_input + max_new_tokens + self.page_size, - swa_needed, - ) - else: - fits = paged_input <= min(self.cur_rem_tokens, self.rem_total_tokens) and ( - not self.is_hybrid_swa or swa_needed <= self.rem_swa_tokens - ) + fits = self.memory_budget.can_allocate_prefill( + paged_input=paged_input, + extend_input_len=cand_extend_input_len, + max_new_tokens=self._swa_new_tokens(req), + chunk_limit=self.rem_chunk_tokens, + ) if not fits: return AddReqResult.NO_TOKEN @@ -1666,7 +1353,7 @@ def preempt_to_schedule(self, req: Req) -> bool: release_counter = 0 for i, running_req in enumerate(self.running_batch.reqs): if running_req in preemptible_reqs: - self.rem_total_token_offset -= ( + self.memory_budget.total_offset -= ( self._get_running_request_total_token_offset(running_req) ) release_counter += 1 diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index e50ddf012669..d6a8ca000d58 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -212,7 +212,6 @@ AddReqResult, PrefillAdder, SchedulePolicy, - estimate_swa_kv_tokens, ) from sglang.srt.managers.scheduler_components.batch_result_processor import ( SchedulerBatchResultProcessor, @@ -280,9 +279,6 @@ validate_input_length, ) from sglang.srt.mem_cache import kv_cache_builder -from sglang.srt.mem_cache.allocator.unified_hybrid_swa import ( - supports_swa_byte_budget, -) from sglang.srt.mem_cache.common import ( maybe_cache_unfinished_req, release_kv_cache, @@ -2520,18 +2516,18 @@ def init_req_max_new_tokens(self, req): self.max_req_len - input_len - 1, ), ) - if supports_swa_byte_budget(self.token_to_kv_pool_allocator): - max_new_tokens = self._clip_unified_swa_max_new_tokens(req, max_new_tokens) - else: - # Non-unified pools keep their fixed token partition. - paged_input_len = -(-input_len // self.page_size) * self.page_size - max_new_tokens = min( - max_new_tokens, - self.max_total_num_tokens * get_parallel().attn_dcp_size - - paged_input_len - - self.page_size - - 1, + max_new_tokens = self.token_to_kv_pool_allocator.max_new_tokens_for_memory( + input_len, + max_new_tokens, + token_capacity=self.max_total_num_tokens * get_parallel().attn_dcp_size, + sliding_window_size=self.sliding_window_size, + chunk_size=self.chunked_prefill_size, + ) + if max_new_tokens is None: + req.set_finish_with_abort( + f"Request prompt exceeds the KV memory budget: input_len={input_len}." ) + max_new_tokens = 0 req.sampling_params.max_new_tokens = max(0, max_new_tokens) # Clipping above can push max_new_tokens below min_new_tokens, which @@ -2539,41 +2535,6 @@ def init_req_max_new_tokens(self, req): if req.sampling_params.min_new_tokens > req.sampling_params.max_new_tokens: req.sampling_params.min_new_tokens = req.sampling_params.max_new_tokens - def _clip_unified_swa_max_new_tokens(self, req: Req, max_new_tokens: int) -> int: - input_len = len(req.origin_input_ids) - - def fits_shared_pool(candidate: int) -> bool: - swa_tokens = estimate_swa_kv_tokens( - input_len, - candidate, - sliding_window_size=self.sliding_window_size, - page_size=self.page_size, - allocation_limit=self.chunked_prefill_size, - ) - return self.token_to_kv_pool_allocator.can_reserve( - input_len + candidate + self.page_size, - swa_tokens, - empty_pool=True, - require_token_slack=True, - ) - - if not fits_shared_pool(0): - req.set_finish_with_abort( - "Request prompt exceeds the unified FULL/SWA KV byte budget: " - f"input_len={input_len}." - ) - return 0 - if fits_shared_pool(max_new_tokens): - return max_new_tokens - lo, hi = 0, max_new_tokens - while lo < hi: - mid = (lo + hi + 1) // 2 - if fits_shared_pool(mid): - lo = mid - else: - hi = mid - 1 - return lo - def _process_and_broadcast_mm_inputs( self, raw_mm_inputs, diff --git a/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py b/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py index 823fe15b67cf..f061506f1182 100644 --- a/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py +++ b/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py @@ -12,10 +12,6 @@ ) from sglang.srt.mem_cache.allocator.swa import is_swa_req_ring -from sglang.srt.mem_cache.allocator.unified_hybrid_swa import ( - UnifiedMambaSWATokenToKVPoolAllocator, - supports_swa_byte_budget, -) if TYPE_CHECKING: from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator @@ -291,28 +287,12 @@ def _get_mamba_token_info(self): mamba_evictable_size=mamba_evictable_size, ) - def _swa_capacity_and_available(self) -> tuple[tuple[int, int], tuple[int, int]]: - allocator = self.token_to_kv_pool_allocator - if supports_swa_byte_budget(allocator): - return ( - (allocator.current_full_capacity, allocator.full_available_size()), - (allocator.current_swa_capacity, allocator.swa_available_size()), - ) - if isinstance(allocator, UnifiedMambaSWATokenToKVPoolAllocator): - # The tri-pool reports static capacities paired with conserve views. - full_available_size = allocator.conserve_full_available_size() - swa_available_size = allocator.conserve_swa_available_size() - else: - full_available_size = allocator.full_available_size() - swa_available_size = allocator.swa_available_size() - return ( - (self.full_tokens_per_layer, full_available_size), - (self.swa_tokens_per_layer, swa_available_size), - ) - def _get_swa_token_info(self) -> PoolStats: (full_capacity, full_available_size), (swa_capacity, swa_available_size) = ( - self._swa_capacity_and_available() + self.token_to_kv_pool_allocator.swa_capacity_and_available( + full_capacity=self.full_tokens_per_layer, + swa_capacity=self.swa_tokens_per_layer, + ) ) full_evictable_size = self.tree_cache.full_evictable_size() swa_evictable_size = self.tree_cache.swa_evictable_size() diff --git a/python/sglang/srt/mem_cache/allocator/base.py b/python/sglang/srt/mem_cache/allocator/base.py index fa640b650672..5bc0b6567f9f 100644 --- a/python/sglang/srt/mem_cache/allocator/base.py +++ b/python/sglang/srt/mem_cache/allocator/base.py @@ -69,12 +69,42 @@ def size_full(self): # The scheduler calls these unconditionally, with no allocator-type branches # on its side; byte-accounted composites override the token-count defaults. + def create_prefill_budget(self, tree_cache, *, num_mixed_decode_tokens=0): + from sglang.srt.mem_cache.prefill_budget import PrefillBudget + + return PrefillBudget( + self, tree_cache, num_mixed_decode_tokens=num_mixed_decode_tokens + ) + + def max_new_tokens_for_memory( + self, + input_tokens: int, + max_new_tokens: int, + *, + token_capacity: int, + sliding_window_size: int | None, + chunk_size: int | None, + ) -> int | None: + """Clip generation to the empty-pool budget; None means prompt cannot fit. + + token_capacity is the scheduler's configured capacity, including its + distributed token scaling. Shared pools use their physical byte layout. + """ + paged_input = -(-input_tokens // self.page_size) * self.page_size + return max( + 0, min(max_new_tokens, token_capacity - paged_input - self.page_size - 1) + ) + def evict_to_free_tokens(self, tree_cache, num_tokens: int) -> None: """Evict unlocked prefix-cache entries until this allocator can serve ``num_tokens`` or nothing evictable remains.""" - from sglang.srt.mem_cache.common import evict_from_tree_cache + from sglang.srt.mem_cache.base_prefix_cache import EvictParams - evict_from_tree_cache(tree_cache, num_tokens) + if tree_cache is None or tree_cache.is_chunk_cache(): + return + shortfall = num_tokens - self.available_size() + if shortfall > 0: + tree_cache.evict_for_alloc(EvictParams(num_tokens=shortfall)) def check_decode_capacity(self, *, num_tokens: int, tree_cache) -> bool: """Whether the next decode step's ``num_tokens`` allocation fits after diff --git a/python/sglang/srt/mem_cache/allocator/hisparse.py b/python/sglang/srt/mem_cache/allocator/hisparse.py index bef9b2a43e68..28f4e6cd511b 100644 --- a/python/sglang/srt/mem_cache/allocator/hisparse.py +++ b/python/sglang/srt/mem_cache/allocator/hisparse.py @@ -346,6 +346,20 @@ def get_kvcache(self): def translate_loc_from_full_to_swa(self, kv_indices: torch.Tensor): return self.logical_attn_allocator.translate_loc_from_full_to_swa(kv_indices) + def create_prefill_budget(self, tree_cache, *, num_mixed_decode_tokens=0): + # Use this wrapper's capacity, which also includes the compressed pool. + from sglang.srt.mem_cache.prefill_budget import SWAPrefillBudget + + return SWAPrefillBudget( + self, tree_cache, num_mixed_decode_tokens=num_mixed_decode_tokens + ) + + def swa_capacity_and_available(self, *, full_capacity, swa_capacity): + return ( + (full_capacity, self.full_available_size()), + (swa_capacity, self.swa_available_size()), + ) + def full_available_size(self): return min( self.logical_attn_allocator.full_available_size(), diff --git a/python/sglang/srt/mem_cache/allocator/swa.py b/python/sglang/srt/mem_cache/allocator/swa.py index 676f2450a004..77b1d8ddc3e1 100644 --- a/python/sglang/srt/mem_cache/allocator/swa.py +++ b/python/sglang/srt/mem_cache/allocator/swa.py @@ -158,6 +158,31 @@ def available_size(self): self.swa_attn_allocator.available_size(), ) + def create_prefill_budget(self, tree_cache, *, num_mixed_decode_tokens=0): + from sglang.srt.mem_cache.prefill_budget import SWAPrefillBudget + + return SWAPrefillBudget( + self, tree_cache, num_mixed_decode_tokens=num_mixed_decode_tokens + ) + + def swa_capacity_and_available(self, *, full_capacity, swa_capacity): + return ( + (full_capacity, self.full_available_size()), + (swa_capacity, self.swa_available_size()), + ) + + def evict_to_free_tokens(self, tree_cache, num_tokens: int) -> None: + from sglang.srt.mem_cache.base_prefix_cache import EvictParams + + if tree_cache is None or tree_cache.is_chunk_cache(): + return + full_shortfall = max(0, num_tokens - self.full_available_size()) + swa_shortfall = max(0, num_tokens - self.swa_available_size()) + if full_shortfall or swa_shortfall: + tree_cache.evict_for_alloc( + EvictParams(num_tokens=full_shortfall, swa_num_tokens=swa_shortfall) + ) + def full_available_size(self): return self.full_attn_allocator.available_size() @@ -672,6 +697,16 @@ def __init__( def available_size(self): return self.swa_attn_allocator.available_size() + def create_prefill_budget(self, tree_cache, *, num_mixed_decode_tokens=0): + from sglang.srt.mem_cache.prefill_budget import SWAPrefillBudget + + return SWAPrefillBudget( + self, + tree_cache, + num_mixed_decode_tokens=num_mixed_decode_tokens, + all_swa=True, + ) + def full_available_size(self): return self.swa_attn_allocator.available_size() diff --git a/python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py b/python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py index 3a7136d9da32..3d72723f6648 100644 --- a/python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py +++ b/python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py @@ -18,7 +18,7 @@ import logging import math -from typing import Callable, List, Optional, Sequence, Tuple, TypeGuard +from typing import Callable, List, Optional, Sequence, Tuple import torch from torch.profiler import record_function @@ -42,16 +42,6 @@ logger = logging.getLogger(__name__) -def supports_swa_byte_budget( - allocator: BaseTokenToKVPoolAllocator | None, -) -> TypeGuard[UnifiedSWATokenToKVPoolAllocator]: - """Whether FULL/SWA demand can be checked against a two-pool byte budget.""" - return ( - isinstance(allocator, UnifiedSWATokenToKVPoolAllocator) - and allocator.supports_asymmetric_reservation - ) - - class UnifiedSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator): """Composite allocator for the hybrid SWA pair (full + swa MHA sub-pools). @@ -59,10 +49,6 @@ class UnifiedSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator): `available_size()` (joint bytes, in TOKENS) is the only safe alloc pre-check. """ - # Unequal FULL/SWA reservations share one byte budget. The tri-pool opts - # out because its Mamba state competes for those same bytes. - supports_asymmetric_reservation = True - # Parent's `size` property has no setter but base init does `self.size = size`; # the no-op setter below absorbs that write. @property @@ -531,6 +517,57 @@ def ensure_capacity(self, full_tokens: int, swa_tokens: int) -> bool: compacted=False, ) + def create_prefill_budget(self, tree_cache, *, num_mixed_decode_tokens=0): + from sglang.srt.mem_cache.prefill_budget import SharedSWAPrefillBudget + + return SharedSWAPrefillBudget( + self, tree_cache, num_mixed_decode_tokens=num_mixed_decode_tokens + ) + + def swa_capacity_and_available(self, *, full_capacity, swa_capacity): + return ( + (self.current_full_capacity, self.full_available_size()), + (self.current_swa_capacity, self.swa_available_size()), + ) + + def max_new_tokens_for_memory( + self, + input_tokens: int, + max_new_tokens: int, + *, + token_capacity: int, + sliding_window_size: int | None, + chunk_size: int | None, + ) -> int | None: + from sglang.srt.mem_cache.prefill_budget import estimate_swa_kv_tokens + + def fits(candidate): + return self.can_reserve( + input_tokens + candidate + self.page_size, + estimate_swa_kv_tokens( + input_tokens, + candidate, + sliding_window_size=sliding_window_size, + page_size=self.page_size, + allocation_limit=chunk_size, + ), + empty_pool=True, + require_token_slack=True, + ) + + if not fits(0): + return None + if fits(max_new_tokens): + return max_new_tokens + lo, hi = 0, max_new_tokens + while lo < hi: + mid = (lo + hi + 1) // 2 + if fits(mid): + lo = mid + else: + hi = mid - 1 + return lo + def evict_to_free_tokens(self, tree_cache, num_tokens: int) -> None: from sglang.srt.mem_cache.base_prefix_cache import EvictParams @@ -961,8 +998,6 @@ class UnifiedMambaSWATokenToKVPoolAllocator(UnifiedSWATokenToKVPoolAllocator): `UnifiedMambaSlotAllocator`. """ - supports_asymmetric_reservation = False - def __init__( self, *, @@ -1216,17 +1251,33 @@ def set_inflight_forward( super().set_inflight_forward(forward_done, out_cache_loc_virtual) self.mamba_allocator.set_inflight_forward(forward_done, None) + def create_prefill_budget(self, tree_cache, *, num_mixed_decode_tokens=0): + # Mamba competes for the shared gap too; retain the tri-pool's existing + # token and state-slot admission until it has a three-way reservation. + return SWATokenToKVPoolAllocator.create_prefill_budget( + self, tree_cache, num_mixed_decode_tokens=num_mixed_decode_tokens + ) + + def max_new_tokens_for_memory(self, *args, **kwargs): + return BaseTokenToKVPoolAllocator.max_new_tokens_for_memory( + self, *args, **kwargs + ) + + def swa_capacity_and_available(self, *, full_capacity, swa_capacity): + return ( + (full_capacity, self.conserve_full_available_size()), + (swa_capacity, self.conserve_swa_available_size()), + ) + def evict_to_free_tokens(self, tree_cache, num_tokens: int) -> None: """Joint-aware eviction: one tri-lifetime node frees bytes on several sides at once, so re-check the JOINT gate instead of the per-side shortfall.""" - from sglang.srt.mem_cache.common import evict_from_tree_cache - # Arbitrary retry bound; a round that frees nothing ends the loop anyway. for _ in range(4): before = self.available_size() if before >= num_tokens: return - evict_from_tree_cache(tree_cache, num_tokens) + SWATokenToKVPoolAllocator.evict_to_free_tokens(self, tree_cache, num_tokens) if self.available_size() <= before: return # no progress diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index 16f1e557c90e..b3e73e6d9722 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/python/sglang/srt/mem_cache/common.py @@ -10,8 +10,7 @@ _get_last_loc_safe_kernel as _get_last_loc_safe_kernel, ) from sglang.kernels.ops.memory.common import get_last_loc_kernel as get_last_loc_kernel -from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator -from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache, EvictParams +from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache from sglang.srt.mem_cache.hicache_storage import PoolTransfer from sglang.srt.mem_cache.memory_pool import HybridReqToTokenPool, ReqToTokenPool from sglang.srt.runtime_context import get_serving, get_spec @@ -150,38 +149,11 @@ def maybe_cache_unfinished_req(req: Req, tree_cache: BasePrefixCache, **kwargs): def evict_from_tree_cache(tree_cache: BasePrefixCache | None, num_tokens: int): - if tree_cache is None: - return - - if tree_cache.is_chunk_cache(): - return - - allocator = tree_cache.token_to_kv_pool_allocator - - if isinstance(allocator, SWATokenToKVPoolAllocator): - # Hybrid allocator - from sglang.srt.mem_cache.allocator.unified_hybrid_swa import ( - supports_swa_byte_budget, + if tree_cache is not None and not tree_cache.is_chunk_cache(): + tree_cache.token_to_kv_pool_allocator.evict_to_free_tokens( + tree_cache, num_tokens ) - if supports_swa_byte_budget(allocator): - allocator.evict_to_free_tokens(tree_cache, num_tokens) - return - - full_num_tokens = max(0, num_tokens - allocator.full_available_size()) - swa_num_tokens = max(0, num_tokens - allocator.swa_available_size()) - if full_num_tokens or swa_num_tokens: - 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_for_alloc( - EvictParams(num_tokens=num_tokens - available_size) - ) - def retraction_backup( req: Req, diff --git a/python/sglang/srt/mem_cache/prefill_budget.py b/python/sglang/srt/mem_cache/prefill_budget.py new file mode 100644 index 000000000000..0811eb763354 --- /dev/null +++ b/python/sglang/srt/mem_cache/prefill_budget.py @@ -0,0 +1,405 @@ +# Copyright 2026 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Memory reservations for one prefill pass. + +The scheduler supplies token demand and its chunk/decode limits. These objects +account for admitted but not yet allocated work and query live cache capacity: +locking a prefix or preempting a request must affect the next admission check. +They neither select requests nor mutate the prefix cache or allocator. +""" + +from typing import Optional + + +def estimate_swa_kv_tokens( + extend_input_len: int, + max_new_tokens: int, + *, + sliding_window_size: Optional[int], + page_size: int, + allocation_limit: Optional[int] = None, + host_hit_length: int = 0, +) -> int: + """Peak SWA reservation for one prefill/decode request.""" + if sliding_window_size is None or sliding_window_size <= 0: + reserved = extend_input_len + max_new_tokens + page_size + else: + allocated = ( + extend_input_len + if allocation_limit is None + else min(extend_input_len, allocation_limit) + ) + allocated_tail = max(allocated - sliding_window_size, 0) + # A full window would double-charge short cached-prefix resumes; + # including extend keeps the reservation above the prefill allocation. + reserved = ( + allocated_tail + + min(extend_input_len + max_new_tokens, sliding_window_size) + + page_size + ) + if host_hit_length > 0: + reserved += -(-host_hit_length // page_size) * page_size + return reserved + + +class PrefillBudget: + """Fixed token pool. Offsets include pending allocations and decode headroom.""" + + def __init__(self, allocator, tree_cache, *, num_mixed_decode_tokens: int = 0): + self.allocator = allocator + self.tree_cache = tree_cache + self.page_size = allocator.page_size + self.total_offset = num_mixed_decode_tokens + self.current_offset = num_mixed_decode_tokens + self.swa_offset = 0 + + def ceil_paged_tokens(self, tokens: int) -> int: + return -(-tokens // self.page_size) * self.page_size + + def _available_and_evictable(self): + evictable = ( + self.tree_cache.full_evictable_size() + if self.tree_cache.supports_mamba() + else self.tree_cache.evictable_size() + ) + return self.allocator.available_size() + evictable + + @property + def remaining_total(self): + return self._available_and_evictable() - self.total_offset + + @property + def remaining_current(self): + return self._available_and_evictable() - self.current_offset + + @property + def remaining_swa(self): + return 0 + + def has_capacity(self) -> bool: + return self.remaining_total > 0 and self.remaining_current > 0 + + def check_prefill( + self, + *, + extend_input_len: int, + total_tokens: int, + max_new_tokens: int, + input_tokens: int, + swa_host_hit_length: int, + chunk_limit: int | None, + ) -> tuple[bool, int | None]: + """Return admission feasibility and a memory bound on the requested chunk.""" + return ( + (True, chunk_limit) + if total_tokens < self.remaining_total + else (False, None) + ) + + def can_allocate_prefill( + self, + *, + paged_input: int, + extend_input_len: int, + max_new_tokens: int, + chunk_limit: int | None, + ) -> bool: + return paged_input <= min(self.remaining_current, self.remaining_total) + + def available_chunk_tokens(self, chunk_limit: int) -> int | None: + available = min(chunk_limit, int(self.remaining_total)) + # Single-pool continuation must make progress to release its KV. + return available if available > 0 else chunk_limit + + def fit_chunk( + self, + *, + extend_input_len: int, + max_new_tokens: int, + chunk_limit: int, + ) -> int | None: + return chunk_limit + + def reserve( + self, + extend_input_len: int, + max_new_tokens: int, + *, + extra_tokens: int = 0, + chunk_limit: int | None = None, + is_chunked_continuation: bool = False, + ) -> None: + extend_input_len = self.ceil_paged_tokens(extend_input_len) + immediate = extend_input_len + self.page_size + extra_tokens + self.total_offset += immediate + max_new_tokens + self.current_offset += immediate + + +class SWAPrefillBudget(PrefillBudget): + """Separate FULL/SWA partitions, including per-request SWA rings.""" + + def __init__(self, *args, all_swa=False, **kwargs): + super().__init__(*args, **kwargs) + self.all_swa = all_swa + self.req_ring = getattr(self.allocator, "swa_req_ring", False) is True + + def _available_and_evictable(self): + if self.all_swa: + return ( + self.allocator.swa_available_size() + + self.tree_cache.swa_evictable_size() + ) + return ( + self.allocator.full_available_size() + self.tree_cache.full_evictable_size() + ) + + @property + def remaining_swa(self): + evictable = 0 if self.req_ring else self.tree_cache.swa_evictable_size() + return self.allocator.swa_available_size() + evictable - self.swa_offset + + def swa_tokens( + self, + extend_input_len, + max_new_tokens, + *, + chunk_limit=None, + swa_host_hit_length=0, + ): + if self.req_ring: + return self.allocator.swa_ring_cost_tokens + return estimate_swa_kv_tokens( + extend_input_len, + max_new_tokens, + sliding_window_size=self.tree_cache.sliding_window_size, + page_size=self.page_size, + allocation_limit=chunk_limit, + host_hit_length=swa_host_hit_length, + ) + + def swa_never_fits(self, extend_input_len, max_new_tokens, **kwargs): + needed = self.swa_tokens(extend_input_len, max_new_tokens, **kwargs) + return ( + needed > self.allocator.size_swa + if self.req_ring + else needed >= self.allocator.size_swa + ) + + def _chunk_cap(self, max_new_tokens, swa_host_hit_length=0): + headroom = self.swa_tokens( + 0, max_new_tokens, swa_host_hit_length=swa_host_hit_length + ) + cap = int(self.remaining_swa) - headroom + return max(0, cap // self.page_size * self.page_size) + + def check_prefill( + self, + *, + extend_input_len: int, + total_tokens: int, + max_new_tokens: int, + input_tokens: int, + swa_host_hit_length: int, + chunk_limit: int | None, + ) -> tuple[bool, int | None]: + if total_tokens >= self.remaining_total: + return False, None + extend_input_len = self.ceil_paged_tokens(extend_input_len) + needed = self.swa_tokens( + extend_input_len, + max_new_tokens, + chunk_limit=chunk_limit, + swa_host_hit_length=swa_host_hit_length, + ) + fits = ( + needed <= self.remaining_swa + if self.req_ring + else needed < self.remaining_swa + ) + if fits: + return True, chunk_limit + # Only permanent shortfalls may shrink a chunk. Transient pressure waits + # so a new prefill does not consume running decodes' window headroom. + cap = 0 + if self.swa_never_fits( + extend_input_len, + max_new_tokens, + chunk_limit=chunk_limit, + swa_host_hit_length=swa_host_hit_length, + ): + cap = self._chunk_cap(max_new_tokens, swa_host_hit_length) + if chunk_limit is None or cap <= 0: + return False, None + return True, min(chunk_limit, cap) + + def has_capacity(self) -> bool: + return super().has_capacity() and self.remaining_swa > 0 + + def available_chunk_tokens(self, chunk_limit: int) -> int | None: + available = min(chunk_limit, int(self.remaining_total)) + if not self.req_ring: + available = min(available, int(self.remaining_swa) - self.page_size) + return available if available > 0 else None + + def can_allocate_prefill( + self, + *, + paged_input: int, + extend_input_len: int, + max_new_tokens: int, + chunk_limit: int | None, + ) -> bool: + return ( + super().can_allocate_prefill( + paged_input=paged_input, + extend_input_len=extend_input_len, + max_new_tokens=max_new_tokens, + chunk_limit=chunk_limit, + ) + and self.swa_tokens( + extend_input_len, max_new_tokens, chunk_limit=chunk_limit + ) + <= self.remaining_swa + ) + + def reserve( + self, + extend_input_len: int, + max_new_tokens: int, + *, + extra_tokens: int = 0, + chunk_limit: int | None = None, + is_chunked_continuation: bool = False, + ) -> None: + super().reserve(extend_input_len, max_new_tokens, extra_tokens=extra_tokens) + # A continuation already owns its ring slot. + if not (self.req_ring and is_chunked_continuation): + self.swa_offset += self.swa_tokens( + self.ceil_paged_tokens(extend_input_len), + max_new_tokens, + chunk_limit=chunk_limit, + ) + + +class SharedSWAPrefillBudget(SWAPrefillBudget): + """FULL and SWA reservations compete for the same physical byte budget.""" + + def __init__(self, *args, num_mixed_decode_tokens=0, **kwargs): + super().__init__( + *args, num_mixed_decode_tokens=num_mixed_decode_tokens, **kwargs + ) + self.swa_offset = num_mixed_decode_tokens + + def _fits(self, full_tokens, swa_tokens, *, empty_pool=False): + return self.allocator.can_reserve( + full_tokens + (0 if empty_pool else self.total_offset), + swa_tokens + (0 if empty_pool else self.swa_offset), + full_evictable_tokens=0 + if empty_pool + else self.tree_cache.full_evictable_size(), + swa_evictable_tokens=0 + if empty_pool + else self.tree_cache.swa_evictable_size(), + empty_pool=empty_pool, + require_token_slack=empty_pool, + ) + + def _joint_chunk_cap(self, *, max_chunk_tokens, chunk_limit, swa_host_hit_length=0): + lo, hi = 0, max(0, max_chunk_tokens) // self.page_size + while lo < hi: + mid = (lo + hi + 1) // 2 + tokens = mid * self.page_size + if self._fits( + tokens + self.page_size, + self.swa_tokens( + tokens, + 0, + chunk_limit=chunk_limit, + swa_host_hit_length=swa_host_hit_length, + ), + ): + lo = mid + else: + hi = mid - 1 + return lo * self.page_size + + def check_prefill( + self, + *, + extend_input_len: int, + total_tokens: int, + max_new_tokens: int, + input_tokens: int, + swa_host_hit_length: int, + chunk_limit: int | None, + ) -> tuple[bool, int | None]: + needed = self.swa_tokens( + extend_input_len, + max_new_tokens, + chunk_limit=chunk_limit, + swa_host_hit_length=swa_host_hit_length, + ) + if self._fits(total_tokens, needed): + return True, chunk_limit + if chunk_limit is None or self._fits( + input_tokens + max_new_tokens + self.page_size, + self.swa_tokens(input_tokens, max_new_tokens, chunk_limit=chunk_limit), + empty_pool=True, + ): + return False, None + cap = self._joint_chunk_cap( + max_chunk_tokens=min(chunk_limit, max(0, extend_input_len - 1)), + chunk_limit=chunk_limit, + swa_host_hit_length=swa_host_hit_length, + ) + return (True, min(chunk_limit, cap)) if cap > 0 else (False, None) + + def has_capacity(self) -> bool: + return self._fits(0, 0) + + def available_chunk_tokens(self, chunk_limit: int) -> int | None: + return chunk_limit + + def fit_chunk( + self, + *, + extend_input_len: int, + max_new_tokens: int, + chunk_limit: int, + ) -> int | None: + candidate = min(extend_input_len, chunk_limit) + finishes = candidate >= extend_input_len + headroom = max_new_tokens if finishes else 0 + if self._fits( + self.ceil_paged_tokens(candidate) + headroom + self.page_size, + self.swa_tokens(candidate, headroom, chunk_limit=chunk_limit), + ): + return chunk_limit + cap = self._joint_chunk_cap( + max_chunk_tokens=max(0, candidate - 1) if finishes else candidate, + chunk_limit=chunk_limit, + ) + return min(chunk_limit, cap) if cap > 0 else None + + def can_allocate_prefill( + self, + *, + paged_input: int, + extend_input_len: int, + max_new_tokens: int, + chunk_limit: int | None, + ) -> bool: + return self._fits( + paged_input + max_new_tokens + self.page_size, + self.swa_tokens(extend_input_len, max_new_tokens, chunk_limit=chunk_limit), + ) diff --git a/test/registered/unit/managers/test_prefill_adder.py b/test/registered/unit/managers/test_prefill_adder.py index c3c1dd1cf4a0..805d561a4bf3 100644 --- a/test/registered/unit/managers/test_prefill_adder.py +++ b/test/registered/unit/managers/test_prefill_adder.py @@ -2,6 +2,8 @@ from types import SimpleNamespace from unittest.mock import MagicMock, patch +import torch + import sglang.srt.managers.schedule_policy as schedule_policy from sglang.srt.managers.schedule_batch import Req from sglang.srt.managers.schedule_policy import ( @@ -13,6 +15,12 @@ DecLockRefResult, IncLockRefResult, ) +from sglang.srt.mem_cache.prefill_budget import ( + PrefillBudget, + SWAPrefillBudget, + estimate_swa_kv_tokens, +) +from sglang.srt.mem_cache.unified_memory_pool import init_unified_swa_pools from sglang.srt.runtime_context import get_context from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler from sglang.srt.utils.common import Range @@ -71,6 +79,10 @@ def create_token_allocator( allocator.swa_available_size.return_value = swa_available_size allocator.available_size.return_value = available_size allocator.size_swa = size_swa + allocator.swa_req_ring = False + allocator.create_prefill_budget.side_effect = lambda tree_cache, **kwargs: ( + PrefillBudget(allocator, tree_cache, **kwargs) + ) return allocator def create_running_batch(self, reqs=None) -> MagicMock: @@ -100,6 +112,7 @@ def create_mock_req(self, rid, priority, max_new_tokens, output_len=0, wait_time req.time_stats = SimpleNamespace(wait_queue_entry_time=wait_time) req.retracted_stain = False req.host_hit_length = 0 + req.swa_host_hit_length = 0 req.storage_hit_length = 0 req.storage_hit_start = None req.host_hit_is_storage = False @@ -123,8 +136,96 @@ def create_adder(self, running_batch, **kwargs): priority_scheduling_preemption_threshold=0, ) defaults.update(kwargs) + defaults["token_to_kv_pool_allocator"].page_size = defaults["page_size"] return PrefillAdder(**defaults) + def create_shared_adder(self): + self.mock_tree_cache.supports_mamba.return_value = False + self.mock_tree_cache.sliding_window_size = 8 + self.mock_tree_cache.is_tree_cache.return_value = False + allocator = init_unified_swa_pools( + device="cpu", + kv_cache_dtype=torch.float16, + head_num=1, + head_dim=4, + v_head_dim=4, + swa_head_num=1, + swa_head_dim=4, + swa_v_head_dim=4, + page_size=4, + start_layer=0, + end_layer=2, + swa_attention_layer_ids=[1], + full_attention_layer_ids=[0], + total_bytes=1024, + enable_memory_saver=False, + need_sort=False, + lazy_compaction=True, + ).token_to_kv_pool_allocator + return self.create_adder( + self.create_running_batch(), + page_size=4, + rem_chunk_tokens=16, + token_to_kv_pool_allocator=allocator, + ) + + def create_shared_req(self, rid, max_new_tokens=4): + req = self.create_mock_req(rid, priority=0, max_new_tokens=max_new_tokens) + req.sampling_params.ignore_eos = False + req.swa_host_hit_length = 0 + req.last_node = MagicMock() + req.full_untruncated_fill_ids = list(range(12)) + req.set_extend_range = MagicMock( + side_effect=lambda start, end: setattr( + req, "extend_range", Range(start, end) + ) + ) + return req + + def test_shared_admission_reserves_all_pending_requests(self): + adder = self.create_shared_adder() + first, second = ( + self.create_shared_req("first"), + self.create_shared_req("second"), + ) + adder.add_one_req(first, has_chunked_req=False, truncation_align_size=None) + self.assertEqual(adder.can_run_list, [first]) + self.assertEqual( + adder.add_one_req( + second, has_chunked_req=False, truncation_align_size=None + ), + AddReqResult.NO_TOKEN, + ) + self.assertEqual(adder.can_run_list, [first]) + + def test_shared_admission_rechecks_after_prefix_lock(self): + adder = self.create_shared_adder() + self.assertIsNotNone(adder.token_to_kv_pool_allocator.alloc(24)) + self.mock_tree_cache.full_evictable_size.return_value = 24 + self.mock_tree_cache.swa_evictable_size.return_value = 24 + + def lock_prefix(_): + self.mock_tree_cache.full_evictable_size.return_value = 0 + self.mock_tree_cache.swa_evictable_size.return_value = 0 + return IncLockRefResult() + + self.mock_tree_cache.inc_lock_ref.side_effect = lock_prefix + req = self.create_shared_req("locked-prefix") + self.assertEqual( + adder.add_one_req(req, has_chunked_req=False, truncation_align_size=None), + AddReqResult.NO_TOKEN, + ) + self.mock_tree_cache.inc_lock_ref.assert_called_once() + self.assertEqual(adder.can_run_list, []) + + def test_shared_continuation_uses_memory_chunk_limit(self): + adder = self.create_shared_adder() + req = self.create_shared_req("continuation", max_new_tokens=80) + self.assertIs(adder.add_chunked_req(req), req) + self.assertEqual(req.extend_range.length, 8) + self.assertEqual(adder.memory_budget.total_offset, 12) + self.assertEqual(adder.memory_budget.swa_offset, 12) + def test_storage_prefetch_fulfillment_resolves_at_admission(self): adder = self.create_adder(self.create_running_batch()) req = self.create_mock_req("storage-hit", priority=0, max_new_tokens=1) @@ -184,7 +285,7 @@ def test_preempt_success_high_priority_values_first(self): running_batch = self.create_running_batch(running_reqs) adder = self.create_adder(running_batch) - self.assertEqual(adder.rem_total_token_offset, 225) + self.assertEqual(adder.memory_budget.total_offset, 225) self.mock_token_allocator.full_available_size.return_value = ( 225 # full occupation of GRam @@ -197,7 +298,9 @@ def test_preempt_success_high_priority_values_first(self): self.assertTrue(success) self.assertIn(running_reqs[0], adder.preempt_list) - self.assertEqual(adder.rem_total_token_offset, 175) # 50 + 75 + 100 - 50 = 175 + self.assertEqual( + adder.memory_budget.total_offset, 175 + ) # 50 + 75 + 100 - 50 = 175 running_batch.release_req.assert_called_once() def test_preempt_success_low_priority_values_first(self): @@ -214,7 +317,7 @@ def test_preempt_success_low_priority_values_first(self): running_batch = self.create_running_batch(running_reqs) adder = self.create_adder(running_batch) - self.assertEqual(adder.rem_total_token_offset, 225) + self.assertEqual(adder.memory_budget.total_offset, 225) self.mock_token_allocator.full_available_size.return_value = ( 225 # full occupation of GRam @@ -227,7 +330,9 @@ def test_preempt_success_low_priority_values_first(self): self.assertTrue(success) self.assertIn(running_reqs[2], adder.preempt_list) - self.assertEqual(adder.rem_total_token_offset, 125) # 50 + 75 + 100 - 100 = 125 + self.assertEqual( + adder.memory_budget.total_offset, 125 + ) # 50 + 75 + 100 - 100 = 125 running_batch.release_req.assert_called_once() def test_preempt_fail_low_priority_values_first(self): @@ -244,7 +349,7 @@ def test_preempt_fail_low_priority_values_first(self): running_batch = self.create_running_batch(running_reqs) adder = self.create_adder(running_batch) - self.assertEqual(adder.rem_total_token_offset, 225) + self.assertEqual(adder.memory_budget.total_offset, 225) self.mock_token_allocator.full_available_size.return_value = ( 225 # full occupation of GRam @@ -282,7 +387,7 @@ def test_preempt_fail_high_priority_values_first(self): running_batch = self.create_running_batch(running_reqs) adder = self.create_adder(running_batch) - self.assertEqual(adder.rem_total_token_offset, 225) + self.assertEqual(adder.memory_budget.total_offset, 225) self.mock_token_allocator.full_available_size.return_value = ( 225 # full occupation of GRam @@ -320,7 +425,7 @@ def test_preempt_skip_already_preempted_request(self): running_batch = self.create_running_batch(running_reqs) adder = self.create_adder(running_batch) - self.assertEqual(adder.rem_total_token_offset, 225) + self.assertEqual(adder.memory_budget.total_offset, 225) self.mock_token_allocator.full_available_size.return_value = 225 self.mock_token_allocator.available_size.return_value = 225 @@ -332,7 +437,7 @@ def test_preempt_skip_already_preempted_request(self): first_success = adder.preempt_to_schedule(first_req) self.assertTrue(first_success) self.assertIn(running_reqs[0], adder.preempt_list) - self.assertEqual(adder.rem_total_token_offset, 175) + self.assertEqual(adder.memory_budget.total_offset, 175) running_batch.release_req.assert_called_once() # Second call needs more tokens than currently free, so it would need to @@ -343,7 +448,7 @@ def test_preempt_skip_already_preempted_request(self): second_success = adder.preempt_to_schedule(second_req) self.assertFalse(second_success) - self.assertEqual(adder.rem_total_token_offset, 175) + self.assertEqual(adder.memory_budget.total_offset, 175) self.assertEqual(adder.preempt_list.count(running_reqs[0]), 1) running_batch.release_req.assert_called_once() @@ -363,7 +468,7 @@ def test_preempt_success_low_priority_values_first_exact_once(self): running_batch = self.create_running_batch(running_reqs) adder = self.create_adder(running_batch) - self.assertEqual(adder.rem_total_token_offset, 475) + self.assertEqual(adder.memory_budget.total_offset, 475) self.mock_token_allocator.full_available_size.return_value = ( 475 # full occupation of GRam @@ -376,7 +481,7 @@ def test_preempt_success_low_priority_values_first_exact_once(self): self.assertTrue(success) self.assertIn(running_reqs[2], adder.preempt_list) self.assertEqual( - adder.rem_total_token_offset, 375 + adder.memory_budget.total_offset, 375 ) # 50 + 75 + 100 + 125 + 125 - 100 = 375 running_batch.release_req.assert_called_once() @@ -396,7 +501,7 @@ def test_preempt_success_low_priority_values_first_exact_twice(self): running_batch = self.create_running_batch(running_reqs) adder = self.create_adder(running_batch) - self.assertEqual(adder.rem_total_token_offset, 475) + self.assertEqual(adder.memory_budget.total_offset, 475) self.mock_token_allocator.full_available_size.return_value = ( 475 # full occupation of GRam @@ -410,7 +515,7 @@ def test_preempt_success_low_priority_values_first_exact_twice(self): self.assertIn(running_reqs[2], adder.preempt_list) self.assertIn(running_reqs[3], adder.preempt_list) self.assertEqual( - adder.rem_total_token_offset, 250 + adder.memory_budget.total_offset, 250 ) # 50 + 75 + 100 + 125 + 125 - 100 - 125 = 250 self.assertEqual(running_batch.release_req.call_count, 2) @@ -432,8 +537,8 @@ def test_mixed_chunk_prefill_budgets(self): self.assertEqual(adder.rem_input_tokens, 192) # 200 - 8 self.assertEqual(adder.rem_chunk_tokens, 56) # 64 - 8 - self.assertEqual(adder.rem_total_token_offset, 408) # 8 + 8 * 50 - self.assertEqual(adder.cur_rem_token_offset, 8) + self.assertEqual(adder.memory_budget.total_offset, 408) # 8 + 8 * 50 + self.assertEqual(adder.memory_budget.current_offset, 8) self.assertEqual(adder.budget_state(), AddReqResult.CONTINUE) # Add a prefill that exactly consumes the chunk budget @@ -466,7 +571,7 @@ def test_mixed_chunk_prefill_budgets(self): self.assertEqual(adder2.rem_input_tokens, 195) # 200 - 5 self.assertEqual(adder2.rem_chunk_tokens, 59) # 64 - 5 - self.assertEqual(adder2.rem_total_token_offset, 255) # 5 + 5 * 50 + self.assertEqual(adder2.memory_budget.total_offset, 255) # 5 + 5 * 50 self.assertEqual(adder2.budget_state(), AddReqResult.CONTINUE) # Same prefill no longer exhausts the chunk budget @@ -521,6 +626,10 @@ def _build_hybrid_swa_chunked_req( rem_chunk_tokens=rem_chunk, ) adder.is_hybrid_swa = is_hybrid_swa + if is_hybrid_swa: + adder.memory_budget = SWAPrefillBudget( + self.mock_token_allocator, self.mock_tree_cache + ) req = self.create_mock_req("chunked", priority=0, max_new_tokens=128) req.prefix_indices = [] @@ -599,7 +708,16 @@ def test_swa_budget_for_req(self): page_size=page, rem_chunk_tokens=rem_chunk, ) - self.assertEqual(adder._swa_budget_for_req(extend, max_new), expected) + self.assertEqual( + estimate_swa_kv_tokens( + extend, + max_new, + sliding_window_size=window, + page_size=page, + allocation_limit=rem_chunk, + ), + expected, + ) def test_swa_admission_admits_short_cached_resume_at_two_window_pool(self): # Livelock regression (real incident). At an SWA pool ~= 2 sliding @@ -619,6 +737,9 @@ def test_swa_admission_admits_short_cached_resume_at_two_window_pool(self): self.mock_tree_cache.is_tree_cache.return_value = False adder = self.create_adder(self.create_running_batch(), page_size=PAGE) adder.is_hybrid_swa = True + adder.memory_budget = SWAPrefillBudget( + self.mock_token_allocator, self.mock_tree_cache + ) req = self.create_mock_req( "resume", priority=0, max_new_tokens=40, output_len=10 @@ -636,7 +757,9 @@ def test_swa_admission_admits_short_cached_resume_at_two_window_pool(self): req.sampling_params = SimpleNamespace(max_new_tokens=40, ignore_eos=False) # Pre-fix: a constant sliding-window reservation rejects the resume. - with patch.object(adder, "_swa_budget_for_req", return_value=WINDOW + PAGE): + with patch.object( + adder.memory_budget, "swa_tokens", return_value=WINDOW + PAGE + ): self.assertIs( adder.add_one_req( req, has_chunked_req=False, truncation_align_size=None @@ -888,6 +1011,11 @@ def create_swa_adder( ) -> PrefillAdder: self.mock_tree_cache.sliding_window_size = sliding_window self.mock_token_allocator = self.create_token_allocator(size_swa=size_swa) + self.mock_token_allocator.create_prefill_budget.side_effect = ( + lambda tree_cache, **kwargs: SWAPrefillBudget( + self.mock_token_allocator, tree_cache, **kwargs + ) + ) return self.create_adder( self.create_running_batch(), page_size=page_size, @@ -899,7 +1027,7 @@ def test_swa_never_fits_false_under_transient_pressure(self): # once running decodes drain -> must wait, not take the hatch. adder = self.create_swa_adder(size_swa=1024, sliding_window=128) self.assertFalse( - adder._swa_req_never_fits(extend_input_len=256, max_new_tokens=64) + adder.memory_budget.swa_never_fits(extend_input_len=256, max_new_tokens=64) ) def test_swa_never_fits_true_when_budget_exceeds_whole_pool(self): @@ -907,7 +1035,7 @@ def test_swa_never_fits_true_when_budget_exceeds_whole_pool(self): # pool: it can never fit however far the pool drains -> hatch. adder = self.create_swa_adder(size_swa=1024, sliding_window=128) self.assertTrue( - adder._swa_req_never_fits( + adder.memory_budget.swa_never_fits( extend_input_len=256, max_new_tokens=64, swa_host_hit_length=4096 ) ) @@ -917,14 +1045,14 @@ def test_swa_never_fits_is_gated_by_pool_capacity(self): # the budget against size_swa (guards against a wrong-accessor bug). req = dict(extend_input_len=256, max_new_tokens=64, swa_host_hit_length=600) self.assertTrue( - self.create_swa_adder(size_swa=512, sliding_window=128)._swa_req_never_fits( - **req - ) + self.create_swa_adder( + size_swa=512, sliding_window=128 + ).memory_budget.swa_never_fits(**req) ) self.assertFalse( self.create_swa_adder( size_swa=4096, sliding_window=128 - )._swa_req_never_fits(**req) + ).memory_budget.swa_never_fits(**req) ) diff --git a/test/registered/unit/managers/test_scheduler_init_req_max_new_tokens.py b/test/registered/unit/managers/test_scheduler_init_req_max_new_tokens.py index 43f182778e26..d651071e5d84 100644 --- a/test/registered/unit/managers/test_scheduler_init_req_max_new_tokens.py +++ b/test/registered/unit/managers/test_scheduler_init_req_max_new_tokens.py @@ -6,6 +6,7 @@ from sglang.srt.environ import envs from sglang.srt.managers.scheduler import Scheduler +from sglang.srt.mem_cache.allocator.token import TokenToKVPoolAllocator from sglang.srt.mem_cache.unified_memory_pool import init_unified_swa_pools from sglang.srt.runtime_context import get_parallel from sglang.test.ci.ci_register import register_cpu_ci @@ -59,7 +60,16 @@ def _new_scheduler( scheduler.max_total_num_tokens = max_total_num_tokens scheduler.page_size = page_size scheduler.max_new_tokens_limit = envs.SGLANG_MAX_NEW_TOKENS_LIMIT.get() - scheduler.token_to_kv_pool_allocator = None + scheduler.sliding_window_size = None + scheduler.chunked_prefill_size = None + scheduler.token_to_kv_pool_allocator = TokenToKVPoolAllocator( + size=max_total_num_tokens, + dtype=torch.int64, + device="cpu", + kvcache=None, + need_sort=False, + ) + scheduler.token_to_kv_pool_allocator.page_size = page_size return scheduler def _new_req(self, max_new_tokens, input_len: int = 8, min_new_tokens: int = 0): diff --git a/test/registered/unit/mem_cache/test_multi_ended_allocator.py b/test/registered/unit/mem_cache/test_multi_ended_allocator.py index aaed067bf844..615dfec53b75 100644 --- a/test/registered/unit/mem_cache/test_multi_ended_allocator.py +++ b/test/registered/unit/mem_cache/test_multi_ended_allocator.py @@ -29,7 +29,6 @@ import torch -from sglang.srt.managers.schedule_policy import estimate_swa_kv_tokens from sglang.srt.mem_cache.allocator.unified_hybrid_swa import ( UnifiedSWATokenToKVPoolAllocator, ) @@ -41,6 +40,7 @@ MultiEndedAllocator, ) from sglang.srt.mem_cache.base_prefix_cache import EvictParams +from sglang.srt.mem_cache.prefill_budget import estimate_swa_kv_tokens from sglang.srt.mem_cache.unified_cache.components import ComponentType from sglang.srt.mem_cache.unified_memory_pool import ( MambaSubPoolSpec, diff --git a/test/registered/unit/mem_cache/test_prefill_memory_budget.py b/test/registered/unit/mem_cache/test_prefill_memory_budget.py new file mode 100644 index 000000000000..3f08e40fdb5f --- /dev/null +++ b/test/registered/unit/mem_cache/test_prefill_memory_budget.py @@ -0,0 +1,224 @@ +"""CPU regressions for allocator-owned prefill admission and pending demand.""" + +import unittest +from types import SimpleNamespace +from unittest.mock import MagicMock + +import torch + +from sglang.srt.mem_cache.allocator.hisparse import ( + DeepSeekV4HiSparseTokenToKVPoolAllocator, +) +from sglang.srt.mem_cache.allocator.swa import ( + PureSWATokenToKVPoolAllocator, + SWATokenToKVPoolAllocator, +) +from sglang.srt.mem_cache.allocator.unified_hybrid_swa import ( + UnifiedMambaSWATokenToKVPoolAllocator, +) +from sglang.srt.mem_cache.common import evict_from_tree_cache +from sglang.srt.mem_cache.prefill_budget import SWAPrefillBudget +from sglang.srt.mem_cache.unified_memory_pool import init_unified_swa_pools +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + + +def _shared_allocator(): + return init_unified_swa_pools( + device="cpu", + kv_cache_dtype=torch.float16, + head_num=1, + head_dim=4, + v_head_dim=4, + swa_head_num=1, + swa_head_dim=4, + swa_v_head_dim=4, + page_size=4, + start_layer=0, + end_layer=2, + swa_attention_layer_ids=[1], + full_attention_layer_ids=[0], + total_bytes=1024, + enable_memory_saver=False, + need_sort=False, + lazy_compaction=True, + ).token_to_kv_pool_allocator + + +def _cache(): + return SimpleNamespace( + sliding_window_size=8, + full_evictable_size=lambda: 0, + swa_evictable_size=lambda: 0, + is_chunk_cache=lambda: False, + ) + + +class TestSharedPrefillMemoryBudget(unittest.TestCase): + def setUp(self): + self.allocator = _shared_allocator() + self.cache = _cache() + self.budget = self.allocator.create_prefill_budget(self.cache) + self.request = dict( + extend_input_len=12, + total_tokens=20, + max_new_tokens=4, + input_tokens=12, + swa_host_hit_length=0, + chunk_limit=16, + ) + + def test_pending_batch_cannot_spend_shared_bytes_twice(self): + self.assertEqual(self.budget.check_prefill(**self.request), (True, 16)) + self.budget.reserve(12, 4, chunk_limit=16) + # Each side separately has room, but their combined reservation does not. + self.assertGreater(self.budget.remaining_total, 20) + self.assertGreater(self.budget.remaining_swa, 16) + self.assertEqual(self.budget.check_prefill(**self.request), (False, None)) + self.assertEqual(self.allocator.full_attn_allocator.allocated_count(), 0) + self.assertEqual(self.allocator.swa_attn_allocator.allocated_count(), 0) + + def test_mixed_decode_reserves_both_sides(self): + budget = self.allocator.create_prefill_budget( + self.cache, num_mixed_decode_tokens=4 + ) + budget.reserve(12, 4, chunk_limit=16) + self.assertEqual( + (budget.total_offset, budget.current_offset, budget.swa_offset), + (24, 20, 20), + ) + self.assertEqual(budget.check_prefill(**self.request), (False, None)) + + def test_prefix_lock_changes_admission_without_rebuilding_budget(self): + self.assertIsNotNone(self.allocator.alloc(24)) + self.cache.full_evictable_size = lambda: 24 + self.cache.swa_evictable_size = lambda: 24 + self.assertEqual(self.budget.check_prefill(**self.request), (True, 16)) + # Locking the cached prefix removes its eviction credit. + self.cache.full_evictable_size = lambda: 0 + self.cache.swa_evictable_size = lambda: 0 + self.assertEqual(self.budget.check_prefill(**self.request), (False, None)) + + def test_final_chunk_reserves_decode_headroom(self): + limit = self.budget.fit_chunk( + extend_input_len=12, max_new_tokens=80, chunk_limit=16 + ) + self.assertEqual(limit, 8) + self.budget.reserve(limit, 0, chunk_limit=16, is_chunked_continuation=True) + self.assertIsNotNone(self.allocator.alloc(limit)) + + def test_host_swa_load_is_part_of_joint_demand(self): + self.assertEqual(self.budget.check_prefill(**self.request), (True, 16)) + request = {**self.request, "swa_host_hit_length": 32} + self.assertEqual(self.budget.check_prefill(**request), (False, None)) + + def test_prompt_clipping_uses_the_empty_pool(self): + kwargs = dict(token_capacity=1, sliding_window_size=8, chunk_size=16) + limit = self.allocator.max_new_tokens_for_memory(12, 80, **kwargs) + self.assertIsNotNone(limit) + self.assertGreater(limit, 0) + self.assertIsNotNone(self.allocator.alloc(24)) + self.assertEqual( + self.allocator.max_new_tokens_for_memory(12, 80, **kwargs), limit + ) + self.assertIsNone(self.allocator.max_new_tokens_for_memory(100, 0, **kwargs)) + + def test_shared_stats_pair_available_tokens_with_current_capacity(self): + self.assertIsNotNone(self.allocator.alloc(12)) + (full_capacity, full_free), (swa_capacity, swa_free) = ( + self.allocator.swa_capacity_and_available(full_capacity=1, swa_capacity=1) + ) + self.assertEqual(full_capacity - full_free, 12) + self.assertEqual(swa_capacity - swa_free, 12) + + def test_common_eviction_dispatches_joint_reclaim(self): + self.cache.token_to_kv_pool_allocator = self.allocator + self.allocator.evict_to_free_tokens = MagicMock() + evict_from_tree_cache(self.cache, 8) + self.allocator.evict_to_free_tokens.assert_called_once_with(self.cache, 8) + + +class TestFixedPrefillMemoryBudget(unittest.TestCase): + def _allocator(self, cls=SWATokenToKVPoolAllocator): + allocator = object.__new__(cls) + allocator.page_size = 4 + allocator._size_full = 128 + allocator._size_swa = 64 + allocator.full_available_size = lambda: 128 + allocator.swa_available_size = lambda: 64 + return allocator + + def test_ring_slot_reserved_once_and_evictable_tokens_give_no_credit(self): + allocator = self._allocator() + allocator._swa_req_ring = True + allocator._swa_ring_cost = 32 + cache = _cache() + cache.swa_evictable_size = lambda: 1000 + budget = allocator.create_prefill_budget(cache) + budget.reserve(12, 4, chunk_limit=16) + self.assertEqual(budget.remaining_swa, 32) + budget.reserve(12, 4, chunk_limit=16, is_chunked_continuation=True) + self.assertEqual(budget.remaining_swa, 32) + # The last exact ring slot remains admissible. + self.assertEqual( + budget.check_prefill( + extend_input_len=4, + total_tokens=12, + max_new_tokens=4, + input_tokens=4, + swa_host_hit_length=0, + chunk_limit=16, + ), + (True, 16), + ) + + def test_pure_swa_budget_reads_swa_capacity(self): + allocator = self._allocator(PureSWATokenToKVPoolAllocator) + allocator.full_available_size = lambda: 0 + budget = allocator.create_prefill_budget(_cache()) + self.assertEqual(budget.remaining_total, 64) + self.assertTrue(budget.has_capacity()) + + def test_hisparse_budget_reads_wrapper_capacity(self): + allocator = self._allocator(DeepSeekV4HiSparseTokenToKVPoolAllocator) + allocator.full_available_size = lambda: 12 + budget = allocator.create_prefill_budget(_cache()) + self.assertEqual(budget.remaining_total, 12) + self.assertEqual(budget.remaining_swa, 64) + + def test_tri_pool_keeps_fixed_admission_and_clipping(self): + allocator = self._allocator(UnifiedMambaSWATokenToKVPoolAllocator) + allocator.can_reserve = MagicMock( + side_effect=AssertionError("two-pool reservation") + ) + budget = allocator.create_prefill_budget(_cache()) + self.assertIs(type(budget), SWAPrefillBudget) + self.assertTrue(budget.has_capacity()) + self.assertEqual( + allocator.max_new_tokens_for_memory( + 5, + 100, + token_capacity=32, + sliding_window_size=8, + chunk_size=16, + ), + 19, + ) + + def test_tri_pool_eviction_does_not_reenter_common(self): + allocator = self._allocator(UnifiedMambaSWATokenToKVPoolAllocator) + allocator.available_size = lambda: 0 + allocator.full_available_size = lambda: 0 + allocator.swa_available_size = lambda: 0 + cache = _cache() + cache.token_to_kv_pool_allocator = allocator + cache.evict_for_alloc = MagicMock() + evict_from_tree_cache(cache, 8) + cache.evict_for_alloc.assert_called_once() + params = cache.evict_for_alloc.call_args.args[0] + self.assertEqual((params.num_tokens, params.swa_num_tokens), (8, 8)) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/mem_cache/test_unified_radix_allocation_eviction.py b/test/registered/unit/mem_cache/test_unified_radix_allocation_eviction.py index 68c8cf1262e5..5b99273aa8c7 100644 --- a/test/registered/unit/mem_cache/test_unified_radix_allocation_eviction.py +++ b/test/registered/unit/mem_cache/test_unified_radix_allocation_eviction.py @@ -554,7 +554,11 @@ def evict_leaf(_node_id, tracker): def test_common_helper_uses_allocation_aware_entry_point(self): tree_cache = MagicMock() tree_cache.is_chunk_cache.return_value = False - tree_cache.token_to_kv_pool_allocator.available_size.return_value = 30 + from sglang.srt.mem_cache.allocator.token import TokenToKVPoolAllocator + + allocator = object.__new__(TokenToKVPoolAllocator) + allocator.available_size = lambda: 30 + tree_cache.token_to_kv_pool_allocator = allocator evict_from_tree_cache(tree_cache, num_tokens=100) From 4fdb2a7fd847865bee68fcea77b66995cb9b9fa5 Mon Sep 17 00:00:00 2001 From: yhzhuang Date: Sun, 13 Sep 2026 18:25:59 -0700 Subject: [PATCH 37/55] Fix unaligned unified SWA prefill admission --- python/sglang/srt/mem_cache/prefill_budget.py | 4 +- .../mem_cache/test_prefill_memory_budget.py | 82 ++++++++++++++++++- 2 files changed, 81 insertions(+), 5 deletions(-) diff --git a/python/sglang/srt/mem_cache/prefill_budget.py b/python/sglang/srt/mem_cache/prefill_budget.py index 0811eb763354..965daaa55f9d 100644 --- a/python/sglang/srt/mem_cache/prefill_budget.py +++ b/python/sglang/srt/mem_cache/prefill_budget.py @@ -381,7 +381,7 @@ def fit_chunk( finishes = candidate >= extend_input_len headroom = max_new_tokens if finishes else 0 if self._fits( - self.ceil_paged_tokens(candidate) + headroom + self.page_size, + candidate + headroom + self.page_size, self.swa_tokens(candidate, headroom, chunk_limit=chunk_limit), ): return chunk_limit @@ -400,6 +400,6 @@ def can_allocate_prefill( chunk_limit: int | None, ) -> bool: return self._fits( - paged_input + max_new_tokens + self.page_size, + extend_input_len + max_new_tokens + self.page_size, self.swa_tokens(extend_input_len, max_new_tokens, chunk_limit=chunk_limit), ) diff --git a/test/registered/unit/mem_cache/test_prefill_memory_budget.py b/test/registered/unit/mem_cache/test_prefill_memory_budget.py index 3f08e40fdb5f..2a73dfffc9e3 100644 --- a/test/registered/unit/mem_cache/test_prefill_memory_budget.py +++ b/test/registered/unit/mem_cache/test_prefill_memory_budget.py @@ -1,11 +1,15 @@ """CPU regressions for allocator-owned prefill admission and pending demand.""" import unittest +from array import array from types import SimpleNamespace from unittest.mock import MagicMock import torch +from sglang.srt.managers.schedule_batch import Req +from sglang.srt.managers.schedule_policy import PrefillAdder +from sglang.srt.managers.scheduler import Scheduler from sglang.srt.mem_cache.allocator.hisparse import ( DeepSeekV4HiSparseTokenToKVPoolAllocator, ) @@ -19,12 +23,14 @@ from sglang.srt.mem_cache.common import evict_from_tree_cache from sglang.srt.mem_cache.prefill_budget import SWAPrefillBudget from sglang.srt.mem_cache.unified_memory_pool import init_unified_swa_pools +from sglang.srt.runtime_context import get_parallel +from sglang.srt.sampling.sampling_params import SamplingParams from sglang.test.ci.ci_register import register_cpu_ci register_cpu_ci(est_time=5, suite="base-a-test-cpu") -def _shared_allocator(): +def _shared_allocator(*, page_size=4, total_bytes=1024): return init_unified_swa_pools( device="cpu", kv_cache_dtype=torch.float16, @@ -34,12 +40,12 @@ def _shared_allocator(): swa_head_num=1, swa_head_dim=4, swa_v_head_dim=4, - page_size=4, + page_size=page_size, start_layer=0, end_layer=2, swa_attention_layer_ids=[1], full_attention_layer_ids=[0], - total_bytes=1024, + total_bytes=total_bytes, enable_memory_saver=False, need_sort=False, lazy_compaction=True, @@ -139,6 +145,76 @@ def test_common_eviction_dispatches_joint_reclaim(self): self.allocator.evict_to_free_tokens.assert_called_once_with(self.cache, 8) +class TestSharedPrefillAdmission(unittest.TestCase): + def _new_admission(self, page_size, pool_pages, *, ignore_eos=False): + allocator = _shared_allocator( + page_size=page_size, total_bytes=pool_pages * page_size * 16 + ) + req = Req( + rid="unaligned-prompt", + origin_input_text=None, + origin_input_ids=array("q", [1] * (page_size + 1)), + sampling_params=SamplingParams(max_new_tokens=1, ignore_eos=ignore_eos), + ) + scheduler = Scheduler.__new__(Scheduler) + scheduler.max_req_len = 16 * page_size + scheduler.max_total_num_tokens = allocator.size_full + scheduler.page_size = page_size + scheduler.max_new_tokens_limit = None + scheduler.sliding_window_size = page_size + scheduler.chunked_prefill_size = page_size + scheduler.token_to_kv_pool_allocator = allocator + with get_parallel().override(attn_dcp_size=1): + scheduler.init_req_max_new_tokens(req) + self.assertEqual(req.sampling_params.max_new_tokens, 1) + req._refresh_fill_ids() + + cache = SimpleNamespace( + sliding_window_size=page_size, + disable=True, + full_evictable_size=lambda: 0, + swa_evictable_size=lambda: 0, + is_chunk_cache=lambda: True, + supports_mamba=lambda: False, + ) + adder = PrefillAdder( + page_size=page_size, + tree_cache=cache, + token_to_kv_pool_allocator=allocator, + running_batch=None, + new_token_ratio=1.0, + rem_input_tokens=16 * page_size, + rem_chunk_tokens=page_size, + ) + return allocator, req, adder + + def test_unaligned_final_chunk_makes_progress(self): + for page_size in (4, 64): + with self.subTest(page_size=page_size): + allocator, req, adder = self._new_admission(page_size, pool_pages=7) + req.prefix_indices = allocator.alloc(page_size) + self.assertIsNotNone(req.prefix_indices) + self.assertTrue(allocator.can_reserve(page_size + 2, page_size + 2)) + + self.assertIsNone(adder.add_chunked_req(req)) + self.assertEqual(adder.can_run_list, [req]) + self.assertEqual(req.extend_range.length, 1) + + def test_unaligned_ignore_eos_enters_empty_pool(self): + for page_size in (4, 64): + with self.subTest(page_size=page_size): + allocator, req, adder = self._new_admission( + page_size, pool_pages=6, ignore_eos=True + ) + self.assertEqual(len(req.prefix_indices), 0) + self.assertTrue(allocator.can_reserve(2 * page_size + 2, 2 * page_size)) + + adder.add_one_req( + req, has_chunked_req=False, truncation_align_size=None + ) + self.assertEqual(adder.can_run_list, [req]) + + class TestFixedPrefillMemoryBudget(unittest.TestCase): def _allocator(self, cls=SWATokenToKVPoolAllocator): allocator = object.__new__(cls) From 089673a0e6067b4a944db17bd9c01fd45f91f40d Mon Sep 17 00:00:00 2001 From: yhzhuang Date: Sun, 13 Sep 2026 18:57:43 -0700 Subject: [PATCH 38/55] Refactor unified SWA allocators into sibling capacity policies Retain shared virtual-ID lifecycle in a common base, gate eager-pool post-capture planning, and preserve capacity realization results. Keep allocator-owned admission policy and existing tri-pool behavior. --- python/sglang/srt/arg_groups/overrides.py | 3 + python/sglang/srt/mem_cache/allocator/base.py | 7 +- .../mem_cache/allocator/unified_hybrid_swa.py | 1314 +++++++++-------- python/sglang/srt/mem_cache/common.py | 6 +- .../srt/mem_cache/kv_cache_configurator.py | 8 +- .../srt/mem_cache/kv_index_translator.py | 6 +- .../unified_cache/components/swa_component.py | 4 +- .../unit/mem_cache/test_unified_tri_pool.py | 8 + 8 files changed, 697 insertions(+), 659 deletions(-) diff --git a/python/sglang/srt/arg_groups/overrides.py b/python/sglang/srt/arg_groups/overrides.py index b49b490a6f69..0fe079fdaa4e 100644 --- a/python/sglang/srt/arg_groups/overrides.py +++ b/python/sglang/srt/arg_groups/overrides.py @@ -1777,6 +1777,9 @@ def post_capture_kv_sizing_planned(server_args: Any) -> bool: mla_enabled = use_mla_backend(server_args) if not envs.SGLANG_ENABLE_POST_CAPTURE_KV_SIZING.get(): return False + # Unified arenas are fully backed before capture and cannot resize afterward. + if cfg.enable_unified_memory: + return False if cfg.device != "cuda": return False if cfg.dcp_size != 1: diff --git a/python/sglang/srt/mem_cache/allocator/base.py b/python/sglang/srt/mem_cache/allocator/base.py index 5bc0b6567f9f..a91f7edc3f29 100644 --- a/python/sglang/srt/mem_cache/allocator/base.py +++ b/python/sglang/srt/mem_cache/allocator/base.py @@ -95,9 +95,12 @@ def max_new_tokens_for_memory( 0, min(max_new_tokens, token_capacity - paged_input - self.page_size - 1) ) - def evict_to_free_tokens(self, tree_cache, num_tokens: int) -> None: + def evict_to_free_tokens(self, tree_cache, num_tokens: int) -> bool | None: """Evict unlocked prefix-cache entries until this allocator can serve - ``num_tokens`` or nothing evictable remains.""" + ``num_tokens`` or nothing evictable remains. + + Return whether capacity was realized, or None if it still needs checking. + """ from sglang.srt.mem_cache.base_prefix_cache import EvictParams if tree_cache is None or tree_cache.is_chunk_cache(): diff --git a/python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py b/python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py index 3d72723f6648..7a2ceb27b4b0 100644 --- a/python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py +++ b/python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py @@ -18,6 +18,7 @@ import logging import math +from abc import abstractmethod from typing import Callable, List, Optional, Sequence, Tuple import torch @@ -42,11 +43,10 @@ logger = logging.getLogger(__name__) -class UnifiedSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator): - """Composite allocator for the hybrid SWA pair (full + swa MHA sub-pools). +class UnifiedSWAAllocatorBase(SWATokenToKVPoolAllocator): + """Shared FULL/SWA virtual IDs, allocation lifecycle, and index translation. - One alloc(N) binds N pages on BOTH sides under the same virtual id, so - `available_size()` (joint bytes, in TOKENS) is the only safe alloc pre-check. + Concrete allocators define the two-ended or Mamba/SWA/FULL capacity policy. """ # Parent's `size` property has no setter but base init does `self.size = size`; @@ -177,18 +177,6 @@ def __init__( # -- construction hooks (the tri-pool subclass overrides both) -- - def _build_swa_attn_allocator(self, **kwargs) -> MultiEndedAllocator: - """The swa sub-allocator: an END pool in the 2-pool pair.""" - return MultiEndedAllocator( - sub_pool_name="swa", - is_id_owner=False, # non-owner; consumes virtuals minted by full - **kwargs, - ) - - def _wire_peers(self) -> None: - self.full_attn_allocator.bind_peer(self.swa_attn_allocator) - self.swa_attn_allocator.bind_peer(self.full_attn_allocator) - # -- capacity reporting (three-way split) -- def available_size(self) -> int: @@ -199,46 +187,6 @@ def available_size(self) -> int: self._joint_avail_memo_epoch = epoch return self._joint_avail_memo_tokens - def _compute_available_size(self) -> int: - """Joint byte budget in TOKENS: each composite alloc(1) consumes one - full-side AND one swa-side page under the same virtual id.""" - fa, sa = self.full_attn_allocator, self.swa_attn_allocator - e_f = fa.entry_bytes_per_page - e_s = sa.entry_bytes_per_page - # Direction-agnostic shared gap: the free byte band between the two pools. - if fa.grow_direction == "up": - gap_bytes = max(0, sa._byte_low_frontier() - fa._byte_high_frontier()) - else: - gap_bytes = max(0, fa._byte_low_frontier() - sa._byte_high_frontier()) - R_f = fa.num_pages - fa.min_page_index - fa._allocated_pages() - R_s = sa.num_pages - sa.min_page_index - sa._allocated_pages() - - if not self.lazy_compaction: - pages_by_bytes = gap_bytes // (e_f + e_s) - return min(pages_by_bytes, R_f, R_s) * self.page_size - - H_f = len(fa._free_phys_pages) - H_s = len(sa._free_phys_pages) - - K1 = min(H_f, H_s) # Phase 1: both drain - - # Phase 2: fewer-holes side extends; more-holes side keeps draining. - if H_f <= H_s: - e_phase2 = e_f - K_phase2_max = H_s - else: - e_phase2 = e_s - K_phase2_max = H_f - K2_room = K_phase2_max - K1 - K2 = min(K2_room, gap_bytes // e_phase2) if e_phase2 > 0 else K2_room - gap_bytes -= K2 * e_phase2 - - K3 = gap_bytes // (e_f + e_s) # Phase 3: both extend - - K_total = K1 + K2 + K3 - K_total = min(K_total, H_f + R_f, H_s + R_s) # index-space caps - return K_total * self.page_size - # Slot-conservation views for the leak invariant only; the byte-coordinated # value would flag spurious leaks. `allocated_count()` is in TOKENS. def _conserve_full_available_size(self) -> int: @@ -281,15 +229,6 @@ def schedulable_full_available_size(self) -> int: def schedulable_swa_available_size(self) -> int: return self.swa_attn_allocator.schedulable_available_size() - def _flush_targets(self): - """Flush ALL members, including ones that are not short themselves: a - one-sided hole is unusable, and compacting it yields SHARED gap.""" - return (self.full_attn_allocator, self.swa_attn_allocator) - - def _ask_float_for_room(self, need_tokens: int) -> None: - """No float in a two-END chain -- nothing can slide.""" - return None - # `size_full` / `size_swa` bound each side independently; current capacities # also account for the peer's live byte usage. @@ -301,618 +240,747 @@ def current_full_capacity(self) -> int: def current_swa_capacity(self) -> int: return self.swa_available_size() + self.swa_attn_allocator.allocated_count() - def reclaim_plan( + @property + def draft_virtual_id_space(self) -> int: + return self.full_attn_allocator.max_slots - 1 + + def debug_print(self) -> str: + return ( + f"#full-available={self.full_attn_allocator.available_size()}, " + f"#swa-available={self.swa_attn_allocator.available_size()}, " + f"#joint-available={self.available_size()}" + ) + + def get_kvcache(self): + return self._kvcache + + def translate_kv_loc( self, - full_tokens: int | float, - swa_tokens: int | float, + loc: torch.Tensor, *, - full_evictable_tokens: int = 0, - swa_evictable_tokens: int = 0, - empty_pool: bool = False, - ) -> Optional[Tuple[int, int]]: - """Return cumulative FULL/SWA eviction targets, or None if impossible. + out: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """Full-layer read path: virtual TOKEN ids -> full-physical TOKEN ids. + ``out=`` writes in place, for cuda-graph buffer stability.""" + result = self.full_attn_allocator.translate_kv_loc(loc, out=out) + return result - The tree evicts FULL before SWA. Minimize required SWA reclaim with all - evictable FULL available, then trim excess FULL reclaim. FULL eviction's - actual SWA cascade is counted by the tree's shared eviction tracker. - """ - if full_tokens < 0 or swa_tokens < 0: - return None + def translate_loc_from_full_to_swa( + self, + kv_indices: torch.Tensor, + *, + out: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """SWA-layer read path: virtual TOKEN ids -> swa kernel-facing ids.""" + return self.swa_attn_allocator.translate_kv_loc_for_kernel(kv_indices, out=out) - page_size = self.page_size - full_pages = (math.ceil(full_tokens) + page_size - 1) // page_size - swa_pages = (math.ceil(swa_tokens) + page_size - 1) // page_size - # Restoring host SWA for device-resident FULL can require more new - # SWA pages than new FULL pages; only the shared budget constrains it. + @property + def kernel_page_multiplier(self) -> int: + return self.full_attn_allocator.kernel_page_multiplier - compacted = empty_pool or not self.lazy_compaction or self._compaction_allowed() + @property + def full_v2p_page_table(self) -> torch.Tensor: + """Page-level virtual->physical table of the full sub-pool.""" + return self.full_attn_allocator.virtual_to_physical - def fits(full_reclaim_pages: int, swa_reclaim_pages: int) -> bool: - return self._fits_page_demand( - full_pages, - swa_pages, - full_reclaim_pages=full_reclaim_pages, - swa_reclaim_pages=swa_reclaim_pages, - compacted=compacted, - empty_pool=empty_pool, - ) + @property + def full_p2v_page_table(self) -> torch.Tensor: + """Page-level physical->virtual table of the full sub-pool.""" + return self.full_attn_allocator.physical_to_virtual - if empty_pool: - return (0, 0) if fits(0, 0) else None - if fits(0, 0): - return (0, 0) + def translate_kv_loc_for_kernel( + self, + loc: torch.Tensor, + *, + out: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """Full-pool virtual TOKEN ids -> kernel-facing ids.""" + return self.full_attn_allocator.translate_kv_loc_for_kernel(loc, out=out) - max_full_pages = min( - self.full_attn_allocator.allocated_count() // page_size, - max(0, int(full_evictable_tokens)) // page_size, - ) - max_swa_pages = min( - self.swa_attn_allocator.allocated_count() // page_size, - max(0, int(swa_evictable_tokens)) // page_size, + def translate_write_loc_for_kernel( + self, + loc: torch.Tensor, + *, + out: Optional[torch.Tensor] = None, + out_width: Optional[int] = None, + ) -> torch.Tensor: + """Widened virtual WRITE loc -> kernel-facing id. DCP is rejected for this + composite at argument validation, so it coincides with the read translate.""" + return self.full_attn_allocator.translate_write_loc_for_kernel( + loc, out=out, out_width=out_width ) - if not fits(max_full_pages, max_swa_pages): - return None - def first_fit(high: int, predicate: Callable[[int], bool]) -> int: - low = 0 - while low < high: - mid = (low + high) // 2 - if predicate(mid): - high = mid - else: - low = mid + 1 - return low + @property + def swa_kernel_page_multiplier(self) -> int: + return self.swa_attn_allocator.kernel_page_multiplier - swa_reclaim_pages = first_fit( - max_swa_pages, lambda value: fits(max_full_pages, value) - ) - full_reclaim_pages = first_fit( - max_full_pages, lambda value: fits(value, swa_reclaim_pages) - ) - return ( - full_reclaim_pages * page_size, - swa_reclaim_pages * page_size, - ) + @property + def swa_v2p_page_table(self) -> torch.Tensor: + """Page-level virtual->physical table of the SWA sub-pool.""" + return self.swa_attn_allocator.virtual_to_physical - def can_reserve( - self, - full_tokens: int | float, - swa_tokens: int | float, - *, - full_evictable_tokens: int = 0, - swa_evictable_tokens: int = 0, - empty_pool: bool = False, - require_token_slack: bool = False, - ) -> bool: - """Check pending FULL/SWA demand against the shared byte envelope. + # -- alloc -- - Scheduler admission keeps the historical one-token strict slack at an - empty-pool boundary. Live admission checks the state reachable after - reclaiming the currently evictable pages from both sides. - """ - if full_tokens < 0 or swa_tokens < 0: - return False - if require_token_slack and ( - full_tokens >= self.size_full or swa_tokens >= self.size_swa - ): - return False + def alloc(self, need_size: int) -> Optional[torch.Tensor]: + with record_function("UnifiedSWAAlloc.alloc"): + if not self.ensure_capacity(need_size, need_size): + return None + # Snapshot the virtual PAGES full will consume, to bind them on swa too. + num_pages = need_size // self.page_size + fa = self.full_attn_allocator + new_virtual_pages = fa.free_virtual_ids[:num_pages].clone() - return ( - self.reclaim_plan( - full_tokens, - swa_tokens, - full_evictable_tokens=full_evictable_tokens, - swa_evictable_tokens=swa_evictable_tokens, - empty_pool=empty_pool, + v_tokens = fa.alloc(need_size) + # Post-pre-check failure can only be internal-state inconsistency. + assert v_tokens is not None, ( + "UnifiedSWA.alloc: full.alloc returned None after joint " + "pre-check passed — internal-state inconsistency" ) - is not None - ) - - def _compaction_allowed(self) -> bool: - return all( - allocator.disagg_move_gate is None or allocator.disagg_move_gate() - for allocator in (self.full_attn_allocator, self.swa_attn_allocator) - ) + self.swa_attn_allocator.alloc_with_virtual(new_virtual_pages) + return v_tokens - def _fits_page_demand( + def alloc_extend( self, - num_full_pages: int, - num_swa_pages: int, - *, - full_reclaim_pages: int = 0, - swa_reclaim_pages: int = 0, - compacted: bool, - empty_pool: bool = False, - ) -> bool: - """Check one FULL/SWA page demand against a single allocator snapshot.""" - if min(num_full_pages, num_swa_pages) < 0: - return False - - fa, sa = self.full_attn_allocator, self.swa_attn_allocator - if empty_pool: - full_live_pages = swa_live_pages = 0 - full_reclaim_pages = swa_reclaim_pages = 0 - else: - full_live_pages = fa.allocated_count() // self.page_size - swa_live_pages = sa.allocated_count() // self.page_size - full_reclaim_pages = min(full_live_pages, max(0, int(full_reclaim_pages))) - swa_reclaim_pages = min(swa_live_pages, max(0, int(swa_reclaim_pages))) + prefix_lens: torch.Tensor, + prefix_lens_cpu: torch.Tensor, + seq_lens: torch.Tensor, + seq_lens_cpu: torch.Tensor, + last_loc: torch.Tensor, + extend_num_tokens: int, + ) -> Optional[torch.Tensor]: + """Paged extend; returns virtual TOKEN ids. The same virtual page maps to + full- and swa-physical, so swa binds exactly what the full kernel consumed.""" + with record_function("UnifiedSWAAlloc.alloc_extend"): + num_new_pages = get_num_new_pages( + seq_lens=seq_lens_cpu, + page_size=self.page_size, + prefix_lens=prefix_lens_cpu, + ) + need_tokens = num_new_pages * self.page_size + if not self.ensure_capacity(need_tokens, need_tokens): + return None - full_live_pages -= full_reclaim_pages - swa_live_pages -= swa_reclaim_pages - full_total_pages = full_live_pages + num_full_pages - swa_total_pages = swa_live_pages + num_swa_pages - virtual_page_capacity = fa.num_virtual_ids - fa.min_page_index - full_page_capacity = min( - virtual_page_capacity, - fa.num_pages - fa.min_page_index, - ) - swa_page_capacity = min( - virtual_page_capacity, - sa.num_pages - sa.min_page_index, - ) - if full_total_pages > full_page_capacity or swa_total_pages > swa_page_capacity: - return False + # Snapshot the virtual PAGES the kernel will consume; clone so swa keeps + # its view after the slice is consumed. + fa = self.full_attn_allocator + new_virtual_pages = fa.free_virtual_ids[:num_new_pages].clone() - if compacted: - full_virtual_room = virtual_page_capacity - full_live_pages - full_holes = swa_holes = 0 - full_index_room = full_page_capacity - full_live_pages - swa_index_room = swa_page_capacity - swa_live_pages - gap_bytes = ( - self._empty_shared_gap_bytes - - full_live_pages * fa.entry_bytes_per_page - - swa_live_pages * sa.entry_bytes_per_page + out_indices = fa.alloc_extend( + prefix_lens, + prefix_lens_cpu, + seq_lens, + seq_lens_cpu, + last_loc, + extend_num_tokens, + num_new_pages=num_new_pages, + ) + assert out_indices is not None, ( + "UnifiedSWA.alloc_extend: full.alloc_extend returned None " + "after joint pre-check passed — internal-state inconsistency" ) + self.swa_attn_allocator.alloc_with_virtual(new_virtual_pages) + return out_indices # virtual TOKEN ids + + def alloc_decode( + self, + seq_lens: torch.Tensor, + seq_lens_cpu: torch.Tensor, + last_loc: torch.Tensor, + ) -> Optional[torch.Tensor]: + """Paged decode: one new token per request, consuming a page only when the + decode wraps.""" + with record_function("UnifiedSWAAlloc.alloc_decode"): + num_new_pages = get_num_new_pages( + seq_lens=seq_lens_cpu, page_size=self.page_size, decode=True + ) + need_tokens = num_new_pages * self.page_size + if not self.ensure_capacity(need_tokens, need_tokens): + return None + + fa = self.full_attn_allocator + new_virtual_pages = fa.free_virtual_ids[:num_new_pages].clone() + + out_indices = fa.alloc_decode(seq_lens, seq_lens_cpu, last_loc) + assert out_indices is not None, ( + "UnifiedSWA.alloc_decode: full.alloc_decode returned None " + "after joint pre-check passed — internal-state inconsistency" + ) + + if new_virtual_pages.numel() > 0: + self.swa_attn_allocator.alloc_with_virtual(new_virtual_pages) + + return out_indices # virtual TOKEN ids + + def is_slot_allocated(self, slot: int) -> bool: + """Token-slot surface = the full side (which owns the virtual ids).""" + return self.full_attn_allocator.is_slot_allocated(slot) + + def allocator_state_str(self) -> str: + return self.full_attn_allocator.allocator_state_str() + + # -- free -- + + def free(self, free_index: torch.Tensor) -> None: + with record_function("UnifiedSWAAlloc.free"): + if free_index is None or free_index.numel() == 0: + return + if self.free_group is not None: + self.free_group.append(self._copy_for_free_group(free_index)) + return + # Order is not load-bearing: the per-sub-pool v2p IS the mapping. Only + # swa needs the tombstone filter; full owns the ids, so all are bound. + v = free_index.detach().to(torch.int64) + v_pages = v // self.page_size + swa_v2p_pages = self.swa_attn_allocator.virtual_to_physical[v_pages] + # `> 0` strict: -1 = tombstoned, 0 = padding-sink page; both skipped. + live_token_mask = swa_v2p_pages > 0 + live_tokens = v[live_token_mask] + if live_tokens.numel() > 0: + self.swa_attn_allocator.free(live_tokens) + self.full_attn_allocator.free(v) + self.full_attn_allocator.clear_inverse_history() + self.swa_attn_allocator.clear_inverse_history() + + def free_swa(self, free_index: torch.Tensor) -> None: + """SWA tombstone path: release swa-physical, keep the virtual id and + full-physical live; `swa.v2p_page[v_page] = -1` IS the tombstone.""" + if free_index is None or free_index.numel() == 0: + return + v = free_index.detach().to(torch.int64) + ps = self.page_size + v_pages = v // ps + # `> 0` strict: -1 = tombstoned, page 0 = padding sink (never freeable). + swa_v2p_pages = self.swa_attn_allocator.virtual_to_physical[v_pages] + live = v[swa_v2p_pages > 0] + if live.numel() == 0: + return + if ps == 1: + # token == page and the live filter just deduped against the v2p + # table, so these ARE unique page ids -- same skip as `_free_lazy`. + self.swa_attn_allocator.free(live, _pages=live) else: - full_virtual_room = len(fa.free_virtual_ids) + full_reclaim_pages - full_holes = len(fa._free_phys_pages) + full_reclaim_pages - swa_holes = len(sa._free_phys_pages) + swa_reclaim_pages - full_index_room = fa.num_pages - fa.min_page_index - fa._allocated_pages() - swa_index_room = sa.num_pages - sa.min_page_index - sa._allocated_pages() - gap_bytes = fa._current_gap_bytes() + self.swa_attn_allocator.free(live) + self.swa_attn_allocator.clear_inverse_history() - if num_full_pages > full_virtual_room: - return False - if num_full_pages > full_holes + full_index_room: - return False - if num_swa_pages > swa_holes + swa_index_room: - return False - full_extensions = max(0, num_full_pages - full_holes) - swa_extensions = max(0, num_swa_pages - swa_holes) - return ( - full_extensions * fa.entry_bytes_per_page - + swa_extensions * sa.entry_bytes_per_page - <= max(0, gap_bytes) + def free_swa_segment(self, free_index: torch.Tensor, *, start_pos: int) -> None: + """free_swa() for a kv-row segment: `start_pos` promises a contiguous + ascending range, so page reps come from stride arithmetic, not `torch.unique`.""" + if free_index is None or free_index.numel() == 0: + return + if self.page_size == 1: + self.free_swa(free_index) + return + ps = self.page_size + reps = self.swa_attn_allocator._page_reps( + free_index.detach().to(torch.int64), start_pos ) + # Keep only pages still bound on swa; freeing a tombstoned one would + # corrupt the hole list. `> 0` strict: -1 tombstoned, 0 padding sink. + rep_pages = reps // ps + swa_v2p_pages = self.swa_attn_allocator.virtual_to_physical[rep_pages] + live_reps = reps[swa_v2p_pages > 0] + if live_reps.numel() == 0: + return + self.swa_attn_allocator.free(live_reps, _pages=live_reps // ps) + self.swa_attn_allocator.clear_inverse_history() + + def free_full(self, free_index: torch.Tensor) -> None: + """Release the full-physical page and the virtual id, leaving the swa + side alone -- the caller already tombstoned it (`swa.v2p_page == -1`).""" + if free_index is None or free_index.numel() == 0: + return + if self.free_group is not None: + self.full_free_group.append(self._copy_for_free_group(free_index)) + return + self.full_attn_allocator.free(free_index.detach().to(torch.int64)) + self.full_attn_allocator.clear_inverse_history() + + def free_full_segment(self, free_index: torch.Tensor, *, start_pos: int) -> None: + if free_index is None or free_index.numel() == 0: + return + if self.page_size == 1: + # token == page: free_full already frees by exact ids, no dedup. + self.free_full(free_index) + return + # The swa v2p is the mapping, so a tombstoned swa page drops out of the + # two-sided segment path by itself; full-only is the same call. + self.free_segment(free_index, start_pos=start_pos) + + def set_full_to_swa_mapping( + self, full_indices: torch.Tensor, swa_indices: torch.Tensor + ) -> None: + """No-op stub for HiCache load-back: in shared mode the swa v2p IS the + mapping, and HiCache for shared SWA is out of scope.""" + return + + def clear_full_to_swa_mapping(self, full_indices: torch.Tensor) -> None: + # Paired with set_full_to_swa_mapping: shared mode has no mapping tensor. + return + + # -- free-group -- + + # Not the SWA parent's hooks: those open the parent's paged full allocator + # as a free group, and this composite's sub-pools defer on their own. + def free_group_begin(self) -> None: + BaseTokenToKVPoolAllocator.free_group_begin(self) + self.free_page_reps_group = [] + self.full_free_group = [] + + def free_group_end(self) -> None: + pending, self.free_page_reps_group = self.free_page_reps_group, None + full_free_group, self.full_free_group = self.full_free_group, [] + BaseTokenToKVPoolAllocator.free_group_end(self) + if full_free_group: + self.full_attn_allocator.free(torch.cat(full_free_group)) + self.full_attn_allocator.clear_inverse_history() + if pending: + self._release_page_reps(pending) + + def free_segment(self, free_index: torch.Tensor, *, start_pos: int) -> None: + """Fixed-shape counterpart of `free()`; see `MultiEndedAllocator._page_reps`. + Both sides share one page-rep derivation instead of dedup'ing twice.""" + if free_index is None or free_index.numel() == 0: + return + if self.page_size == 1: + self.free(free_index) + return + reps = self.full_attn_allocator._page_reps( + free_index.detach().to(torch.int64), start_pos + ) + if self.free_page_reps_group is None: + self._release_page_reps((reps,)) + else: + self.free_page_reps_group.append(reps) + + def _release_page_reps(self, pieces: Sequence[torch.Tensor]) -> None: + reps = pieces[0] if len(pieces) == 1 else torch.cat(tuple(pieces)) + v_pages = reps // self.page_size + # Same tombstone filter as `free`, but at PAGE granularity (page_size + # times smaller): `> 0` strict -- -1 = tombstoned, 0 = padding sink. + swa_v2p_pages = self.swa_attn_allocator.virtual_to_physical[v_pages] + live_pages = v_pages[swa_v2p_pages > 0] + if live_pages.numel() > 0: + self.swa_attn_allocator.free(live_pages * self.page_size, _pages=live_pages) + self.full_attn_allocator.free(reps, _pages=v_pages) + self.full_attn_allocator.clear_inverse_history() + self.swa_attn_allocator.clear_inverse_history() - def ensure_capacity(self, full_tokens: int, swa_tokens: int) -> bool: - """Gate one allocation and compact both sides on shortfall.""" - if full_tokens < 0 or swa_tokens < 0: - return False - if full_tokens == 0 and swa_tokens == 0: - return True - page_size = self.page_size - num_full_pages = (int(full_tokens) + page_size - 1) // page_size - num_swa_pages = (int(swa_tokens) + page_size - 1) // page_size - if self._fits_page_demand( - num_full_pages, - num_swa_pages, - compacted=False, - ): - return True - if not self.lazy_compaction or not self._compaction_allowed(): - return False - if not self._fits_page_demand( - num_full_pages, - num_swa_pages, - compacted=True, + def _joint_capacity_memo_violations(self) -> List[str]: + """Idle-time twin of `MultiEndedAllocator._capacity_memo_violations` + for the composite joint view. Empty == healthy.""" + if ( + self._joint_avail_memo_epoch + != self.full_attn_allocator._chain_capacity_epoch() ): - return False - self.full_attn_allocator.flush_for_allocation() - self.swa_attn_allocator.flush_for_allocation() - return self._fits_page_demand( - num_full_pages, - num_swa_pages, - compacted=False, - ) + return [] + actual = self._compute_available_size() + if self._joint_avail_memo_tokens == actual: + return [] + return [ + f"[joint] stale available_size memo: " + f"cached={self._joint_avail_memo_tokens}, actual={actual}" + ] - def create_prefill_budget(self, tree_cache, *, num_mixed_decode_tokens=0): - from sglang.srt.mem_cache.prefill_budget import SharedSWAPrefillBudget + def clear(self) -> None: + self.full_attn_allocator.clear() + self.swa_attn_allocator.clear() + self.free_group = None + self.free_page_reps_group = None + self.full_free_group = [] - return SharedSWAPrefillBudget( - self, tree_cache, num_mixed_decode_tokens=num_mixed_decode_tokens - ) + # -- Lazy compaction hooks -- - def swa_capacity_and_available(self, *, full_capacity, swa_capacity): - return ( - (self.current_full_capacity, self.full_available_size()), - (self.current_swa_capacity, self.swa_available_size()), - ) + def set_latest_forward_done_event(self, event: Optional[torch.cuda.Event]) -> None: + """Forward the per-batch `forward_done` event to BOTH sub-allocators.""" + with record_function("UnifiedSWAAlloc.set_latest_forward_done_event"): + self.full_attn_allocator.set_latest_forward_done_event(event) + self.swa_attn_allocator.set_latest_forward_done_event(event) - def max_new_tokens_for_memory( + def set_inflight_forward( self, - input_tokens: int, - max_new_tokens: int, - *, - token_capacity: int, - sliding_window_size: int | None, - chunk_size: int | None, - ) -> int | None: - from sglang.srt.mem_cache.prefill_budget import estimate_swa_kv_tokens - - def fits(candidate): - return self.can_reserve( - input_tokens + candidate + self.page_size, - estimate_swa_kv_tokens( - input_tokens, - candidate, - sliding_window_size=sliding_window_size, - page_size=self.page_size, - allocation_limit=chunk_size, - ), - empty_pool=True, - require_token_slack=True, + forward_done: torch.cuda.Event, + out_cache_loc_virtual: Optional[torch.Tensor], + ) -> None: + """Hand the forward's metadata to BOTH sub-pools; each materializes its own + write-set via its OWN v2p, and the forward writes both sides per token.""" + with record_function("UnifiedSWAAlloc.set_inflight_forward"): + self.full_attn_allocator.set_inflight_forward( + forward_done, out_cache_loc_virtual + ) + self.swa_attn_allocator.set_inflight_forward( + forward_done, out_cache_loc_virtual ) - if not fits(0): - return None - if fits(max_new_tokens): - return max_new_tokens - lo, hi = 0, max_new_tokens - while lo < hi: - mid = (lo + hi + 1) // 2 - if fits(mid): - lo = mid - else: - hi = mid - 1 - return lo + @abstractmethod + def _build_swa_attn_allocator(self, **kwargs) -> MultiEndedAllocator: ... - def evict_to_free_tokens(self, tree_cache, num_tokens: int) -> None: - from sglang.srt.mem_cache.base_prefix_cache import EvictParams + @abstractmethod + def _wire_peers(self) -> None: ... - if tree_cache is None or tree_cache.is_chunk_cache(): - return - reclaim_plan = self.reclaim_plan( - num_tokens, - num_tokens, - full_evictable_tokens=tree_cache.full_evictable_size(), - swa_evictable_tokens=tree_cache.swa_evictable_size(), - ) - if reclaim_plan is None: - return - full_reclaim, swa_reclaim = reclaim_plan - if full_reclaim or swa_reclaim: - tree_cache.evict_for_alloc( - EvictParams(num_tokens=full_reclaim, swa_num_tokens=swa_reclaim) - ) - # A zero-reclaim plan can still depend on compaction before allocation. - self.ensure_capacity(num_tokens, num_tokens) + @abstractmethod + def _compute_available_size(self) -> int: ... - @property - def draft_virtual_id_space(self) -> int: - return self.full_attn_allocator.max_slots - 1 + @abstractmethod + def ensure_capacity(self, full_tokens: int, swa_tokens: int) -> bool: ... - def debug_print(self) -> str: - return ( - f"#full-available={self.full_attn_allocator.available_size()}, " - f"#swa-available={self.swa_attn_allocator.available_size()}, " - f"#joint-available={self.available_size()}" + @abstractmethod + def _flush_targets(self): ... + + @abstractmethod + def _ask_float_for_room(self, need_tokens: int) -> None: ... + + +class UnifiedSWATokenToKVPoolAllocator(UnifiedSWAAllocatorBase): + """Two-ended FULL/SWA allocator with asymmetric shared-byte reservations.""" + + def _build_swa_attn_allocator(self, **kwargs) -> MultiEndedAllocator: + """The swa sub-allocator: an END pool in the 2-pool pair.""" + return MultiEndedAllocator( + sub_pool_name="swa", + is_id_owner=False, # non-owner; consumes virtuals minted by full + **kwargs, ) - def get_kvcache(self): - return self._kvcache + def _wire_peers(self) -> None: + self.full_attn_allocator.bind_peer(self.swa_attn_allocator) + self.swa_attn_allocator.bind_peer(self.full_attn_allocator) - def translate_kv_loc( - self, - loc: torch.Tensor, - *, - out: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - """Full-layer read path: virtual TOKEN ids -> full-physical TOKEN ids. - ``out=`` writes in place, for cuda-graph buffer stability.""" - result = self.full_attn_allocator.translate_kv_loc(loc, out=out) - return result + def _compute_available_size(self) -> int: + """Joint byte budget in TOKENS: each composite alloc(1) consumes one + full-side AND one swa-side page under the same virtual id.""" + fa, sa = self.full_attn_allocator, self.swa_attn_allocator + e_f = fa.entry_bytes_per_page + e_s = sa.entry_bytes_per_page + # Direction-agnostic shared gap: the free byte band between the two pools. + if fa.grow_direction == "up": + gap_bytes = max(0, sa._byte_low_frontier() - fa._byte_high_frontier()) + else: + gap_bytes = max(0, fa._byte_low_frontier() - sa._byte_high_frontier()) + R_f = fa.num_pages - fa.min_page_index - fa._allocated_pages() + R_s = sa.num_pages - sa.min_page_index - sa._allocated_pages() - def translate_loc_from_full_to_swa( - self, - kv_indices: torch.Tensor, - *, - out: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - """SWA-layer read path: virtual TOKEN ids -> swa kernel-facing ids.""" - return self.swa_attn_allocator.translate_kv_loc_for_kernel(kv_indices, out=out) + if not self.lazy_compaction: + pages_by_bytes = gap_bytes // (e_f + e_s) + return min(pages_by_bytes, R_f, R_s) * self.page_size - @property - def kernel_page_multiplier(self) -> int: - return self.full_attn_allocator.kernel_page_multiplier + H_f = len(fa._free_phys_pages) + H_s = len(sa._free_phys_pages) - @property - def full_v2p_page_table(self) -> torch.Tensor: - """Page-level virtual->physical table of the full sub-pool.""" - return self.full_attn_allocator.virtual_to_physical + K1 = min(H_f, H_s) # Phase 1: both drain - @property - def full_p2v_page_table(self) -> torch.Tensor: - """Page-level physical->virtual table of the full sub-pool.""" - return self.full_attn_allocator.physical_to_virtual + # Phase 2: fewer-holes side extends; more-holes side keeps draining. + if H_f <= H_s: + e_phase2 = e_f + K_phase2_max = H_s + else: + e_phase2 = e_s + K_phase2_max = H_f + K2_room = K_phase2_max - K1 + K2 = min(K2_room, gap_bytes // e_phase2) if e_phase2 > 0 else K2_room + gap_bytes -= K2 * e_phase2 - def translate_kv_loc_for_kernel( - self, - loc: torch.Tensor, - *, - out: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - """Full-pool virtual TOKEN ids -> kernel-facing ids.""" - return self.full_attn_allocator.translate_kv_loc_for_kernel(loc, out=out) + K3 = gap_bytes // (e_f + e_s) # Phase 3: both extend - def translate_write_loc_for_kernel( + K_total = K1 + K2 + K3 + K_total = min(K_total, H_f + R_f, H_s + R_s) # index-space caps + return K_total * self.page_size + + def _flush_targets(self): + """Flush ALL members, including ones that are not short themselves: a + one-sided hole is unusable, and compacting it yields SHARED gap.""" + return (self.full_attn_allocator, self.swa_attn_allocator) + + def _ask_float_for_room(self, need_tokens: int) -> None: + """No float in a two-END chain -- nothing can slide.""" + return None + + def reclaim_plan( self, - loc: torch.Tensor, + full_tokens: int | float, + swa_tokens: int | float, *, - out: Optional[torch.Tensor] = None, - out_width: Optional[int] = None, - ) -> torch.Tensor: - """Widened virtual WRITE loc -> kernel-facing id. DCP is rejected for this - composite at argument validation, so it coincides with the read translate.""" - return self.full_attn_allocator.translate_write_loc_for_kernel( - loc, out=out, out_width=out_width - ) - - @property - def swa_kernel_page_multiplier(self) -> int: - return self.swa_attn_allocator.kernel_page_multiplier + full_evictable_tokens: int = 0, + swa_evictable_tokens: int = 0, + empty_pool: bool = False, + ) -> Optional[Tuple[int, int]]: + """Return cumulative FULL/SWA eviction targets, or None if impossible. - @property - def swa_v2p_page_table(self) -> torch.Tensor: - """Page-level virtual->physical table of the SWA sub-pool.""" - return self.swa_attn_allocator.virtual_to_physical + The tree evicts FULL before SWA. Minimize required SWA reclaim with all + evictable FULL available, then trim excess FULL reclaim. FULL eviction's + actual SWA cascade is counted by the tree's shared eviction tracker. + """ + if full_tokens < 0 or swa_tokens < 0: + return None - # -- alloc -- + page_size = self.page_size + full_pages = (math.ceil(full_tokens) + page_size - 1) // page_size + swa_pages = (math.ceil(swa_tokens) + page_size - 1) // page_size + # Restoring host SWA for device-resident FULL can require more new + # SWA pages than new FULL pages; only the shared budget constrains it. - def alloc(self, need_size: int) -> Optional[torch.Tensor]: - with record_function("UnifiedSWAAlloc.alloc"): - if not self.ensure_capacity(need_size, need_size): - return None - # Snapshot the virtual PAGES full will consume, to bind them on swa too. - num_pages = need_size // self.page_size - fa = self.full_attn_allocator - new_virtual_pages = fa.free_virtual_ids[:num_pages].clone() + compacted = empty_pool or not self.lazy_compaction or self._compaction_allowed() - v_tokens = fa.alloc(need_size) - # Post-pre-check failure can only be internal-state inconsistency. - assert v_tokens is not None, ( - "UnifiedSWA.alloc: full.alloc returned None after joint " - "pre-check passed — internal-state inconsistency" + def fits(full_reclaim_pages: int, swa_reclaim_pages: int) -> bool: + return self._fits_page_demand( + full_pages, + swa_pages, + full_reclaim_pages=full_reclaim_pages, + swa_reclaim_pages=swa_reclaim_pages, + compacted=compacted, + empty_pool=empty_pool, ) - self.swa_attn_allocator.alloc_with_virtual(new_virtual_pages) - return v_tokens - def alloc_extend( - self, - prefix_lens: torch.Tensor, - prefix_lens_cpu: torch.Tensor, - seq_lens: torch.Tensor, - seq_lens_cpu: torch.Tensor, - last_loc: torch.Tensor, - extend_num_tokens: int, - ) -> Optional[torch.Tensor]: - """Paged extend; returns virtual TOKEN ids. The same virtual page maps to - full- and swa-physical, so swa binds exactly what the full kernel consumed.""" - with record_function("UnifiedSWAAlloc.alloc_extend"): - num_new_pages = get_num_new_pages( - seq_lens=seq_lens_cpu, - page_size=self.page_size, - prefix_lens=prefix_lens_cpu, - ) - need_tokens = num_new_pages * self.page_size - if not self.ensure_capacity(need_tokens, need_tokens): - return None + if empty_pool: + return (0, 0) if fits(0, 0) else None + if fits(0, 0): + return (0, 0) - # Snapshot the virtual PAGES the kernel will consume; clone so swa keeps - # its view after the slice is consumed. - fa = self.full_attn_allocator - new_virtual_pages = fa.free_virtual_ids[:num_new_pages].clone() + max_full_pages = min( + self.full_attn_allocator.allocated_count() // page_size, + max(0, int(full_evictable_tokens)) // page_size, + ) + max_swa_pages = min( + self.swa_attn_allocator.allocated_count() // page_size, + max(0, int(swa_evictable_tokens)) // page_size, + ) + if not fits(max_full_pages, max_swa_pages): + return None - out_indices = fa.alloc_extend( - prefix_lens, - prefix_lens_cpu, - seq_lens, - seq_lens_cpu, - last_loc, - extend_num_tokens, - num_new_pages=num_new_pages, - ) - assert out_indices is not None, ( - "UnifiedSWA.alloc_extend: full.alloc_extend returned None " - "after joint pre-check passed — internal-state inconsistency" - ) - self.swa_attn_allocator.alloc_with_virtual(new_virtual_pages) - return out_indices # virtual TOKEN ids + def first_fit(high: int, predicate: Callable[[int], bool]) -> int: + low = 0 + while low < high: + mid = (low + high) // 2 + if predicate(mid): + high = mid + else: + low = mid + 1 + return low - def alloc_decode( + swa_reclaim_pages = first_fit( + max_swa_pages, lambda value: fits(max_full_pages, value) + ) + full_reclaim_pages = first_fit( + max_full_pages, lambda value: fits(value, swa_reclaim_pages) + ) + return ( + full_reclaim_pages * page_size, + swa_reclaim_pages * page_size, + ) + + def can_reserve( self, - seq_lens: torch.Tensor, - seq_lens_cpu: torch.Tensor, - last_loc: torch.Tensor, - ) -> Optional[torch.Tensor]: - """Paged decode: one new token per request, consuming a page only when the - decode wraps.""" - with record_function("UnifiedSWAAlloc.alloc_decode"): - num_new_pages = get_num_new_pages( - seq_lens=seq_lens_cpu, page_size=self.page_size, decode=True - ) - need_tokens = num_new_pages * self.page_size - if not self.ensure_capacity(need_tokens, need_tokens): - return None + full_tokens: int | float, + swa_tokens: int | float, + *, + full_evictable_tokens: int = 0, + swa_evictable_tokens: int = 0, + empty_pool: bool = False, + require_token_slack: bool = False, + ) -> bool: + """Check pending FULL/SWA demand against the shared byte envelope. - fa = self.full_attn_allocator - new_virtual_pages = fa.free_virtual_ids[:num_new_pages].clone() + Scheduler admission keeps the historical one-token strict slack at an + empty-pool boundary. Live admission checks the state reachable after + reclaiming the currently evictable pages from both sides. + """ + if full_tokens < 0 or swa_tokens < 0: + return False + if require_token_slack and ( + full_tokens >= self.size_full or swa_tokens >= self.size_swa + ): + return False - out_indices = fa.alloc_decode(seq_lens, seq_lens_cpu, last_loc) - assert out_indices is not None, ( - "UnifiedSWA.alloc_decode: full.alloc_decode returned None " - "after joint pre-check passed — internal-state inconsistency" + return ( + self.reclaim_plan( + full_tokens, + swa_tokens, + full_evictable_tokens=full_evictable_tokens, + swa_evictable_tokens=swa_evictable_tokens, + empty_pool=empty_pool, ) + is not None + ) - if new_virtual_pages.numel() > 0: - self.swa_attn_allocator.alloc_with_virtual(new_virtual_pages) - - return out_indices # virtual TOKEN ids - - def is_slot_allocated(self, slot: int) -> bool: - """Token-slot surface = the full side (which owns the virtual ids).""" - return self.full_attn_allocator.is_slot_allocated(slot) + def _compaction_allowed(self) -> bool: + return all( + allocator.disagg_move_gate is None or allocator.disagg_move_gate() + for allocator in (self.full_attn_allocator, self.swa_attn_allocator) + ) - def allocator_state_str(self) -> str: - return self.full_attn_allocator.allocator_state_str() + def _fits_page_demand( + self, + num_full_pages: int, + num_swa_pages: int, + *, + full_reclaim_pages: int = 0, + swa_reclaim_pages: int = 0, + compacted: bool, + empty_pool: bool = False, + ) -> bool: + """Check one FULL/SWA page demand against a single allocator snapshot.""" + if min(num_full_pages, num_swa_pages) < 0: + return False - # -- free -- + fa, sa = self.full_attn_allocator, self.swa_attn_allocator + if empty_pool: + full_live_pages = swa_live_pages = 0 + full_reclaim_pages = swa_reclaim_pages = 0 + else: + full_live_pages = fa.allocated_count() // self.page_size + swa_live_pages = sa.allocated_count() // self.page_size + full_reclaim_pages = min(full_live_pages, max(0, int(full_reclaim_pages))) + swa_reclaim_pages = min(swa_live_pages, max(0, int(swa_reclaim_pages))) - def free(self, free_index: torch.Tensor) -> None: - with record_function("UnifiedSWAAlloc.free"): - if free_index is None or free_index.numel() == 0: - return - if self.free_group is not None: - self.free_group.append(self._copy_for_free_group(free_index)) - return - # Order is not load-bearing: the per-sub-pool v2p IS the mapping. Only - # swa needs the tombstone filter; full owns the ids, so all are bound. - v = free_index.detach().to(torch.int64) - v_pages = v // self.page_size - swa_v2p_pages = self.swa_attn_allocator.virtual_to_physical[v_pages] - # `> 0` strict: -1 = tombstoned, 0 = padding-sink page; both skipped. - live_token_mask = swa_v2p_pages > 0 - live_tokens = v[live_token_mask] - if live_tokens.numel() > 0: - self.swa_attn_allocator.free(live_tokens) - self.full_attn_allocator.free(v) - self.full_attn_allocator.clear_inverse_history() - self.swa_attn_allocator.clear_inverse_history() + full_live_pages -= full_reclaim_pages + swa_live_pages -= swa_reclaim_pages + full_total_pages = full_live_pages + num_full_pages + swa_total_pages = swa_live_pages + num_swa_pages + virtual_page_capacity = fa.num_virtual_ids - fa.min_page_index + full_page_capacity = min( + virtual_page_capacity, + fa.num_pages - fa.min_page_index, + ) + swa_page_capacity = min( + virtual_page_capacity, + sa.num_pages - sa.min_page_index, + ) + if full_total_pages > full_page_capacity or swa_total_pages > swa_page_capacity: + return False - def free_swa(self, free_index: torch.Tensor) -> None: - """SWA tombstone path: release swa-physical, keep the virtual id and - full-physical live; `swa.v2p_page[v_page] = -1` IS the tombstone.""" - if free_index is None or free_index.numel() == 0: - return - v = free_index.detach().to(torch.int64) - ps = self.page_size - v_pages = v // ps - # `> 0` strict: -1 = tombstoned, page 0 = padding sink (never freeable). - swa_v2p_pages = self.swa_attn_allocator.virtual_to_physical[v_pages] - live = v[swa_v2p_pages > 0] - if live.numel() == 0: - return - if ps == 1: - # token == page and the live filter just deduped against the v2p - # table, so these ARE unique page ids -- same skip as `_free_lazy`. - self.swa_attn_allocator.free(live, _pages=live) + if compacted: + full_virtual_room = virtual_page_capacity - full_live_pages + full_holes = swa_holes = 0 + full_index_room = full_page_capacity - full_live_pages + swa_index_room = swa_page_capacity - swa_live_pages + gap_bytes = ( + self._empty_shared_gap_bytes + - full_live_pages * fa.entry_bytes_per_page + - swa_live_pages * sa.entry_bytes_per_page + ) else: - self.swa_attn_allocator.free(live) - self.swa_attn_allocator.clear_inverse_history() + full_virtual_room = len(fa.free_virtual_ids) + full_reclaim_pages + full_holes = len(fa._free_phys_pages) + full_reclaim_pages + swa_holes = len(sa._free_phys_pages) + swa_reclaim_pages + full_index_room = fa.num_pages - fa.min_page_index - fa._allocated_pages() + swa_index_room = sa.num_pages - sa.min_page_index - sa._allocated_pages() + gap_bytes = fa._current_gap_bytes() - def free_swa_segment(self, free_index: torch.Tensor, *, start_pos: int) -> None: - """free_swa() for a kv-row segment: `start_pos` promises a contiguous - ascending range, so page reps come from stride arithmetic, not `torch.unique`.""" - if free_index is None or free_index.numel() == 0: - return - if self.page_size == 1: - self.free_swa(free_index) - return - ps = self.page_size - reps = self.swa_attn_allocator._page_reps( - free_index.detach().to(torch.int64), start_pos + if num_full_pages > full_virtual_room: + return False + if num_full_pages > full_holes + full_index_room: + return False + if num_swa_pages > swa_holes + swa_index_room: + return False + full_extensions = max(0, num_full_pages - full_holes) + swa_extensions = max(0, num_swa_pages - swa_holes) + return ( + full_extensions * fa.entry_bytes_per_page + + swa_extensions * sa.entry_bytes_per_page + <= max(0, gap_bytes) ) - # Keep only pages still bound on swa; freeing a tombstoned one would - # corrupt the hole list. `> 0` strict: -1 tombstoned, 0 padding sink. - rep_pages = reps // ps - swa_v2p_pages = self.swa_attn_allocator.virtual_to_physical[rep_pages] - live_reps = reps[swa_v2p_pages > 0] - if live_reps.numel() == 0: - return - self.swa_attn_allocator.free(live_reps, _pages=live_reps // ps) - self.swa_attn_allocator.clear_inverse_history() - def free_full(self, free_index: torch.Tensor) -> None: - """Release the full-physical page and the virtual id, leaving the swa - side alone -- the caller already tombstoned it (`swa.v2p_page == -1`).""" - if free_index is None or free_index.numel() == 0: - return - if self.free_group is not None: - self.full_free_group.append(self._copy_for_free_group(free_index)) - return - self.full_attn_allocator.free(free_index.detach().to(torch.int64)) - self.full_attn_allocator.clear_inverse_history() + def ensure_capacity(self, full_tokens: int, swa_tokens: int) -> bool: + """Gate one allocation and compact both sides on shortfall.""" + if full_tokens < 0 or swa_tokens < 0: + return False + if full_tokens == 0 and swa_tokens == 0: + return True + page_size = self.page_size + num_full_pages = (int(full_tokens) + page_size - 1) // page_size + num_swa_pages = (int(swa_tokens) + page_size - 1) // page_size + if self._fits_page_demand( + num_full_pages, + num_swa_pages, + compacted=False, + ): + return True + if not self.lazy_compaction or not self._compaction_allowed(): + return False + if not self._fits_page_demand( + num_full_pages, + num_swa_pages, + compacted=True, + ): + return False + self.full_attn_allocator.flush_for_allocation() + self.swa_attn_allocator.flush_for_allocation() + return self._fits_page_demand( + num_full_pages, + num_swa_pages, + compacted=False, + ) - def free_full_segment(self, free_index: torch.Tensor, *, start_pos: int) -> None: - if free_index is None or free_index.numel() == 0: - return - if self.page_size == 1: - # token == page: free_full already frees by exact ids, no dedup. - self.free_full(free_index) - return - # The swa v2p is the mapping, so a tombstoned swa page drops out of the - # two-sided segment path by itself; full-only is the same call. - self.free_segment(free_index, start_pos=start_pos) + def create_prefill_budget(self, tree_cache, *, num_mixed_decode_tokens=0): + from sglang.srt.mem_cache.prefill_budget import SharedSWAPrefillBudget - def set_full_to_swa_mapping( - self, full_indices: torch.Tensor, swa_indices: torch.Tensor - ) -> None: - """No-op stub for HiCache load-back: in shared mode the swa v2p IS the - mapping, and HiCache for shared SWA is out of scope.""" - return + return SharedSWAPrefillBudget( + self, tree_cache, num_mixed_decode_tokens=num_mixed_decode_tokens + ) - def clear_full_to_swa_mapping(self, full_indices: torch.Tensor) -> None: - # Paired with set_full_to_swa_mapping: shared mode has no mapping tensor. - return + def swa_capacity_and_available(self, *, full_capacity, swa_capacity): + return ( + (self.current_full_capacity, self.full_available_size()), + (self.current_swa_capacity, self.swa_available_size()), + ) - # -- free-group -- + def max_new_tokens_for_memory( + self, + input_tokens: int, + max_new_tokens: int, + *, + token_capacity: int, + sliding_window_size: int | None, + chunk_size: int | None, + ) -> int | None: + from sglang.srt.mem_cache.prefill_budget import estimate_swa_kv_tokens - # Not the SWA parent's hooks: those open the parent's paged full allocator - # as a free group, and this composite's sub-pools defer on their own. - def free_group_begin(self) -> None: - BaseTokenToKVPoolAllocator.free_group_begin(self) - self.free_page_reps_group = [] - self.full_free_group = [] + def fits(candidate): + return self.can_reserve( + input_tokens + candidate + self.page_size, + estimate_swa_kv_tokens( + input_tokens, + candidate, + sliding_window_size=sliding_window_size, + page_size=self.page_size, + allocation_limit=chunk_size, + ), + empty_pool=True, + require_token_slack=True, + ) - def free_group_end(self) -> None: - pending, self.free_page_reps_group = self.free_page_reps_group, None - full_free_group, self.full_free_group = self.full_free_group, [] - BaseTokenToKVPoolAllocator.free_group_end(self) - if full_free_group: - self.full_attn_allocator.free(torch.cat(full_free_group)) - self.full_attn_allocator.clear_inverse_history() - if pending: - self._release_page_reps(pending) + if not fits(0): + return None + if fits(max_new_tokens): + return max_new_tokens + lo, hi = 0, max_new_tokens + while lo < hi: + mid = (lo + hi + 1) // 2 + if fits(mid): + lo = mid + else: + hi = mid - 1 + return lo - def free_segment(self, free_index: torch.Tensor, *, start_pos: int) -> None: - """Fixed-shape counterpart of `free()`; see `MultiEndedAllocator._page_reps`. - Both sides share one page-rep derivation instead of dedup'ing twice.""" - if free_index is None or free_index.numel() == 0: - return - if self.page_size == 1: - self.free(free_index) + def evict_to_free_tokens(self, tree_cache, num_tokens: int) -> bool | None: + from sglang.srt.mem_cache.base_prefix_cache import EvictParams + + if tree_cache is None or tree_cache.is_chunk_cache(): return - reps = self.full_attn_allocator._page_reps( - free_index.detach().to(torch.int64), start_pos + reclaim_plan = self.reclaim_plan( + num_tokens, + num_tokens, + full_evictable_tokens=tree_cache.full_evictable_size(), + swa_evictable_tokens=tree_cache.swa_evictable_size(), ) - if self.free_page_reps_group is None: - self._release_page_reps((reps,)) - else: - self.free_page_reps_group.append(reps) - - def _release_page_reps(self, pieces: Sequence[torch.Tensor]) -> None: - reps = pieces[0] if len(pieces) == 1 else torch.cat(tuple(pieces)) - v_pages = reps // self.page_size - # Same tombstone filter as `free`, but at PAGE granularity (page_size - # times smaller): `> 0` strict -- -1 = tombstoned, 0 = padding sink. - swa_v2p_pages = self.swa_attn_allocator.virtual_to_physical[v_pages] - live_pages = v_pages[swa_v2p_pages > 0] - if live_pages.numel() > 0: - self.swa_attn_allocator.free(live_pages * self.page_size, _pages=live_pages) - self.full_attn_allocator.free(reps, _pages=v_pages) - self.full_attn_allocator.clear_inverse_history() - self.swa_attn_allocator.clear_inverse_history() + if reclaim_plan is None: + return + full_reclaim, swa_reclaim = reclaim_plan + if full_reclaim or swa_reclaim: + tree_cache.evict_for_alloc( + EvictParams(num_tokens=full_reclaim, swa_num_tokens=swa_reclaim) + ) + # A zero-reclaim plan can still depend on compaction before allocation. + return self.ensure_capacity(num_tokens, num_tokens) def verify_byte_accounting(self) -> List[str]: return ( @@ -922,52 +990,6 @@ def verify_byte_accounting(self) -> List[str]: + self._joint_capacity_memo_violations() ) - def _joint_capacity_memo_violations(self) -> List[str]: - """Idle-time twin of `MultiEndedAllocator._capacity_memo_violations` - for the composite joint view. Empty == healthy.""" - if ( - self._joint_avail_memo_epoch - != self.full_attn_allocator._chain_capacity_epoch() - ): - return [] - actual = self._compute_available_size() - if self._joint_avail_memo_tokens == actual: - return [] - return [ - f"[joint] stale available_size memo: " - f"cached={self._joint_avail_memo_tokens}, actual={actual}" - ] - - def clear(self) -> None: - self.full_attn_allocator.clear() - self.swa_attn_allocator.clear() - self.free_group = None - self.free_page_reps_group = None - self.full_free_group = [] - - # -- Lazy compaction hooks -- - - def set_latest_forward_done_event(self, event: Optional[torch.cuda.Event]) -> None: - """Forward the per-batch `forward_done` event to BOTH sub-allocators.""" - with record_function("UnifiedSWAAlloc.set_latest_forward_done_event"): - self.full_attn_allocator.set_latest_forward_done_event(event) - self.swa_attn_allocator.set_latest_forward_done_event(event) - - def set_inflight_forward( - self, - forward_done: torch.cuda.Event, - out_cache_loc_virtual: Optional[torch.Tensor], - ) -> None: - """Hand the forward's metadata to BOTH sub-pools; each materializes its own - write-set via its OWN v2p, and the forward writes both sides per token.""" - with record_function("UnifiedSWAAlloc.set_inflight_forward"): - self.full_attn_allocator.set_inflight_forward( - forward_done, out_cache_loc_virtual - ) - self.swa_attn_allocator.set_inflight_forward( - forward_done, out_cache_loc_virtual - ) - def flush_opportunistic(self) -> int: """Non-urgent flush of BOTH sub-allocators; sync-free.""" with record_function("UnifiedSWAAlloc.flush_opportunistic"): @@ -983,7 +1005,7 @@ def flush_opportunistic(self) -> int: return fa.flush_opportunistic() + sa.flush_opportunistic() -class UnifiedMambaSWATokenToKVPoolAllocator(UnifiedSWATokenToKVPoolAllocator): +class UnifiedMambaSWATokenToKVPoolAllocator(UnifiedSWAAllocatorBase): """Tri-pool composite for models with full KV + SWA KV + mamba/conv state (both `mambaish_config` and `is_hybrid_swa`). diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index b3e73e6d9722..c080231385d3 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/python/sglang/srt/mem_cache/common.py @@ -148,9 +148,11 @@ def maybe_cache_unfinished_req(req: Req, tree_cache: BasePrefixCache, **kwargs): tree_cache.cache_unfinished_req(req, **kwargs) -def evict_from_tree_cache(tree_cache: BasePrefixCache | None, num_tokens: int): +def evict_from_tree_cache( + tree_cache: BasePrefixCache | None, num_tokens: int +) -> bool | None: if tree_cache is not None and not tree_cache.is_chunk_cache(): - tree_cache.token_to_kv_pool_allocator.evict_to_free_tokens( + return tree_cache.token_to_kv_pool_allocator.evict_to_free_tokens( tree_cache, num_tokens ) diff --git a/python/sglang/srt/mem_cache/kv_cache_configurator.py b/python/sglang/srt/mem_cache/kv_cache_configurator.py index 2ddb58efb21f..a199a7d6a390 100644 --- a/python/sglang/srt/mem_cache/kv_cache_configurator.py +++ b/python/sglang/srt/mem_cache/kv_cache_configurator.py @@ -51,7 +51,7 @@ is_swa_req_ring, ) from sglang.srt.mem_cache.allocator.unified_hybrid_swa import ( - UnifiedSWATokenToKVPoolAllocator, + UnifiedSWAAllocatorBase, ) from sglang.srt.mem_cache.allocator.unified_mamba import ( UnifiedMambaTokenToKVPoolAllocator, @@ -544,7 +544,7 @@ def _init_pools( token_to_kv_pool_allocator, ( UnifiedMambaTokenToKVPoolAllocator, - UnifiedSWATokenToKVPoolAllocator, + UnifiedSWAAllocatorBase, ), ): draft_virtual_id_space = ( @@ -567,7 +567,7 @@ def _init_pools( if ( isinstance( token_to_kv_pool_allocator, - UnifiedSWATokenToKVPoolAllocator, + UnifiedSWAAllocatorBase, ) and self.is_hybrid_swa ): @@ -2153,7 +2153,7 @@ def _build_token_to_kv_pool_allocator( else: swa_allocator = token_to_kv_pool_allocator uses_unified_virtual_ids = isinstance( - swa_allocator, UnifiedSWATokenToKVPoolAllocator + swa_allocator, UnifiedSWAAllocatorBase ) has_draft_swa_layers = ( not self.is_hybrid_swa_mtp_draft or self.draft_swa_full_capacity diff --git a/python/sglang/srt/mem_cache/kv_index_translator.py b/python/sglang/srt/mem_cache/kv_index_translator.py index 1ccb5b35077c..b7c71ae096e4 100644 --- a/python/sglang/srt/mem_cache/kv_index_translator.py +++ b/python/sglang/srt/mem_cache/kv_index_translator.py @@ -65,7 +65,7 @@ build_kv_read_table_packed, ) from sglang.srt.mem_cache.allocator.unified_hybrid_swa import ( - UnifiedSWATokenToKVPoolAllocator, + UnifiedSWAAllocatorBase, ) from sglang.srt.mem_cache.allocator.unified_mamba import ( UnifiedMambaTokenToKVPoolAllocator, @@ -121,7 +121,7 @@ def __init__( self.is_translating = ( isinstance( token_to_kv_pool_allocator, - (UnifiedMambaTokenToKVPoolAllocator, UnifiedSWATokenToKVPoolAllocator), + (UnifiedMambaTokenToKVPoolAllocator, UnifiedSWAAllocatorBase), ) and token_to_kv_pool_allocator.get_kvcache() is token_to_kv_pool ) @@ -139,7 +139,7 @@ def __init__( # DCP read ids stay WIDENED to the consumer: selecting this rank's # share changes the length, so only the production site can do it. self.defer_read_translate = get_parallel().attn_dcp_size > 1 - if isinstance(alloc, UnifiedSWATokenToKVPoolAllocator): + if isinstance(alloc, UnifiedSWAAllocatorBase): self._swa_v2p_table = alloc.swa_v2p_page_table self._swa_page_multiplier = alloc.swa_kernel_page_multiplier self._swa_write_loc_from_full = self._swa_write_loc_unified diff --git a/python/sglang/srt/mem_cache/unified_cache/components/swa_component.py b/python/sglang/srt/mem_cache/unified_cache/components/swa_component.py index 705cdf0f32ce..fb5d3cc2158c 100644 --- a/python/sglang/srt/mem_cache/unified_cache/components/swa_component.py +++ b/python/sglang/srt/mem_cache/unified_cache/components/swa_component.py @@ -211,11 +211,11 @@ def _translate_full_to_swa(self, full_indices: torch.Tensor) -> torch.Tensor: def _unified_allocator(self): """The unified SWA composite, or None when running on the static pool.""" from sglang.srt.mem_cache.allocator.unified_hybrid_swa import ( - UnifiedSWATokenToKVPoolAllocator, + UnifiedSWAAllocatorBase, ) allocator = self.cache.token_to_kv_pool_allocator - if isinstance(allocator, UnifiedSWATokenToKVPoolAllocator): + if isinstance(allocator, UnifiedSWAAllocatorBase): return allocator return None diff --git a/test/registered/unit/mem_cache/test_unified_tri_pool.py b/test/registered/unit/mem_cache/test_unified_tri_pool.py index 755d96848f29..0bf7ecaa644e 100644 --- a/test/registered/unit/mem_cache/test_unified_tri_pool.py +++ b/test/registered/unit/mem_cache/test_unified_tri_pool.py @@ -961,11 +961,19 @@ def _factory_kwargs(self, **over): return kw def test_budget_sizing_and_boot_signature(self): + from sglang.srt.mem_cache.allocator.unified_hybrid_swa import ( + UnifiedSWAAllocatorBase, + UnifiedSWATokenToKVPoolAllocator, + ) + budget = 1 << 20 bundle = init_unified_mamba_swa_pools( **self._factory_kwargs(unified_total_bytes=budget) ) pool = bundle.unified_memory_pool + allocator = bundle.token_to_kv_pool_allocator + self.assertIsInstance(allocator, UnifiedSWAAllocatorBase) + self.assertNotIsInstance(allocator, UnifiedSWATokenToKVPoolAllocator) # Buffer = budget + the state pool's bytes (budget captured AFTER the # state carve-out), never the token-count re-sum. state_bytes = 4 * pool.spec("mamba").entry_bytes() From 121fddedc3486a30ba4d2a3d8324a88d52056e3d Mon Sep 17 00:00:00 2001 From: yhzhuang Date: Sun, 13 Sep 2026 18:58:03 -0700 Subject: [PATCH 39/55] Clean up unified PD rejection and transfer interfaces --- python/sglang/srt/arg_groups/overrides.py | 3 + python/sglang/srt/disaggregation/base/conn.py | 24 - python/sglang/srt/disaggregation/decode.py | 81 +- python/sglang/srt/disaggregation/prefill.py | 20 +- python/sglang/srt/mem_cache/allocator/base.py | 7 +- .../mem_cache/allocator/unified_hybrid_swa.py | 1241 +++++++++-------- python/sglang/srt/mem_cache/common.py | 6 +- .../srt/mem_cache/kv_cache_configurator.py | 8 +- .../srt/mem_cache/kv_index_translator.py | 6 +- .../unified_cache/components/swa_component.py | 4 +- ...test_priority_scheduling_disaggregation.py | 1 - .../unit/mem_cache/test_unified_tri_pool.py | 8 + 12 files changed, 691 insertions(+), 718 deletions(-) diff --git a/python/sglang/srt/arg_groups/overrides.py b/python/sglang/srt/arg_groups/overrides.py index b49b490a6f69..0fe079fdaa4e 100644 --- a/python/sglang/srt/arg_groups/overrides.py +++ b/python/sglang/srt/arg_groups/overrides.py @@ -1777,6 +1777,9 @@ def post_capture_kv_sizing_planned(server_args: Any) -> bool: mla_enabled = use_mla_backend(server_args) if not envs.SGLANG_ENABLE_POST_CAPTURE_KV_SIZING.get(): return False + # Unified arenas are fully backed before capture and cannot resize afterward. + if cfg.enable_unified_memory: + return False if cfg.device != "cuda": return False if cfg.dcp_size != 1: diff --git a/python/sglang/srt/disaggregation/base/conn.py b/python/sglang/srt/disaggregation/base/conn.py index ad698fc4cddc..733128b3cfde 100644 --- a/python/sglang/srt/disaggregation/base/conn.py +++ b/python/sglang/srt/disaggregation/base/conn.py @@ -108,7 +108,6 @@ class KVPoll: class BaseKVManager(ABC): """Base class for managing transfer states""" - uses_separate_draft_kv_indices: bool = False enable_deferred_decode_kv_release: bool = False @abstractmethod @@ -157,17 +156,6 @@ def send( """ ... - def send_with_draft_indices( - self, - kv_indices: npt.NDArray[np.int32], - draft_kv_indices: npt.NDArray[np.int32], - state_indices: Optional[List] = None, - num_kv_tokens: Optional[int] = None, - ): - raise NotImplementedError( - f"{type(self).__name__} does not support separate draft KV indices" - ) - def pop_decode_prefix_len(self) -> int: return 0 @@ -238,18 +226,6 @@ def send_metadata( """ ... - def send_metadata_with_draft_indices( - self, - kv_indices: npt.NDArray[np.int32], - draft_kv_indices: npt.NDArray[np.int32], - aux_index: Optional[int] = None, - state_indices: Optional[List] = None, - decode_prefix_len: Optional[int] = None, - ): - raise NotImplementedError( - f"{type(self).__name__} does not support separate draft KV indices" - ) - @abstractmethod def poll(self) -> KVPoll: """ diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 3a8b6b6e2a54..6b6c75e26ef0 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -510,12 +510,14 @@ def _reclaim_swa_tail_capacity( if self._supports_unified_swa_reservation(): full_required = ceil_align(full_len, page_size) swa_required = ceil_align(swa_tail_len, page_size) - allocator.evict_to_free_tokens( + capacity_ready = allocator.evict_to_free_tokens( self.tree_cache, full_required, swa_num_tokens=swa_required, ) - if allocator.ensure_capacity(full_required, swa_required): + if capacity_ready is None: + capacity_ready = allocator.ensure_capacity(full_required, swa_required) + if capacity_ready: return None return ( "Unified FULL/SWA byte reclamation insufficient: " @@ -852,12 +854,7 @@ def _pre_alloc_fill_len(req: Req) -> int: return len(req.origin_input_ids) + max(len(req.output_ids) - 1, 0) def _check_if_req_exceed_kv_capacity(self, req: Req) -> bool: - # HiSparse admits up to the host-backed logical capacity. - if self.scheduler.enable_hisparse: - capacity = self.scheduler.tp_worker.model_runner.max_token_pool_size - else: - capacity = self.max_total_num_tokens - input_len = self._rebootstrap_prefill_len(req) + message = None if self._supports_unified_swa_reservation(): full_required, swa_required = self._prealloc_required_tokens(req) if not self._uses_swa_tail_prealloc(): @@ -869,29 +866,29 @@ def _check_if_req_exceed_kv_capacity(self, req: Req) -> bool: f"Request {req.rid} exceeds the unified FULL/SWA KV byte " f"budget: full={full_required}, swa={swa_required}" ) - logger.error(message) - prepare_abort(req, message, status_code=HTTPStatus.BAD_REQUEST) - self.scheduler.output_streamer.stream_output([req], req.return_logprob) - return True - return False - if input_len > capacity: - message = f"Request {req.rid} exceeds the maximum number of tokens: {input_len} > {capacity}" + else: + # HiSparse admits up to the host-backed logical capacity. + capacity = ( + self.scheduler.tp_worker.model_runner.max_token_pool_size + if self.scheduler.enable_hisparse + else self.max_total_num_tokens + ) + input_len = self._rebootstrap_prefill_len(req) + if input_len > capacity: + message = f"Request {req.rid} exceeds the maximum number of tokens: {input_len} > {capacity}" + elif self._uses_swa_tail_prealloc(): + _, swa_required = self._prealloc_required_tokens(req) + swa_capacity = self.token_to_kv_pool_allocator.size_swa + if swa_required > swa_capacity: + message = ( + f"Request {req.rid} requires too many SWA KV tokens for " + f"decode preallocation: {swa_required} > {swa_capacity}" + ) + if message is not None: logger.error(message) prepare_abort(req, message, status_code=HTTPStatus.BAD_REQUEST) self.scheduler.output_streamer.stream_output([req], req.return_logprob) return True - if self._uses_swa_tail_prealloc(): - _, swa_required = self._prealloc_required_tokens(req) - swa_capacity = self.token_to_kv_pool_allocator.size_swa - if swa_required > swa_capacity: - message = ( - f"Request {req.rid} requires too many SWA KV tokens for " - f"decode preallocation: {swa_required} > {swa_capacity}" - ) - logger.error(message) - prepare_abort(req, message, status_code=HTTPStatus.BAD_REQUEST) - self.scheduler.output_streamer.stream_output([req], req.return_logprob) - return True return False def extend(self, reqs: List[Req], is_retracted: bool = False) -> None: @@ -1603,15 +1600,6 @@ def _c128_state_payload(): page_indices = kv_to_page_indices(kv_indices, kv_transfer_page_size).astype( np.int32 ) - draft_page_indices = None - if self.kv_manager.uses_separate_draft_kv_indices: - if self.scheduler.enable_hisparse: - raise NotImplementedError( - "separate draft KV indices are not supported with HiSparse" - ) - draft_page_indices = kv_to_page_indices( - raw_kv_indices, page_size - ).astype(np.int32) device_page_indices = None if ( self.scheduler.enable_hisparse @@ -1648,21 +1636,12 @@ def _c128_state_payload(): self.transfer_queue.staging_handler.register_decode_req( decode_req.req.bootstrap_room, decode_req ) - if draft_page_indices is not None: - decode_req.kv_receiver.send_metadata_with_draft_indices( - page_indices, - draft_page_indices, - decode_req.metadata_buffer_index, - state_indices, - **metadata_kwargs, - ) - else: - decode_req.kv_receiver.send_metadata( - page_indices, - decode_req.metadata_buffer_index, - state_indices, - **metadata_kwargs, - ) + decode_req.kv_receiver.send_metadata( + page_indices, + decode_req.metadata_buffer_index, + state_indices, + **metadata_kwargs, + ) if decode_req.is_rebootstrap: self.kv_manager.submit_prefill_recompute( decode_req.kv_receiver, diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index cdaf15261fef..fdefaaeb07d6 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -1422,7 +1422,6 @@ def _c128_state_payload(): else: segments = [(start_idx, end_idx)] - separate_draft_indices = self.disagg_prefill_bootstrap_queue.kv_manager.uses_separate_draft_kv_indices for seg_start, seg_end in segments: is_final_segment = seg_end == end_idx raw_kv_indices = self.req_to_token_pool.req_to_token[ @@ -1442,20 +1441,11 @@ def _c128_state_payload(): ): continue send_state_indices = state_indices if segment_is_last else None - if separate_draft_indices: - draft_page_indices = kv_to_page_indices(raw_kv_indices, page_size) - req.disagg_kv_sender.send_with_draft_indices( - page_indices, - draft_page_indices, - send_state_indices, - num_kv_tokens=seg_end - seg_start, - ) - else: - req.disagg_kv_sender.send( - page_indices, - send_state_indices, - num_kv_tokens=seg_end - seg_start, - ) + req.disagg_kv_sender.send( + page_indices, + send_state_indices, + num_kv_tokens=seg_end - seg_start, + ) req.start_send_idx = end_idx # A last chunk needs no entry: every `last_chunk=True` call site has # already put the request on `disagg_prefill_inflight_queue`. diff --git a/python/sglang/srt/mem_cache/allocator/base.py b/python/sglang/srt/mem_cache/allocator/base.py index 5bc0b6567f9f..a91f7edc3f29 100644 --- a/python/sglang/srt/mem_cache/allocator/base.py +++ b/python/sglang/srt/mem_cache/allocator/base.py @@ -95,9 +95,12 @@ def max_new_tokens_for_memory( 0, min(max_new_tokens, token_capacity - paged_input - self.page_size - 1) ) - def evict_to_free_tokens(self, tree_cache, num_tokens: int) -> None: + def evict_to_free_tokens(self, tree_cache, num_tokens: int) -> bool | None: """Evict unlocked prefix-cache entries until this allocator can serve - ``num_tokens`` or nothing evictable remains.""" + ``num_tokens`` or nothing evictable remains. + + Return whether capacity was realized, or None if it still needs checking. + """ from sglang.srt.mem_cache.base_prefix_cache import EvictParams if tree_cache is None or tree_cache.is_chunk_cache(): diff --git a/python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py b/python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py index be837800d757..dd7ea6d8d2d3 100644 --- a/python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py +++ b/python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py @@ -18,6 +18,7 @@ import logging import math +from abc import abstractmethod from typing import Callable, List, Optional, Sequence, Tuple, TypeGuard import torch @@ -46,23 +47,15 @@ def supports_swa_byte_budget( allocator: BaseTokenToKVPoolAllocator | None, ) -> TypeGuard[UnifiedSWATokenToKVPoolAllocator]: """Whether FULL/SWA demand can be checked against a two-pool byte budget.""" - return ( - isinstance(allocator, UnifiedSWATokenToKVPoolAllocator) - and allocator.supports_asymmetric_reservation - ) + return isinstance(allocator, UnifiedSWATokenToKVPoolAllocator) -class UnifiedSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator): - """Composite allocator for the hybrid SWA pair (full + swa MHA sub-pools). +class UnifiedSWAAllocatorBase(SWATokenToKVPoolAllocator): + """Shared FULL/SWA virtual IDs, allocation lifecycle, and index translation. - One alloc(N) binds N pages on BOTH sides under the same virtual id, so - `available_size()` (joint bytes, in TOKENS) is the only safe alloc pre-check. + Concrete allocators define the two-ended or Mamba/SWA/FULL capacity policy. """ - # Unequal FULL/SWA reservations share one byte budget. The tri-pool opts - # out because its Mamba state competes for those same bytes. - supports_asymmetric_reservation = True - # Parent's `size` property has no setter but base init does `self.size = size`; # the no-op setter below absorbs that write. @property @@ -191,18 +184,6 @@ def __init__( # -- construction hooks (the tri-pool subclass overrides both) -- - def _build_swa_attn_allocator(self, **kwargs) -> MultiEndedAllocator: - """The swa sub-allocator: an END pool in the 2-pool pair.""" - return MultiEndedAllocator( - sub_pool_name="swa", - is_id_owner=False, # non-owner; consumes virtuals minted by full - **kwargs, - ) - - def _wire_peers(self) -> None: - self.full_attn_allocator.bind_peer(self.swa_attn_allocator) - self.swa_attn_allocator.bind_peer(self.full_attn_allocator) - # -- capacity reporting (three-way split) -- def available_size(self) -> int: @@ -213,46 +194,6 @@ def available_size(self) -> int: self._joint_avail_memo_epoch = epoch return self._joint_avail_memo_tokens - def _compute_available_size(self) -> int: - """Joint byte budget in TOKENS: each composite alloc(1) consumes one - full-side AND one swa-side page under the same virtual id.""" - fa, sa = self.full_attn_allocator, self.swa_attn_allocator - e_f = fa.entry_bytes_per_page - e_s = sa.entry_bytes_per_page - # Direction-agnostic shared gap: the free byte band between the two pools. - if fa.grow_direction == "up": - gap_bytes = max(0, sa._byte_low_frontier() - fa._byte_high_frontier()) - else: - gap_bytes = max(0, fa._byte_low_frontier() - sa._byte_high_frontier()) - R_f = fa.num_pages - fa.min_page_index - fa._allocated_pages() - R_s = sa.num_pages - sa.min_page_index - sa._allocated_pages() - - if not self.lazy_compaction: - pages_by_bytes = gap_bytes // (e_f + e_s) - return min(pages_by_bytes, R_f, R_s) * self.page_size - - H_f = len(fa._free_phys_pages) - H_s = len(sa._free_phys_pages) - - K1 = min(H_f, H_s) # Phase 1: both drain - - # Phase 2: fewer-holes side extends; more-holes side keeps draining. - if H_f <= H_s: - e_phase2 = e_f - K_phase2_max = H_s - else: - e_phase2 = e_s - K_phase2_max = H_f - K2_room = K_phase2_max - K1 - K2 = min(K2_room, gap_bytes // e_phase2) if e_phase2 > 0 else K2_room - gap_bytes -= K2 * e_phase2 - - K3 = gap_bytes // (e_f + e_s) # Phase 3: both extend - - K_total = K1 + K2 + K3 - K_total = min(K_total, H_f + R_f, H_s + R_s) # index-space caps - return K_total * self.page_size - # Slot-conservation views for the leak invariant only; the byte-coordinated # value would flag spurious leaks. `allocated_count()` is in TOKENS. def _conserve_full_available_size(self) -> int: @@ -295,15 +236,6 @@ def schedulable_full_available_size(self) -> int: def schedulable_swa_available_size(self) -> int: return self.swa_attn_allocator.schedulable_available_size() - def _flush_targets(self): - """Flush ALL members, including ones that are not short themselves: a - one-sided hole is unusable, and compacting it yields SHARED gap.""" - return (self.full_attn_allocator, self.swa_attn_allocator) - - def _ask_float_for_room(self, need_tokens: int) -> None: - """No float in a two-END chain -- nothing can slide.""" - return None - # `size_full` / `size_swa` bound each side independently; current capacities # also account for the peer's live byte usage. @@ -315,144 +247,620 @@ def current_full_capacity(self) -> int: def current_swa_capacity(self) -> int: return self.swa_available_size() + self.swa_attn_allocator.allocated_count() - def reclaim_plan( - self, - full_tokens: int | float, - swa_tokens: int | float, - *, - full_evictable_tokens: int = 0, - swa_evictable_tokens: int = 0, - empty_pool: bool = False, - ) -> Optional[Tuple[int, int]]: - """Return cumulative FULL/SWA eviction targets, or None if impossible. - - The tree evicts FULL before SWA. Minimize required SWA reclaim with all - evictable FULL available, then trim excess FULL reclaim. FULL eviction's - actual SWA cascade is counted by the tree's shared eviction tracker. - """ - if full_tokens < 0 or swa_tokens < 0: - return None - - page_size = self.page_size - full_pages = (math.ceil(full_tokens) + page_size - 1) // page_size - swa_pages = (math.ceil(swa_tokens) + page_size - 1) // page_size - # Restoring host SWA for device-resident FULL can require more new - # SWA pages than new FULL pages; only the shared budget constrains it. + @property + def draft_virtual_id_space(self) -> int: + return self.full_attn_allocator.max_slots - 1 - compacted = empty_pool or not self.lazy_compaction or self._compaction_allowed() + def debug_print(self) -> str: + return ( + f"#full-available={self.full_attn_allocator.available_size()}, " + f"#swa-available={self.swa_attn_allocator.available_size()}, " + f"#joint-available={self.available_size()}" + ) - def fits(full_reclaim_pages: int, swa_reclaim_pages: int) -> bool: - return self._fits_page_demand( - full_pages, - swa_pages, - full_reclaim_pages=full_reclaim_pages, - swa_reclaim_pages=swa_reclaim_pages, - compacted=compacted, - empty_pool=empty_pool, - ) + def get_kvcache(self): + return self._kvcache - if empty_pool: - return (0, 0) if fits(0, 0) else None - if fits(0, 0): - return (0, 0) + def translate_kv_loc( + self, + loc: torch.Tensor, + *, + out: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """Full-layer read path: virtual TOKEN ids -> full-physical TOKEN ids. + ``out=`` writes in place, for cuda-graph buffer stability.""" + result = self.full_attn_allocator.translate_kv_loc(loc, out=out) + return result - max_full_pages = min( - self.full_attn_allocator.allocated_count() // page_size, - max(0, int(full_evictable_tokens)) // page_size, - ) - max_swa_pages = min( - self.swa_attn_allocator.allocated_count() // page_size, - max(0, int(swa_evictable_tokens)) // page_size, - ) - if not fits(max_full_pages, max_swa_pages): - return None + def translate_kv_indices_for_transfer( + self, kv_indices: torch.Tensor + ) -> torch.Tensor: + """Translate virtual token ids to full-pool physical ids for PD.""" + return self.full_attn_allocator.translate_kv_loc(kv_indices.to(torch.int64)) - def first_fit(high: int, predicate: Callable[[int], bool]) -> int: - low = 0 - while low < high: - mid = (low + high) // 2 - if predicate(mid): - high = mid - else: - low = mid + 1 - return low + def translate_swa_indices_for_transfer( + self, kv_indices: torch.Tensor + ) -> torch.Tensor: + """Virtual token ids -> physical, not kernel-facing, SWA ids for PD.""" + return self.swa_attn_allocator.translate_kv_loc(kv_indices.to(torch.int64)) - swa_reclaim_pages = first_fit( - max_swa_pages, lambda value: fits(max_full_pages, value) - ) - full_reclaim_pages = first_fit( - max_full_pages, lambda value: fits(value, swa_reclaim_pages) - ) - return ( - full_reclaim_pages * page_size, - swa_reclaim_pages * page_size, + def set_disagg_move_gate(self, gate: Callable[[], bool]) -> None: + """Install the PD compaction gate on both physical sub-allocators.""" + assert self.lazy_compaction, ( + "PD disaggregation with the unified memory pool requires lazy " + "compaction (eager free-path compaction moves pages under " + "in-flight transfers)." ) + self.full_attn_allocator.disagg_move_gate = gate + self.swa_attn_allocator.disagg_move_gate = gate - def can_reserve( + def translate_loc_from_full_to_swa( self, - full_tokens: int | float, - swa_tokens: int | float, + kv_indices: torch.Tensor, *, - full_evictable_tokens: int = 0, - swa_evictable_tokens: int = 0, - empty_pool: bool = False, - require_token_slack: bool = False, - ) -> bool: - """Check pending FULL/SWA demand against the shared byte envelope. + out: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """SWA-layer read path: virtual TOKEN ids -> swa kernel-facing ids.""" + return self.swa_attn_allocator.translate_kv_loc_for_kernel(kv_indices, out=out) - Scheduler admission keeps the historical one-token strict slack at an - empty-pool boundary. Live admission checks the state reachable after - reclaiming the currently evictable pages from both sides. - """ - if full_tokens < 0 or swa_tokens < 0: - return False - if require_token_slack and ( - full_tokens >= self.size_full or swa_tokens >= self.size_swa - ): - return False + @property + def kernel_page_multiplier(self) -> int: + return self.full_attn_allocator.kernel_page_multiplier - return ( - self.reclaim_plan( - full_tokens, - swa_tokens, - full_evictable_tokens=full_evictable_tokens, - swa_evictable_tokens=swa_evictable_tokens, - empty_pool=empty_pool, - ) - is not None - ) + @property + def full_v2p_page_table(self) -> torch.Tensor: + """Page-level virtual->physical table of the full sub-pool.""" + return self.full_attn_allocator.virtual_to_physical - def _compaction_allowed(self) -> bool: - return all( - allocator.disagg_move_gate is None or allocator.disagg_move_gate() - for allocator in (self.full_attn_allocator, self.swa_attn_allocator) - ) + @property + def full_p2v_page_table(self) -> torch.Tensor: + """Page-level physical->virtual table of the full sub-pool.""" + return self.full_attn_allocator.physical_to_virtual - def _fits_page_demand( + def translate_kv_loc_for_kernel( self, - num_full_pages: int, - num_swa_pages: int, + loc: torch.Tensor, *, - full_reclaim_pages: int = 0, - swa_reclaim_pages: int = 0, - compacted: bool, - empty_pool: bool = False, - ) -> bool: - """Check one FULL/SWA page demand against a single allocator snapshot.""" - if min(num_full_pages, num_swa_pages) < 0: - return False + out: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """Full-pool virtual TOKEN ids -> kernel-facing ids.""" + return self.full_attn_allocator.translate_kv_loc_for_kernel(loc, out=out) - fa, sa = self.full_attn_allocator, self.swa_attn_allocator - if empty_pool: - full_live_pages = swa_live_pages = 0 - full_reclaim_pages = swa_reclaim_pages = 0 - else: - full_live_pages = fa.allocated_count() // self.page_size - swa_live_pages = sa.allocated_count() // self.page_size - full_reclaim_pages = min(full_live_pages, max(0, int(full_reclaim_pages))) - swa_reclaim_pages = min(swa_live_pages, max(0, int(swa_reclaim_pages))) + def translate_write_loc_for_kernel( + self, + loc: torch.Tensor, + *, + out: Optional[torch.Tensor] = None, + out_width: Optional[int] = None, + ) -> torch.Tensor: + """Widened virtual WRITE loc -> kernel-facing id. DCP is rejected for this + composite at argument validation, so it coincides with the read translate.""" + return self.full_attn_allocator.translate_write_loc_for_kernel( + loc, out=out, out_width=out_width + ) - full_live_pages -= full_reclaim_pages + @property + def swa_kernel_page_multiplier(self) -> int: + return self.swa_attn_allocator.kernel_page_multiplier + + @property + def swa_v2p_page_table(self) -> torch.Tensor: + """Page-level virtual->physical table of the SWA sub-pool.""" + return self.swa_attn_allocator.virtual_to_physical + + # -- alloc -- + + def alloc(self, need_size: int) -> Optional[torch.Tensor]: + with record_function("UnifiedSWAAlloc.alloc"): + if not self.ensure_capacity(need_size, need_size): + return None + # Snapshot the virtual PAGES full will consume, to bind them on swa too. + num_pages = need_size // self.page_size + fa = self.full_attn_allocator + new_virtual_pages = fa.free_virtual_ids[:num_pages].clone() + + v_tokens = fa.alloc(need_size) + # Post-pre-check failure can only be internal-state inconsistency. + assert v_tokens is not None, ( + "UnifiedSWA.alloc: full.alloc returned None after joint " + "pre-check passed — internal-state inconsistency" + ) + self.swa_attn_allocator.alloc_with_virtual(new_virtual_pages) + return v_tokens + + def alloc_extend( + self, + prefix_lens: torch.Tensor, + prefix_lens_cpu: torch.Tensor, + seq_lens: torch.Tensor, + seq_lens_cpu: torch.Tensor, + last_loc: torch.Tensor, + extend_num_tokens: int, + ) -> Optional[torch.Tensor]: + """Paged extend; returns virtual TOKEN ids. The same virtual page maps to + full- and swa-physical, so swa binds exactly what the full kernel consumed.""" + with record_function("UnifiedSWAAlloc.alloc_extend"): + num_new_pages = get_num_new_pages( + seq_lens=seq_lens_cpu, + page_size=self.page_size, + prefix_lens=prefix_lens_cpu, + ) + need_tokens = num_new_pages * self.page_size + if not self.ensure_capacity(need_tokens, need_tokens): + return None + + # Snapshot the virtual PAGES the kernel will consume; clone so swa keeps + # its view after the slice is consumed. + fa = self.full_attn_allocator + new_virtual_pages = fa.free_virtual_ids[:num_new_pages].clone() + + out_indices = fa.alloc_extend( + prefix_lens, + prefix_lens_cpu, + seq_lens, + seq_lens_cpu, + last_loc, + extend_num_tokens, + num_new_pages=num_new_pages, + ) + assert out_indices is not None, ( + "UnifiedSWA.alloc_extend: full.alloc_extend returned None " + "after joint pre-check passed — internal-state inconsistency" + ) + self.swa_attn_allocator.alloc_with_virtual(new_virtual_pages) + return out_indices # virtual TOKEN ids + + def alloc_decode( + self, + seq_lens: torch.Tensor, + seq_lens_cpu: torch.Tensor, + last_loc: torch.Tensor, + ) -> Optional[torch.Tensor]: + """Paged decode: one new token per request, consuming a page only when the + decode wraps.""" + with record_function("UnifiedSWAAlloc.alloc_decode"): + num_new_pages = get_num_new_pages( + seq_lens=seq_lens_cpu, page_size=self.page_size, decode=True + ) + need_tokens = num_new_pages * self.page_size + if not self.ensure_capacity(need_tokens, need_tokens): + return None + + fa = self.full_attn_allocator + new_virtual_pages = fa.free_virtual_ids[:num_new_pages].clone() + + out_indices = fa.alloc_decode(seq_lens, seq_lens_cpu, last_loc) + assert out_indices is not None, ( + "UnifiedSWA.alloc_decode: full.alloc_decode returned None " + "after joint pre-check passed — internal-state inconsistency" + ) + + if new_virtual_pages.numel() > 0: + self.swa_attn_allocator.alloc_with_virtual(new_virtual_pages) + + return out_indices # virtual TOKEN ids + + def is_slot_allocated(self, slot: int) -> bool: + """Token-slot surface = the full side (which owns the virtual ids).""" + return self.full_attn_allocator.is_slot_allocated(slot) + + def allocator_state_str(self) -> str: + return self.full_attn_allocator.allocator_state_str() + + # -- free -- + + def free(self, free_index: torch.Tensor) -> None: + with record_function("UnifiedSWAAlloc.free"): + if free_index is None or free_index.numel() == 0: + return + if self.free_group is not None: + self.free_group.append(self._copy_for_free_group(free_index)) + return + # Order is not load-bearing: the per-sub-pool v2p IS the mapping. Only + # swa needs the tombstone filter; full owns the ids, so all are bound. + v = free_index.detach().to(torch.int64) + v_pages = v // self.page_size + swa_v2p_pages = self.swa_attn_allocator.virtual_to_physical[v_pages] + # `> 0` strict: -1 = tombstoned, 0 = padding-sink page; both skipped. + live_token_mask = swa_v2p_pages > 0 + live_tokens = v[live_token_mask] + if live_tokens.numel() > 0: + self.swa_attn_allocator.free(live_tokens) + self.full_attn_allocator.free(v) + self.full_attn_allocator.clear_inverse_history() + self.swa_attn_allocator.clear_inverse_history() + + def free_swa(self, free_index: torch.Tensor) -> None: + """SWA tombstone path: release swa-physical, keep the virtual id and + full-physical live; `swa.v2p_page[v_page] = -1` IS the tombstone.""" + if free_index is None or free_index.numel() == 0: + return + v = free_index.detach().to(torch.int64) + ps = self.page_size + v_pages = v // ps + # `> 0` strict: -1 = tombstoned, page 0 = padding sink (never freeable). + swa_v2p_pages = self.swa_attn_allocator.virtual_to_physical[v_pages] + live = v[swa_v2p_pages > 0] + if live.numel() == 0: + return + if ps == 1: + # token == page and the live filter just deduped against the v2p + # table, so these ARE unique page ids -- same skip as `_free_lazy`. + self.swa_attn_allocator.free(live, _pages=live) + else: + self.swa_attn_allocator.free(live) + self.swa_attn_allocator.clear_inverse_history() + + def free_swa_segment(self, free_index: torch.Tensor, *, start_pos: int) -> None: + """free_swa() for a kv-row segment: `start_pos` promises a contiguous + ascending range, so page reps come from stride arithmetic, not `torch.unique`.""" + if free_index is None or free_index.numel() == 0: + return + if self.page_size == 1: + self.free_swa(free_index) + return + ps = self.page_size + reps = self.swa_attn_allocator._page_reps( + free_index.detach().to(torch.int64), start_pos + ) + # Keep only pages still bound on swa; freeing a tombstoned one would + # corrupt the hole list. `> 0` strict: -1 tombstoned, 0 padding sink. + rep_pages = reps // ps + swa_v2p_pages = self.swa_attn_allocator.virtual_to_physical[rep_pages] + live_reps = reps[swa_v2p_pages > 0] + if live_reps.numel() == 0: + return + self.swa_attn_allocator.free(live_reps, _pages=live_reps // ps) + self.swa_attn_allocator.clear_inverse_history() + + def free_full(self, free_index: torch.Tensor) -> None: + """Release the full-physical page and the virtual id, leaving the swa + side alone -- the caller already tombstoned it (`swa.v2p_page == -1`).""" + if free_index is None or free_index.numel() == 0: + return + if self.free_group is not None: + self.full_free_group.append(self._copy_for_free_group(free_index)) + return + self.full_attn_allocator.free(free_index.detach().to(torch.int64)) + self.full_attn_allocator.clear_inverse_history() + + def free_full_segment(self, free_index: torch.Tensor, *, start_pos: int) -> None: + if free_index is None or free_index.numel() == 0: + return + if self.page_size == 1: + # token == page: free_full already frees by exact ids, no dedup. + self.free_full(free_index) + return + # The swa v2p is the mapping, so a tombstoned swa page drops out of the + # two-sided segment path by itself; full-only is the same call. + self.free_segment(free_index, start_pos=start_pos) + + def set_full_to_swa_mapping( + self, full_indices: torch.Tensor, swa_indices: torch.Tensor + ) -> None: + """No-op stub for HiCache load-back: in shared mode the swa v2p IS the + mapping, and HiCache for shared SWA is out of scope.""" + return + + def clear_full_to_swa_mapping(self, full_indices: torch.Tensor) -> None: + # Paired with set_full_to_swa_mapping: shared mode has no mapping tensor. + return + + # -- free-group -- + + # Not the SWA parent's hooks: those open the parent's paged full allocator + # as a free group, and this composite's sub-pools defer on their own. + def free_group_begin(self) -> None: + BaseTokenToKVPoolAllocator.free_group_begin(self) + self.free_page_reps_group = [] + self.full_free_group = [] + + def free_group_end(self) -> None: + pending, self.free_page_reps_group = self.free_page_reps_group, None + full_free_group, self.full_free_group = self.full_free_group, [] + BaseTokenToKVPoolAllocator.free_group_end(self) + if full_free_group: + self.full_attn_allocator.free(torch.cat(full_free_group)) + self.full_attn_allocator.clear_inverse_history() + if pending: + self._release_page_reps(pending) + + def free_segment(self, free_index: torch.Tensor, *, start_pos: int) -> None: + """Fixed-shape counterpart of `free()`; see `MultiEndedAllocator._page_reps`. + Both sides share one page-rep derivation instead of dedup'ing twice.""" + if free_index is None or free_index.numel() == 0: + return + if self.page_size == 1: + self.free(free_index) + return + reps = self.full_attn_allocator._page_reps( + free_index.detach().to(torch.int64), start_pos + ) + if self.free_page_reps_group is None: + self._release_page_reps((reps,)) + else: + self.free_page_reps_group.append(reps) + + def _release_page_reps(self, pieces: Sequence[torch.Tensor]) -> None: + reps = pieces[0] if len(pieces) == 1 else torch.cat(tuple(pieces)) + v_pages = reps // self.page_size + # Same tombstone filter as `free`, but at PAGE granularity (page_size + # times smaller): `> 0` strict -- -1 = tombstoned, 0 = padding sink. + swa_v2p_pages = self.swa_attn_allocator.virtual_to_physical[v_pages] + live_pages = v_pages[swa_v2p_pages > 0] + if live_pages.numel() > 0: + self.swa_attn_allocator.free(live_pages * self.page_size, _pages=live_pages) + self.full_attn_allocator.free(reps, _pages=v_pages) + self.full_attn_allocator.clear_inverse_history() + self.swa_attn_allocator.clear_inverse_history() + + def _joint_capacity_memo_violations(self) -> List[str]: + """Idle-time twin of `MultiEndedAllocator._capacity_memo_violations` + for the composite joint view. Empty == healthy.""" + if ( + self._joint_avail_memo_epoch + != self.full_attn_allocator._chain_capacity_epoch() + ): + return [] + actual = self._compute_available_size() + if self._joint_avail_memo_tokens == actual: + return [] + return [ + f"[joint] stale available_size memo: " + f"cached={self._joint_avail_memo_tokens}, actual={actual}" + ] + + def clear(self) -> None: + self.full_attn_allocator.clear() + self.swa_attn_allocator.clear() + self.free_group = None + self.free_page_reps_group = None + self.full_free_group = [] + + # -- Lazy compaction hooks -- + + def set_latest_forward_done_event(self, event: Optional[torch.cuda.Event]) -> None: + """Forward the per-batch `forward_done` event to BOTH sub-allocators.""" + with record_function("UnifiedSWAAlloc.set_latest_forward_done_event"): + self.full_attn_allocator.set_latest_forward_done_event(event) + self.swa_attn_allocator.set_latest_forward_done_event(event) + + def set_inflight_forward( + self, + forward_done: torch.cuda.Event, + out_cache_loc_virtual: Optional[torch.Tensor], + ) -> None: + """Hand the forward's metadata to BOTH sub-pools; each materializes its own + write-set via its OWN v2p, and the forward writes both sides per token.""" + with record_function("UnifiedSWAAlloc.set_inflight_forward"): + self.full_attn_allocator.set_inflight_forward( + forward_done, out_cache_loc_virtual + ) + self.swa_attn_allocator.set_inflight_forward( + forward_done, out_cache_loc_virtual + ) + + @abstractmethod + def _build_swa_attn_allocator(self, **kwargs) -> MultiEndedAllocator: ... + + @abstractmethod + def _wire_peers(self) -> None: ... + + @abstractmethod + def _compute_available_size(self) -> int: ... + + @abstractmethod + def ensure_capacity(self, full_tokens: int, swa_tokens: int) -> bool: ... + + @abstractmethod + def _flush_targets(self): ... + + @abstractmethod + def _ask_float_for_room(self, need_tokens: int) -> None: ... + + +class UnifiedSWATokenToKVPoolAllocator(UnifiedSWAAllocatorBase): + """Two-ended FULL/SWA allocator with asymmetric shared-byte reservations.""" + + def _build_swa_attn_allocator(self, **kwargs) -> MultiEndedAllocator: + """The swa sub-allocator: an END pool in the 2-pool pair.""" + return MultiEndedAllocator( + sub_pool_name="swa", + is_id_owner=False, # non-owner; consumes virtuals minted by full + **kwargs, + ) + + def _wire_peers(self) -> None: + self.full_attn_allocator.bind_peer(self.swa_attn_allocator) + self.swa_attn_allocator.bind_peer(self.full_attn_allocator) + + def _compute_available_size(self) -> int: + """Joint byte budget in TOKENS: each composite alloc(1) consumes one + full-side AND one swa-side page under the same virtual id.""" + fa, sa = self.full_attn_allocator, self.swa_attn_allocator + e_f = fa.entry_bytes_per_page + e_s = sa.entry_bytes_per_page + # Direction-agnostic shared gap: the free byte band between the two pools. + if fa.grow_direction == "up": + gap_bytes = max(0, sa._byte_low_frontier() - fa._byte_high_frontier()) + else: + gap_bytes = max(0, fa._byte_low_frontier() - sa._byte_high_frontier()) + R_f = fa.num_pages - fa.min_page_index - fa._allocated_pages() + R_s = sa.num_pages - sa.min_page_index - sa._allocated_pages() + + if not self.lazy_compaction: + pages_by_bytes = gap_bytes // (e_f + e_s) + return min(pages_by_bytes, R_f, R_s) * self.page_size + + H_f = len(fa._free_phys_pages) + H_s = len(sa._free_phys_pages) + + K1 = min(H_f, H_s) # Phase 1: both drain + + # Phase 2: fewer-holes side extends; more-holes side keeps draining. + if H_f <= H_s: + e_phase2 = e_f + K_phase2_max = H_s + else: + e_phase2 = e_s + K_phase2_max = H_f + K2_room = K_phase2_max - K1 + K2 = min(K2_room, gap_bytes // e_phase2) if e_phase2 > 0 else K2_room + gap_bytes -= K2 * e_phase2 + + K3 = gap_bytes // (e_f + e_s) # Phase 3: both extend + + K_total = K1 + K2 + K3 + K_total = min(K_total, H_f + R_f, H_s + R_s) # index-space caps + return K_total * self.page_size + + def _flush_targets(self): + """Flush ALL members, including ones that are not short themselves: a + one-sided hole is unusable, and compacting it yields SHARED gap.""" + return (self.full_attn_allocator, self.swa_attn_allocator) + + def _ask_float_for_room(self, need_tokens: int) -> None: + """No float in a two-END chain -- nothing can slide.""" + return None + + def reclaim_plan( + self, + full_tokens: int | float, + swa_tokens: int | float, + *, + full_evictable_tokens: int = 0, + swa_evictable_tokens: int = 0, + empty_pool: bool = False, + ) -> Optional[Tuple[int, int]]: + """Return cumulative FULL/SWA eviction targets, or None if impossible. + + The tree evicts FULL before SWA. Minimize required SWA reclaim with all + evictable FULL available, then trim excess FULL reclaim. FULL eviction's + actual SWA cascade is counted by the tree's shared eviction tracker. + """ + if full_tokens < 0 or swa_tokens < 0: + return None + + page_size = self.page_size + full_pages = (math.ceil(full_tokens) + page_size - 1) // page_size + swa_pages = (math.ceil(swa_tokens) + page_size - 1) // page_size + # Restoring host SWA for device-resident FULL can require more new + # SWA pages than new FULL pages; only the shared budget constrains it. + + compacted = empty_pool or not self.lazy_compaction or self._compaction_allowed() + + def fits(full_reclaim_pages: int, swa_reclaim_pages: int) -> bool: + return self._fits_page_demand( + full_pages, + swa_pages, + full_reclaim_pages=full_reclaim_pages, + swa_reclaim_pages=swa_reclaim_pages, + compacted=compacted, + empty_pool=empty_pool, + ) + + if empty_pool: + return (0, 0) if fits(0, 0) else None + if fits(0, 0): + return (0, 0) + + max_full_pages = min( + self.full_attn_allocator.allocated_count() // page_size, + max(0, int(full_evictable_tokens)) // page_size, + ) + max_swa_pages = min( + self.swa_attn_allocator.allocated_count() // page_size, + max(0, int(swa_evictable_tokens)) // page_size, + ) + if not fits(max_full_pages, max_swa_pages): + return None + + def first_fit(high: int, predicate: Callable[[int], bool]) -> int: + low = 0 + while low < high: + mid = (low + high) // 2 + if predicate(mid): + high = mid + else: + low = mid + 1 + return low + + swa_reclaim_pages = first_fit( + max_swa_pages, lambda value: fits(max_full_pages, value) + ) + full_reclaim_pages = first_fit( + max_full_pages, lambda value: fits(value, swa_reclaim_pages) + ) + return ( + full_reclaim_pages * page_size, + swa_reclaim_pages * page_size, + ) + + def can_reserve( + self, + full_tokens: int | float, + swa_tokens: int | float, + *, + full_evictable_tokens: int = 0, + swa_evictable_tokens: int = 0, + empty_pool: bool = False, + require_token_slack: bool = False, + ) -> bool: + """Check pending FULL/SWA demand against the shared byte envelope. + + Scheduler admission keeps the historical one-token strict slack at an + empty-pool boundary. Live admission checks the state reachable after + reclaiming the currently evictable pages from both sides. + """ + if full_tokens < 0 or swa_tokens < 0: + return False + if require_token_slack and ( + full_tokens >= self.size_full or swa_tokens >= self.size_swa + ): + return False + + return ( + self.reclaim_plan( + full_tokens, + swa_tokens, + full_evictable_tokens=full_evictable_tokens, + swa_evictable_tokens=swa_evictable_tokens, + empty_pool=empty_pool, + ) + is not None + ) + + def _compaction_allowed(self) -> bool: + return all( + allocator.disagg_move_gate is None or allocator.disagg_move_gate() + for allocator in (self.full_attn_allocator, self.swa_attn_allocator) + ) + + def _fits_page_demand( + self, + num_full_pages: int, + num_swa_pages: int, + *, + full_reclaim_pages: int = 0, + swa_reclaim_pages: int = 0, + compacted: bool, + empty_pool: bool = False, + ) -> bool: + """Check one FULL/SWA page demand against a single allocator snapshot.""" + if min(num_full_pages, num_swa_pages) < 0: + return False + + fa, sa = self.full_attn_allocator, self.swa_attn_allocator + if empty_pool: + full_live_pages = swa_live_pages = 0 + full_reclaim_pages = swa_reclaim_pages = 0 + else: + full_live_pages = fa.allocated_count() // self.page_size + swa_live_pages = sa.allocated_count() // self.page_size + full_reclaim_pages = min(full_live_pages, max(0, int(full_reclaim_pages))) + swa_reclaim_pages = min(swa_live_pages, max(0, int(swa_reclaim_pages))) + + full_live_pages -= full_reclaim_pages swa_live_pages -= swa_reclaim_pages full_total_pages = full_live_pages + num_full_pages swa_total_pages = swa_live_pages + num_swa_pages @@ -579,195 +987,32 @@ def fits(candidate): if fits(mid): lo = mid else: - hi = mid - 1 - return lo - - def evict_to_free_tokens( - self, tree_cache, num_tokens: int, *, swa_num_tokens: Optional[int] = None - ) -> None: - from sglang.srt.mem_cache.base_prefix_cache import EvictParams - - if tree_cache is None or tree_cache.is_chunk_cache(): - return - required_swa = num_tokens if swa_num_tokens is None else swa_num_tokens - reclaim_plan = self.reclaim_plan( - num_tokens, - required_swa, - full_evictable_tokens=tree_cache.full_evictable_size(), - swa_evictable_tokens=tree_cache.swa_evictable_size(), - ) - if reclaim_plan is None: - return - full_reclaim, swa_reclaim = reclaim_plan - if full_reclaim or swa_reclaim: - tree_cache.evict_for_alloc( - EvictParams(num_tokens=full_reclaim, swa_num_tokens=swa_reclaim) - ) - # A zero-reclaim plan can still depend on compaction before allocation. - self.ensure_capacity(num_tokens, required_swa) - - @property - def draft_virtual_id_space(self) -> int: - return self.full_attn_allocator.max_slots - 1 - - def debug_print(self) -> str: - return ( - f"#full-available={self.full_attn_allocator.available_size()}, " - f"#swa-available={self.swa_attn_allocator.available_size()}, " - f"#joint-available={self.available_size()}" - ) - - def get_kvcache(self): - return self._kvcache - - def translate_kv_loc( - self, - loc: torch.Tensor, - *, - out: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - """Full-layer read path: virtual TOKEN ids -> full-physical TOKEN ids. - ``out=`` writes in place, for cuda-graph buffer stability.""" - result = self.full_attn_allocator.translate_kv_loc(loc, out=out) - return result - - def translate_kv_indices_for_transfer( - self, kv_indices: torch.Tensor - ) -> torch.Tensor: - """Translate virtual token ids to full-pool physical ids for PD.""" - return self.full_attn_allocator.translate_kv_loc(kv_indices.to(torch.int64)) - - def translate_swa_indices_for_transfer( - self, kv_indices: torch.Tensor - ) -> torch.Tensor: - """Virtual token ids -> physical, not kernel-facing, SWA ids for PD.""" - return self.swa_attn_allocator.translate_kv_loc(kv_indices.to(torch.int64)) - - def set_disagg_move_gate(self, gate: Callable[[], bool]) -> None: - """Install the PD compaction gate on both physical sub-allocators.""" - assert self.lazy_compaction, ( - "PD disaggregation with the unified memory pool requires lazy " - "compaction (eager free-path compaction moves pages under " - "in-flight transfers)." - ) - self.full_attn_allocator.disagg_move_gate = gate - self.swa_attn_allocator.disagg_move_gate = gate - - def translate_loc_from_full_to_swa( - self, - kv_indices: torch.Tensor, - *, - out: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - """SWA-layer read path: virtual TOKEN ids -> swa kernel-facing ids.""" - return self.swa_attn_allocator.translate_kv_loc_for_kernel(kv_indices, out=out) - - @property - def kernel_page_multiplier(self) -> int: - return self.full_attn_allocator.kernel_page_multiplier - - @property - def full_v2p_page_table(self) -> torch.Tensor: - """Page-level virtual->physical table of the full sub-pool.""" - return self.full_attn_allocator.virtual_to_physical - - @property - def full_p2v_page_table(self) -> torch.Tensor: - """Page-level physical->virtual table of the full sub-pool.""" - return self.full_attn_allocator.physical_to_virtual - - def translate_kv_loc_for_kernel( - self, - loc: torch.Tensor, - *, - out: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - """Full-pool virtual TOKEN ids -> kernel-facing ids.""" - return self.full_attn_allocator.translate_kv_loc_for_kernel(loc, out=out) - - def translate_write_loc_for_kernel( - self, - loc: torch.Tensor, - *, - out: Optional[torch.Tensor] = None, - out_width: Optional[int] = None, - ) -> torch.Tensor: - """Widened virtual WRITE loc -> kernel-facing id. DCP is rejected for this - composite at argument validation, so it coincides with the read translate.""" - return self.full_attn_allocator.translate_write_loc_for_kernel( - loc, out=out, out_width=out_width - ) - - @property - def swa_kernel_page_multiplier(self) -> int: - return self.swa_attn_allocator.kernel_page_multiplier - - @property - def swa_v2p_page_table(self) -> torch.Tensor: - """Page-level virtual->physical table of the SWA sub-pool.""" - return self.swa_attn_allocator.virtual_to_physical - - # -- alloc -- - - def alloc(self, need_size: int) -> Optional[torch.Tensor]: - with record_function("UnifiedSWAAlloc.alloc"): - if not self.ensure_capacity(need_size, need_size): - return None - # Snapshot the virtual PAGES full will consume, to bind them on swa too. - num_pages = need_size // self.page_size - fa = self.full_attn_allocator - new_virtual_pages = fa.free_virtual_ids[:num_pages].clone() - - v_tokens = fa.alloc(need_size) - # Post-pre-check failure can only be internal-state inconsistency. - assert v_tokens is not None, ( - "UnifiedSWA.alloc: full.alloc returned None after joint " - "pre-check passed — internal-state inconsistency" - ) - self.swa_attn_allocator.alloc_with_virtual(new_virtual_pages) - return v_tokens - - def alloc_extend( - self, - prefix_lens: torch.Tensor, - prefix_lens_cpu: torch.Tensor, - seq_lens: torch.Tensor, - seq_lens_cpu: torch.Tensor, - last_loc: torch.Tensor, - extend_num_tokens: int, - ) -> Optional[torch.Tensor]: - """Paged extend; returns virtual TOKEN ids. The same virtual page maps to - full- and swa-physical, so swa binds exactly what the full kernel consumed.""" - with record_function("UnifiedSWAAlloc.alloc_extend"): - num_new_pages = get_num_new_pages( - seq_lens=seq_lens_cpu, - page_size=self.page_size, - prefix_lens=prefix_lens_cpu, - ) - need_tokens = num_new_pages * self.page_size - if not self.ensure_capacity(need_tokens, need_tokens): - return None - - # Snapshot the virtual PAGES the kernel will consume; clone so swa keeps - # its view after the slice is consumed. - fa = self.full_attn_allocator - new_virtual_pages = fa.free_virtual_ids[:num_new_pages].clone() - - out_indices = fa.alloc_extend( - prefix_lens, - prefix_lens_cpu, - seq_lens, - seq_lens_cpu, - last_loc, - extend_num_tokens, - num_new_pages=num_new_pages, - ) - assert out_indices is not None, ( - "UnifiedSWA.alloc_extend: full.alloc_extend returned None " - "after joint pre-check passed — internal-state inconsistency" + hi = mid - 1 + return lo + + def evict_to_free_tokens( + self, tree_cache, num_tokens: int, *, swa_num_tokens: Optional[int] = None + ) -> bool | None: + from sglang.srt.mem_cache.base_prefix_cache import EvictParams + + if tree_cache is None or tree_cache.is_chunk_cache(): + return + required_swa = num_tokens if swa_num_tokens is None else swa_num_tokens + reclaim_plan = self.reclaim_plan( + num_tokens, + required_swa, + full_evictable_tokens=tree_cache.full_evictable_size(), + swa_evictable_tokens=tree_cache.swa_evictable_size(), + ) + if reclaim_plan is None: + return + full_reclaim, swa_reclaim = reclaim_plan + if full_reclaim or swa_reclaim: + tree_cache.evict_for_alloc( + EvictParams(num_tokens=full_reclaim, swa_num_tokens=swa_reclaim) ) - self.swa_attn_allocator.alloc_with_virtual(new_virtual_pages) - return out_indices # virtual TOKEN ids + # A zero-reclaim plan can still depend on compaction before allocation. + return self.ensure_capacity(num_tokens, required_swa) def alloc_extend_swa_tail( self, @@ -831,190 +1076,6 @@ def alloc_extend_swa_tail( self.swa_attn_allocator.alloc_with_virtual(tail_virtual_pages) return out_indices - def alloc_decode( - self, - seq_lens: torch.Tensor, - seq_lens_cpu: torch.Tensor, - last_loc: torch.Tensor, - ) -> Optional[torch.Tensor]: - """Paged decode: one new token per request, consuming a page only when the - decode wraps.""" - with record_function("UnifiedSWAAlloc.alloc_decode"): - num_new_pages = get_num_new_pages( - seq_lens=seq_lens_cpu, page_size=self.page_size, decode=True - ) - need_tokens = num_new_pages * self.page_size - if not self.ensure_capacity(need_tokens, need_tokens): - return None - - fa = self.full_attn_allocator - new_virtual_pages = fa.free_virtual_ids[:num_new_pages].clone() - - out_indices = fa.alloc_decode(seq_lens, seq_lens_cpu, last_loc) - assert out_indices is not None, ( - "UnifiedSWA.alloc_decode: full.alloc_decode returned None " - "after joint pre-check passed — internal-state inconsistency" - ) - - if new_virtual_pages.numel() > 0: - self.swa_attn_allocator.alloc_with_virtual(new_virtual_pages) - - return out_indices # virtual TOKEN ids - - def is_slot_allocated(self, slot: int) -> bool: - """Token-slot surface = the full side (which owns the virtual ids).""" - return self.full_attn_allocator.is_slot_allocated(slot) - - def allocator_state_str(self) -> str: - return self.full_attn_allocator.allocator_state_str() - - # -- free -- - - def free(self, free_index: torch.Tensor) -> None: - with record_function("UnifiedSWAAlloc.free"): - if free_index is None or free_index.numel() == 0: - return - if self.free_group is not None: - self.free_group.append(self._copy_for_free_group(free_index)) - return - # Order is not load-bearing: the per-sub-pool v2p IS the mapping. Only - # swa needs the tombstone filter; full owns the ids, so all are bound. - v = free_index.detach().to(torch.int64) - v_pages = v // self.page_size - swa_v2p_pages = self.swa_attn_allocator.virtual_to_physical[v_pages] - # `> 0` strict: -1 = tombstoned, 0 = padding-sink page; both skipped. - live_token_mask = swa_v2p_pages > 0 - live_tokens = v[live_token_mask] - if live_tokens.numel() > 0: - self.swa_attn_allocator.free(live_tokens) - self.full_attn_allocator.free(v) - self.full_attn_allocator.clear_inverse_history() - self.swa_attn_allocator.clear_inverse_history() - - def free_swa(self, free_index: torch.Tensor) -> None: - """SWA tombstone path: release swa-physical, keep the virtual id and - full-physical live; `swa.v2p_page[v_page] = -1` IS the tombstone.""" - if free_index is None or free_index.numel() == 0: - return - v = free_index.detach().to(torch.int64) - ps = self.page_size - v_pages = v // ps - # `> 0` strict: -1 = tombstoned, page 0 = padding sink (never freeable). - swa_v2p_pages = self.swa_attn_allocator.virtual_to_physical[v_pages] - live = v[swa_v2p_pages > 0] - if live.numel() == 0: - return - if ps == 1: - # token == page and the live filter just deduped against the v2p - # table, so these ARE unique page ids -- same skip as `_free_lazy`. - self.swa_attn_allocator.free(live, _pages=live) - else: - self.swa_attn_allocator.free(live) - self.swa_attn_allocator.clear_inverse_history() - - def free_swa_segment(self, free_index: torch.Tensor, *, start_pos: int) -> None: - """free_swa() for a kv-row segment: `start_pos` promises a contiguous - ascending range, so page reps come from stride arithmetic, not `torch.unique`.""" - if free_index is None or free_index.numel() == 0: - return - if self.page_size == 1: - self.free_swa(free_index) - return - ps = self.page_size - reps = self.swa_attn_allocator._page_reps( - free_index.detach().to(torch.int64), start_pos - ) - # Keep only pages still bound on swa; freeing a tombstoned one would - # corrupt the hole list. `> 0` strict: -1 tombstoned, 0 padding sink. - rep_pages = reps // ps - swa_v2p_pages = self.swa_attn_allocator.virtual_to_physical[rep_pages] - live_reps = reps[swa_v2p_pages > 0] - if live_reps.numel() == 0: - return - self.swa_attn_allocator.free(live_reps, _pages=live_reps // ps) - self.swa_attn_allocator.clear_inverse_history() - - def free_full(self, free_index: torch.Tensor) -> None: - """Release the full-physical page and the virtual id, leaving the swa - side alone -- the caller already tombstoned it (`swa.v2p_page == -1`).""" - if free_index is None or free_index.numel() == 0: - return - if self.free_group is not None: - self.full_free_group.append(self._copy_for_free_group(free_index)) - return - self.full_attn_allocator.free(free_index.detach().to(torch.int64)) - self.full_attn_allocator.clear_inverse_history() - - def free_full_segment(self, free_index: torch.Tensor, *, start_pos: int) -> None: - if free_index is None or free_index.numel() == 0: - return - if self.page_size == 1: - # token == page: free_full already frees by exact ids, no dedup. - self.free_full(free_index) - return - # The swa v2p is the mapping, so a tombstoned swa page drops out of the - # two-sided segment path by itself; full-only is the same call. - self.free_segment(free_index, start_pos=start_pos) - - def set_full_to_swa_mapping( - self, full_indices: torch.Tensor, swa_indices: torch.Tensor - ) -> None: - """No-op stub for HiCache load-back: in shared mode the swa v2p IS the - mapping, and HiCache for shared SWA is out of scope.""" - return - - def clear_full_to_swa_mapping(self, full_indices: torch.Tensor) -> None: - # Paired with set_full_to_swa_mapping: shared mode has no mapping tensor. - return - - # -- free-group -- - - # Not the SWA parent's hooks: those open the parent's paged full allocator - # as a free group, and this composite's sub-pools defer on their own. - def free_group_begin(self) -> None: - BaseTokenToKVPoolAllocator.free_group_begin(self) - self.free_page_reps_group = [] - self.full_free_group = [] - - def free_group_end(self) -> None: - pending, self.free_page_reps_group = self.free_page_reps_group, None - full_free_group, self.full_free_group = self.full_free_group, [] - BaseTokenToKVPoolAllocator.free_group_end(self) - if full_free_group: - self.full_attn_allocator.free(torch.cat(full_free_group)) - self.full_attn_allocator.clear_inverse_history() - if pending: - self._release_page_reps(pending) - - def free_segment(self, free_index: torch.Tensor, *, start_pos: int) -> None: - """Fixed-shape counterpart of `free()`; see `MultiEndedAllocator._page_reps`. - Both sides share one page-rep derivation instead of dedup'ing twice.""" - if free_index is None or free_index.numel() == 0: - return - if self.page_size == 1: - self.free(free_index) - return - reps = self.full_attn_allocator._page_reps( - free_index.detach().to(torch.int64), start_pos - ) - if self.free_page_reps_group is None: - self._release_page_reps((reps,)) - else: - self.free_page_reps_group.append(reps) - - def _release_page_reps(self, pieces: Sequence[torch.Tensor]) -> None: - reps = pieces[0] if len(pieces) == 1 else torch.cat(tuple(pieces)) - v_pages = reps // self.page_size - # Same tombstone filter as `free`, but at PAGE granularity (page_size - # times smaller): `> 0` strict -- -1 = tombstoned, 0 = padding sink. - swa_v2p_pages = self.swa_attn_allocator.virtual_to_physical[v_pages] - live_pages = v_pages[swa_v2p_pages > 0] - if live_pages.numel() > 0: - self.swa_attn_allocator.free(live_pages * self.page_size, _pages=live_pages) - self.full_attn_allocator.free(reps, _pages=v_pages) - self.full_attn_allocator.clear_inverse_history() - self.swa_attn_allocator.clear_inverse_history() - def verify_byte_accounting(self) -> List[str]: return ( _chain_byte_accounting_violations( @@ -1023,52 +1084,6 @@ def verify_byte_accounting(self) -> List[str]: + self._joint_capacity_memo_violations() ) - def _joint_capacity_memo_violations(self) -> List[str]: - """Idle-time twin of `MultiEndedAllocator._capacity_memo_violations` - for the composite joint view. Empty == healthy.""" - if ( - self._joint_avail_memo_epoch - != self.full_attn_allocator._chain_capacity_epoch() - ): - return [] - actual = self._compute_available_size() - if self._joint_avail_memo_tokens == actual: - return [] - return [ - f"[joint] stale available_size memo: " - f"cached={self._joint_avail_memo_tokens}, actual={actual}" - ] - - def clear(self) -> None: - self.full_attn_allocator.clear() - self.swa_attn_allocator.clear() - self.free_group = None - self.free_page_reps_group = None - self.full_free_group = [] - - # -- Lazy compaction hooks -- - - def set_latest_forward_done_event(self, event: Optional[torch.cuda.Event]) -> None: - """Forward the per-batch `forward_done` event to BOTH sub-allocators.""" - with record_function("UnifiedSWAAlloc.set_latest_forward_done_event"): - self.full_attn_allocator.set_latest_forward_done_event(event) - self.swa_attn_allocator.set_latest_forward_done_event(event) - - def set_inflight_forward( - self, - forward_done: torch.cuda.Event, - out_cache_loc_virtual: Optional[torch.Tensor], - ) -> None: - """Hand the forward's metadata to BOTH sub-pools; each materializes its own - write-set via its OWN v2p, and the forward writes both sides per token.""" - with record_function("UnifiedSWAAlloc.set_inflight_forward"): - self.full_attn_allocator.set_inflight_forward( - forward_done, out_cache_loc_virtual - ) - self.swa_attn_allocator.set_inflight_forward( - forward_done, out_cache_loc_virtual - ) - def flush_opportunistic(self) -> int: """Non-urgent flush of BOTH sub-allocators; sync-free.""" with record_function("UnifiedSWAAlloc.flush_opportunistic"): @@ -1084,7 +1099,7 @@ def flush_opportunistic(self) -> int: return fa.flush_opportunistic() + sa.flush_opportunistic() -class UnifiedMambaSWATokenToKVPoolAllocator(UnifiedSWATokenToKVPoolAllocator): +class UnifiedMambaSWATokenToKVPoolAllocator(UnifiedSWAAllocatorBase): """Tri-pool composite for models with full KV + SWA KV + mamba/conv state (both `mambaish_config` and `is_hybrid_swa`). @@ -1099,8 +1114,6 @@ class UnifiedMambaSWATokenToKVPoolAllocator(UnifiedSWATokenToKVPoolAllocator): `UnifiedMambaSlotAllocator`. """ - supports_asymmetric_reservation = False - def __init__( self, *, diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index b3e73e6d9722..c080231385d3 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/python/sglang/srt/mem_cache/common.py @@ -148,9 +148,11 @@ def maybe_cache_unfinished_req(req: Req, tree_cache: BasePrefixCache, **kwargs): tree_cache.cache_unfinished_req(req, **kwargs) -def evict_from_tree_cache(tree_cache: BasePrefixCache | None, num_tokens: int): +def evict_from_tree_cache( + tree_cache: BasePrefixCache | None, num_tokens: int +) -> bool | None: if tree_cache is not None and not tree_cache.is_chunk_cache(): - tree_cache.token_to_kv_pool_allocator.evict_to_free_tokens( + return tree_cache.token_to_kv_pool_allocator.evict_to_free_tokens( tree_cache, num_tokens ) diff --git a/python/sglang/srt/mem_cache/kv_cache_configurator.py b/python/sglang/srt/mem_cache/kv_cache_configurator.py index 7a256af7f432..8cf0ec1d66dc 100644 --- a/python/sglang/srt/mem_cache/kv_cache_configurator.py +++ b/python/sglang/srt/mem_cache/kv_cache_configurator.py @@ -51,7 +51,7 @@ is_swa_req_ring, ) from sglang.srt.mem_cache.allocator.unified_hybrid_swa import ( - UnifiedSWATokenToKVPoolAllocator, + UnifiedSWAAllocatorBase, ) from sglang.srt.mem_cache.allocator.unified_mamba import ( UnifiedMambaTokenToKVPoolAllocator, @@ -537,7 +537,7 @@ def _init_pools( token_to_kv_pool_allocator, ( UnifiedMambaTokenToKVPoolAllocator, - UnifiedSWATokenToKVPoolAllocator, + UnifiedSWAAllocatorBase, ), ): draft_virtual_id_space = ( @@ -560,7 +560,7 @@ def _init_pools( if ( isinstance( token_to_kv_pool_allocator, - UnifiedSWATokenToKVPoolAllocator, + UnifiedSWAAllocatorBase, ) and self.is_hybrid_swa ): @@ -2137,7 +2137,7 @@ def _build_token_to_kv_pool_allocator( else: swa_allocator = token_to_kv_pool_allocator uses_unified_virtual_ids = isinstance( - swa_allocator, UnifiedSWATokenToKVPoolAllocator + swa_allocator, UnifiedSWAAllocatorBase ) has_draft_swa_layers = ( not self.is_hybrid_swa_mtp_draft or self.draft_swa_full_capacity diff --git a/python/sglang/srt/mem_cache/kv_index_translator.py b/python/sglang/srt/mem_cache/kv_index_translator.py index 1ccb5b35077c..b7c71ae096e4 100644 --- a/python/sglang/srt/mem_cache/kv_index_translator.py +++ b/python/sglang/srt/mem_cache/kv_index_translator.py @@ -65,7 +65,7 @@ build_kv_read_table_packed, ) from sglang.srt.mem_cache.allocator.unified_hybrid_swa import ( - UnifiedSWATokenToKVPoolAllocator, + UnifiedSWAAllocatorBase, ) from sglang.srt.mem_cache.allocator.unified_mamba import ( UnifiedMambaTokenToKVPoolAllocator, @@ -121,7 +121,7 @@ def __init__( self.is_translating = ( isinstance( token_to_kv_pool_allocator, - (UnifiedMambaTokenToKVPoolAllocator, UnifiedSWATokenToKVPoolAllocator), + (UnifiedMambaTokenToKVPoolAllocator, UnifiedSWAAllocatorBase), ) and token_to_kv_pool_allocator.get_kvcache() is token_to_kv_pool ) @@ -139,7 +139,7 @@ def __init__( # DCP read ids stay WIDENED to the consumer: selecting this rank's # share changes the length, so only the production site can do it. self.defer_read_translate = get_parallel().attn_dcp_size > 1 - if isinstance(alloc, UnifiedSWATokenToKVPoolAllocator): + if isinstance(alloc, UnifiedSWAAllocatorBase): self._swa_v2p_table = alloc.swa_v2p_page_table self._swa_page_multiplier = alloc.swa_kernel_page_multiplier self._swa_write_loc_from_full = self._swa_write_loc_unified diff --git a/python/sglang/srt/mem_cache/unified_cache/components/swa_component.py b/python/sglang/srt/mem_cache/unified_cache/components/swa_component.py index 705cdf0f32ce..fb5d3cc2158c 100644 --- a/python/sglang/srt/mem_cache/unified_cache/components/swa_component.py +++ b/python/sglang/srt/mem_cache/unified_cache/components/swa_component.py @@ -211,11 +211,11 @@ def _translate_full_to_swa(self, full_indices: torch.Tensor) -> torch.Tensor: def _unified_allocator(self): """The unified SWA composite, or None when running on the static pool.""" from sglang.srt.mem_cache.allocator.unified_hybrid_swa import ( - UnifiedSWATokenToKVPoolAllocator, + UnifiedSWAAllocatorBase, ) allocator = self.cache.token_to_kv_pool_allocator - if isinstance(allocator, UnifiedSWATokenToKVPoolAllocator): + if isinstance(allocator, UnifiedSWAAllocatorBase): return allocator return None diff --git a/test/registered/unit/managers/test_priority_scheduling_disaggregation.py b/test/registered/unit/managers/test_priority_scheduling_disaggregation.py index b723f4f52563..6f707c47f0c1 100644 --- a/test/registered/unit/managers/test_priority_scheduling_disaggregation.py +++ b/test/registered/unit/managers/test_priority_scheduling_disaggregation.py @@ -162,7 +162,6 @@ def pre_alloc_mock(req, prefix_indices=None, prefix_len=0, total_prefix_len=0): queue.transfer_queue = SimpleNamespace(queue=[], enable_staging=False) queue.kv_manager = SimpleNamespace( kv_args=SimpleNamespace(state_types=[]), - uses_separate_draft_kv_indices=False, ) queue.tree_cache = MagicMock() diff --git a/test/registered/unit/mem_cache/test_unified_tri_pool.py b/test/registered/unit/mem_cache/test_unified_tri_pool.py index 755d96848f29..0bf7ecaa644e 100644 --- a/test/registered/unit/mem_cache/test_unified_tri_pool.py +++ b/test/registered/unit/mem_cache/test_unified_tri_pool.py @@ -961,11 +961,19 @@ def _factory_kwargs(self, **over): return kw def test_budget_sizing_and_boot_signature(self): + from sglang.srt.mem_cache.allocator.unified_hybrid_swa import ( + UnifiedSWAAllocatorBase, + UnifiedSWATokenToKVPoolAllocator, + ) + budget = 1 << 20 bundle = init_unified_mamba_swa_pools( **self._factory_kwargs(unified_total_bytes=budget) ) pool = bundle.unified_memory_pool + allocator = bundle.token_to_kv_pool_allocator + self.assertIsInstance(allocator, UnifiedSWAAllocatorBase) + self.assertNotIsInstance(allocator, UnifiedSWATokenToKVPoolAllocator) # Buffer = budget + the state pool's bytes (budget captured AFTER the # state carve-out), never the token-count re-sum. state_bytes = 4 * pool.spec("mamba").entry_bytes() From 564845d35c12f8391c85d966f21637d9a7a44236 Mon Sep 17 00:00:00 2001 From: yhzhuang Date: Sun, 13 Sep 2026 19:40:59 -0700 Subject: [PATCH 40/55] Use shared byte capacity for unified SWA pools Consolidate the historical PR iterations through the last upstream merge (f83adaf9c0), preserving its exact tree. Later review and refactor commits remain separate. --- python/sglang/srt/arg_groups/kv_cache_hook.py | 9 +- python/sglang/srt/disaggregation/decode.py | 2 - python/sglang/srt/managers/schedule_policy.py | 355 +++++++++++++----- python/sglang/srt/managers/scheduler.py | 60 ++- .../scheduler_components/invariant_checker.py | 46 ++- .../pool_stats_observer.py | 33 +- python/sglang/srt/mem_cache/allocator/base.py | 10 + .../mem_cache/allocator/unified_hybrid_swa.py | 308 +++++++++++++-- .../mem_cache/allocator/unified_sub_pool.py | 23 +- python/sglang/srt/mem_cache/common.py | 31 +- .../srt/mem_cache/kv_cache_configurator.py | 16 +- .../srt/mem_cache/unified_memory_pool.py | 36 +- .../sglang/srt/model_executor/model_runner.py | 9 +- .../srt/model_executor/pool_configurator.py | 45 ++- .../test_scheduler_init_req_max_new_tokens.py | 1 + .../test_hisparse_max_token_pool_size.py | 4 + .../mem_cache/test_multi_ended_allocator.py | 96 +++++ .../model_executor/test_pool_configurator.py | 74 +++- 18 files changed, 948 insertions(+), 210 deletions(-) diff --git a/python/sglang/srt/arg_groups/kv_cache_hook.py b/python/sglang/srt/arg_groups/kv_cache_hook.py index 4d473c79ef23..b5ba4eac39c4 100644 --- a/python/sglang/srt/arg_groups/kv_cache_hook.py +++ b/python/sglang/srt/arg_groups/kv_cache_hook.py @@ -200,6 +200,10 @@ def handle_unified_memory_pool(server_args: Any) -> None: if not cfg.enable_unified_memory: return if cfg.disaggregation_mode != "null": + assert cfg.pp_size == 1, ( + "--enable-unified-memory with PD disaggregation does not support " + "pipeline parallelism (--pp-size > 1)." + ) # Constraints of the whole-envelope transfer; see # UnifiedMLATokenToKVPool.get_contiguous_buf_infos. assert cfg.disaggregation_transfer_backend == "mooncake", ( @@ -207,11 +211,6 @@ def handle_unified_memory_pool(server_args: Any) -> None: "the mooncake transfer backend; got " f"{cfg.disaggregation_transfer_backend!r}." ) - assert cfg.pp_size == 1, ( - "--enable-unified-memory with PD disaggregation does not support " - "pipeline parallelism (whole-envelope transfer has no per-layer " - "entries to subset)." - ) assert not envs.SGLANG_DISABLE_LAZY_COMPACTION.get(), ( "--enable-unified-memory with PD disaggregation requires lazy " "compaction; unset SGLANG_DISABLE_LAZY_COMPACTION." diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index c1cf16029d7a..98ce8dadac44 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -454,7 +454,6 @@ def _reclaim_swa_tail_capacity( EvictParams(swa_num_tokens=required - available) ) available = self.token_to_kv_pool_allocator.swa_available_size() - if available < required: return ( f"SWA eviction insufficient: needed={required}, " @@ -856,7 +855,6 @@ def resume_retracted_reqs( break if uses_swa_tail_prealloc and swa_required > swa_allocatable_tokens: break - resumed_reqs.append(req) indices_to_remove.add(i) req.is_retracted = False diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index 9fa67710b31a..85810debb4ac 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -54,6 +54,7 @@ ) from sglang.srt.mem_cache.allocator.unified_hybrid_swa import ( UnifiedMambaSWATokenToKVPoolAllocator, + UnifiedSWATokenToKVPoolAllocator, ) from sglang.srt.mem_cache.allocator.unified_mamba import ( UnifiedMambaTokenToKVPoolAllocator, @@ -78,6 +79,39 @@ os.environ.get("SGLANG_CLIP_MAX_NEW_TOKENS_ESTIMATION", "4096") ) + +def estimate_swa_kv_tokens( + extend_input_len: int, + max_new_tokens: int, + *, + sliding_window_size: Optional[int], + page_size: int, + allocation_limit: Optional[int] = None, + host_hit_length: int = 0, + include_allocated_tail: bool = True, +) -> int: + """Peak SWA reservation for one prefill/decode request.""" + if sliding_window_size is None or sliding_window_size <= 0: + reserved = extend_input_len + max_new_tokens + page_size + else: + allocated = ( + extend_input_len + if allocation_limit is None + else min(extend_input_len, allocation_limit) + ) + allocated_tail = ( + max(allocated - sliding_window_size, 0) if include_allocated_tail else 0 + ) + reserved = ( + allocated_tail + + min(extend_input_len + max_new_tokens, sliding_window_size) + + page_size + ) + if host_hit_length > 0: + reserved += -(-host_hit_length // page_size) * page_size + return reserved + + # Threshold for in-batch prefix cache. # If a request has a matched prefix length (against existing cache) less than this value, # the scheduler runs the in-batch prefix caching check for this request. @@ -605,8 +639,16 @@ def __init__( self.token_to_kv_pool_allocator, PureSWATokenToKVPoolAllocator ) self.is_hybrid_ssm_cache = self.tree_cache.supports_mamba() + self.is_unified_swa = ( + isinstance( + self.token_to_kv_pool_allocator, UnifiedSWATokenToKVPoolAllocator + ) + and self.token_to_kv_pool_allocator.supports_asymmetric_reservation + ) - self.rem_swa_token_offset = 0 + self.rem_swa_token_offset = ( + num_mixed_decode_tokens if self.is_unified_swa else 0 + ) # A new state slot eats shared-gap bytes that `rem_total_tokens` counts # as free, so reserve per slot or admission over-commits. Gate on the @@ -787,21 +829,14 @@ def _swa_budget_for_req( def _swa_reserved_tokens( self, extend_input_len: int, max_new_tokens: int, swa_host_hit_length: int = 0 ) -> int: - """SWA slots a request adds to its own sliding window + page slack + the - load-back charge. Shared floor of _swa_budget_for_req and _swa_chunk_cap. - - The headroom is min(extend + decode, window), not a constant window: a - request contributes only extend + decode fresh tokens to its window and - a cached SWA prefix funds the rest. Charging a full window double-counted - a short cached-prefix resume and livelocked admission at a ~2-window - pool; keeping extend in the min() holds the reservation >= the prefill - allocation so an admitted request cannot OOM.""" - window = self.tree_cache.sliding_window_size - headroom = min(extend_input_len + max_new_tokens, window) - reserved = headroom + self.page_size - if swa_host_hit_length > 0: - reserved += self.ceil_paged_tokens(swa_host_hit_length) - return reserved + return estimate_swa_kv_tokens( + extend_input_len, + max_new_tokens, + sliding_window_size=self.tree_cache.sliding_window_size, + page_size=self.page_size, + host_hit_length=swa_host_hit_length, + include_allocated_tail=False, + ) def _swa_new_tokens(self, req: Req) -> int: """Tokens a request may still decode, for SWA headroom sizing. Mirrors @@ -813,10 +848,45 @@ def _swa_new_tokens(self, req: Req) -> int: CLIP_MAX_NEW_TOKENS, ) - def _swa_chunk_cap(self, max_new_tokens: int, swa_host_hit_length: int = 0) -> int: - """Largest page-aligned extend chunk the SWA pool can admit right now, - keeping a sliding window of headroom below rem_swa_tokens; 0 if not - even one page fits. Only valid when is_hybrid_swa is True. + def _unified_swa_reservation_fits( + self, + full_tokens: int, + swa_tokens: int, + *, + empty_pool: bool = False, + ) -> bool: + if not self.is_unified_swa: + return True + allocator = self.token_to_kv_pool_allocator + if empty_pool: + full_evictable_tokens = 0 + swa_evictable_tokens = 0 + else: + full_tokens += self.rem_total_token_offset + swa_tokens += self.rem_swa_token_offset + full_evictable_tokens = self.tree_cache.full_evictable_size() + swa_evictable_tokens = self.tree_cache.swa_evictable_size() + return allocator.can_reserve( + full_tokens, + swa_tokens, + full_evictable_tokens=full_evictable_tokens, + swa_evictable_tokens=swa_evictable_tokens, + empty_pool=empty_pool, + require_token_slack=empty_pool, + ) + + def _swa_chunk_cap( + self, + max_new_tokens: int, + swa_host_hit_length: int = 0, + *, + max_chunk_tokens: Optional[int] = None, + ) -> int: + """Largest page-aligned extend chunk the SWA pool can admit right now. + + Split pools reserve headroom below rem_swa_tokens; unified pools use + the joint byte budget. Returns 0 if not even one page fits. Only valid + when is_hybrid_swa is True. Escape hatch for a request whose budget can never pass the _swa_budget_for_req gate (extend near/above the pool size, or a large @@ -824,17 +894,41 @@ def _swa_chunk_cap(self, max_new_tokens: int, swa_host_hit_length: int = 0) -> i forever (head-of-line livelock). Shrinking is sound because past a chunk boundary only the sliding window stays locked — the rest turns evictable — so each pass's transient footprint fits the pool.""" - # extend_input_len=0: this solves for the extend chunk itself, so the - # reserved headroom is the post-chunk decode window only. - cap = int(self.rem_swa_tokens) - self._swa_reserved_tokens( - 0, max_new_tokens, swa_host_hit_length - ) - if cap <= 0: - return 0 - return cap // self.page_size * self.page_size + if not self.is_unified_swa: + # extend_input_len=0: this solves for the extend chunk itself, so + # the reserved headroom is the post-chunk decode window only. + cap = int(self.rem_swa_tokens) - self._swa_reserved_tokens( + 0, max_new_tokens, swa_host_hit_length + ) + if cap <= 0: + return 0 + cap = cap // self.page_size * self.page_size + return cap + + assert max_chunk_tokens is not None + cap = max_chunk_tokens // self.page_size * self.page_size + lo, hi = 0, cap // self.page_size + while lo < hi: + mid = (lo + hi + 1) // 2 + chunk_tokens = mid * self.page_size + full_tokens = self.ceil_paged_tokens(chunk_tokens) + self.page_size + swa_tokens = self._swa_budget_for_req( + chunk_tokens, max_new_tokens, swa_host_hit_length + ) + if self._unified_swa_reservation_fits(full_tokens, swa_tokens): + lo = mid + else: + hi = mid - 1 + return lo * self.page_size def _swa_req_never_fits( - self, extend_input_len: int, max_new_tokens: int, swa_host_hit_length: int = 0 + self, + extend_input_len: int, + max_new_tokens: int, + swa_host_hit_length: int = 0, + *, + full_tokens: Optional[int] = None, + unified_swa_tokens: Optional[int] = None, ) -> bool: """True when a request's SWA budget exceeds the *entire* SWA pool, so it can never be admitted whole no matter how far the pool drains. @@ -846,13 +940,78 @@ def _swa_req_never_fits( decode headroom collapses the SWA evictable cushion and forces running requests to retract (observed as a severe retraction/re-prefill storm on hybrid-SWA models at high concurrency).""" - capacity = self.token_to_kv_pool_allocator.size_swa - return ( - self._swa_budget_for_req( - extend_input_len, max_new_tokens, swa_host_hit_length + swa_tokens = self._swa_budget_for_req( + extend_input_len, max_new_tokens, swa_host_hit_length + ) + if self.is_unified_swa: + assert full_tokens is not None + assert unified_swa_tokens is not None + return not self._unified_swa_reservation_fits( + full_tokens, unified_swa_tokens, empty_pool=True + ) + if self._swa_req_ring: + return swa_tokens > self.token_to_kv_pool_allocator.size_swa + return swa_tokens >= self.token_to_kv_pool_allocator.size_swa + + def _check_swa_admission( + self, + req: Req, + *, + extend_input_len: int, + total_tokens: int, + ) -> tuple[bool, Optional[int]]: + """Check SWA pressure and return the usable chunk limit.""" + if not self.is_unified_swa: + extend_input_len = self.ceil_paged_tokens(extend_input_len) + max_new_tokens = self._swa_new_tokens(req) + swa_needed = self._swa_budget_for_req( + extend_input_len, + max_new_tokens, + swa_host_hit_length=req.swa_host_hit_length, + ) + reservation_fits = ( + self._unified_swa_reservation_fits(total_tokens, swa_needed) + if self.is_unified_swa + else swa_needed <= self.rem_swa_tokens + if self._swa_req_ring + else swa_needed < self.rem_swa_tokens + ) + if reservation_fits: + return True, self.rem_chunk_tokens + + full_ever_tokens = unified_swa_tokens = None + if self.is_unified_swa: + full_input_tokens = len(req.full_untruncated_fill_ids) + full_ever_tokens = ( + self.ceil_paged_tokens(full_input_tokens) + + max_new_tokens + + self.page_size ) - >= capacity + unified_swa_tokens = self._swa_budget_for_req( + full_input_tokens, max_new_tokens + ) + if not self._swa_req_never_fits( + extend_input_len, + max_new_tokens, + req.swa_host_hit_length, + full_tokens=full_ever_tokens, + unified_swa_tokens=unified_swa_tokens, + ): + return False, None + + max_chunk_tokens = self.rem_chunk_tokens + chunk_max_new = max_new_tokens + if self.is_unified_swa: + max_chunk_tokens = min(max_chunk_tokens or 0, max(0, extend_input_len - 1)) + chunk_max_new = 0 + swa_cap = self._swa_chunk_cap( + chunk_max_new, + req.swa_host_hit_length, + max_chunk_tokens=max_chunk_tokens, ) + if self.rem_chunk_tokens is None or swa_cap <= 0: + return False, None + return True, min(self.rem_chunk_tokens, swa_cap) def _mamba_gap_budget_for_req(self, req: Req) -> int: """Shared-gap reservation (full-token-equivalents) for a request's new @@ -873,9 +1032,12 @@ def ceil_paged_tokens(self, tokens: int) -> int: return -(-tokens // self.page_size) * self.page_size def budget_state(self): - no_token = self.rem_total_tokens <= 0 or self.cur_rem_tokens <= 0 - if not no_token and self.is_hybrid_swa: - no_token = self.rem_swa_tokens <= 0 + if self.is_unified_swa: + no_token = not self._unified_swa_reservation_fits(0, 0) + else: + no_token = self.rem_total_tokens <= 0 or self.cur_rem_tokens <= 0 + if not no_token and self.is_hybrid_swa: + no_token = self.rem_swa_tokens <= 0 # Gate new mamba slots separately: rem_total_tokens' full_evictable can't # cover a mamba slot, which needs mamba-recoverable bytes (see __init__). if not no_token and self.rem_mamba_slots is not None: @@ -1062,8 +1224,15 @@ def add_chunked_req(self, req: Req): if self.dllm_config is not None: _rem_tokens = self._get_dllm_remain_tokens() else: - _rem_tokens = min(self.rem_chunk_tokens, int(self.rem_total_tokens)) - if self.is_hybrid_swa and not self._swa_req_ring: + if self.is_unified_swa: + _rem_tokens = self.rem_chunk_tokens + else: + _rem_tokens = min(self.rem_chunk_tokens, int(self.rem_total_tokens)) + if ( + self.is_hybrid_swa + and not self.is_unified_swa + and not self._swa_req_ring + ): # alloc_extend needs extend_num_tokens + page_size per request, # so reserve one page here to avoid OOM. # Ring mode skips it: rem_swa_tokens counts slots, not chunk tokens. @@ -1091,6 +1260,23 @@ def add_chunked_req(self, req: Req): cand_extend_input_len = len(req.full_untruncated_fill_ids) - len( req.prefix_indices ) + if self.is_unified_swa: + candidate_tokens = min(cand_extend_input_len, _rem_tokens) + candidate_finishes = candidate_tokens >= cand_extend_input_len + candidate_max_new = self._swa_new_tokens(req) if candidate_finishes else 0 + if not self._unified_swa_reservation_fits( + self.ceil_paged_tokens(candidate_tokens) + + candidate_max_new + + self.page_size, + self._swa_budget_for_req(candidate_tokens, candidate_max_new), + ): + max_partial = candidate_tokens + if candidate_finishes: + max_partial = max(0, candidate_tokens - 1) + joint_cap = self._swa_chunk_cap(0, max_chunk_tokens=max_partial) + if joint_cap <= 0: + return req + _rem_tokens = min(_rem_tokens, joint_cap) truncated = cand_extend_input_len > _rem_tokens new_len = min(cand_extend_input_len, _rem_tokens) req.set_extend_range(len(req.prefix_indices), len(req.prefix_indices) + new_len) @@ -1135,16 +1321,22 @@ def add_one_req_ignore_eos(self, req: Req): # Shared Mamba pool: fold the new mamba state's shared-gap cost into the # budget gate so admission can't over-commit (0 for baseline / non-Mamba). paged_input += self._mamba_gap_budget_for_req(req) - if paged_input > min(self.cur_rem_tokens, self.rem_total_tokens): - return AddReqResult.NO_TOKEN if self.is_hybrid_swa: - if ( - self._swa_budget_for_req( - cand_extend_input_len, self._swa_new_tokens(req) - ) - > self.rem_swa_tokens + max_new_tokens = self._swa_new_tokens(req) + swa_needed = self._swa_budget_for_req(cand_extend_input_len, max_new_tokens) + if self.is_unified_swa: + if not self._unified_swa_reservation_fits( + paged_input + max_new_tokens + self.page_size, + swa_needed, + ): + return AddReqResult.NO_TOKEN + elif ( + paged_input > min(self.cur_rem_tokens, self.rem_total_tokens) + or swa_needed > self.rem_swa_tokens ): return AddReqResult.NO_TOKEN + elif paged_input > min(self.cur_rem_tokens, self.rem_total_tokens): + return AddReqResult.NO_TOKEN def add_req_state(r, insert_sort=False): new_token_ratio = ( @@ -1289,44 +1481,23 @@ def add_one_req( # this returns 0, so the debit sites below reuse the value. mamba_gap_reserve = self._mamba_gap_budget_for_req(req) total_tokens += mamba_gap_reserve - # adjusting the input_tokens based on host_hit_length and page_size - real_input_tokens = cand_extend_input_len - req.host_hit_length - real_input_tokens = self.ceil_paged_tokens(real_input_tokens) + swa_extend_input_len = cand_extend_input_len - req.host_hit_length + real_input_tokens = self.ceil_paged_tokens(swa_extend_input_len) prefix_len = len(req.prefix_indices) - if total_tokens >= self.rem_total_tokens: + if not self.is_unified_swa and total_tokens >= self.rem_total_tokens: return AddReqResult.NO_TOKEN chunk_tokens_limit = self.rem_chunk_tokens if self.is_hybrid_swa: - # host-hit prefix is loaded back, not re-prefilled, so the SWA peak is - # driven only by the freshly-prefilled tail (the loaded window is - # charged separately via swa_host_hit_length). - swa_needed = self._swa_budget_for_req( - real_input_tokens, - self._swa_new_tokens(req), - swa_host_hit_length=req.swa_host_hit_length, + can_admit, chunk_tokens_limit = self._check_swa_admission( + req, + extend_input_len=swa_extend_input_len, + total_tokens=total_tokens, ) - # Ring-slot capacity is exact, so needing exactly what is left still - # fits; the legacy SWA-token path keeps its conservative `>=`. - if ( - swa_needed > self.rem_swa_tokens - if self._swa_req_ring - else swa_needed >= self.rem_swa_tokens - ): - if not self._swa_req_never_fits( - real_input_tokens, - self._swa_new_tokens(req), - req.swa_host_hit_length, - ): - return AddReqResult.NO_TOKEN - swa_cap = self._swa_chunk_cap( - self._swa_new_tokens(req), req.swa_host_hit_length - ) - if self.rem_chunk_tokens is None or swa_cap <= 0: - return AddReqResult.NO_TOKEN - chunk_tokens_limit = min(self.rem_chunk_tokens, swa_cap) + if not can_admit: + return AddReqResult.NO_TOKEN if ( self.rem_chunk_tokens is None @@ -1339,34 +1510,18 @@ def add_one_req( return AddReqResult.OTHER with self._lock_node(req.last_node): - # self.rem_total_tokens may decrease after the lock acquisition - if total_tokens >= self.rem_total_tokens: + # Capacity may decrease after the lock acquisition. + if not self.is_unified_swa and total_tokens >= self.rem_total_tokens: return AddReqResult.NO_TOKEN if self.is_hybrid_swa: - # self.rem_swa_tokens may decrease after the lock acquisition - swa_needed = self._swa_budget_for_req( - real_input_tokens, - self._swa_new_tokens(req), - swa_host_hit_length=req.swa_host_hit_length, + can_admit, chunk_tokens_limit = self._check_swa_admission( + req, + extend_input_len=swa_extend_input_len, + total_tokens=total_tokens, ) - if ( - swa_needed > self.rem_swa_tokens - if self._swa_req_ring - else swa_needed >= self.rem_swa_tokens - ): - if not self._swa_req_never_fits( - real_input_tokens, - self._swa_new_tokens(req), - req.swa_host_hit_length, - ): - return AddReqResult.NO_TOKEN - swa_cap = self._swa_chunk_cap( - self._swa_new_tokens(req), req.swa_host_hit_length - ) - if self.rem_chunk_tokens is None or swa_cap <= 0: - return AddReqResult.NO_TOKEN - chunk_tokens_limit = min(self.rem_chunk_tokens, swa_cap) + if not can_admit: + return AddReqResult.NO_TOKEN # Negotiate only after every KV-budget gate (a NO_TOKEN rank must # report not-prefillable via finalize()) and before init_load_back diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index f6bf9b53ad9c..4e849d6b98e4 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -212,6 +212,7 @@ AddReqResult, PrefillAdder, SchedulePolicy, + estimate_swa_kv_tokens, ) from sglang.srt.managers.scheduler_components.batch_result_processor import ( SchedulerBatchResultProcessor, @@ -279,6 +280,9 @@ validate_input_length, ) from sglang.srt.mem_cache import kv_cache_builder +from sglang.srt.mem_cache.allocator.unified_hybrid_swa import ( + UnifiedSWATokenToKVPoolAllocator, +) from sglang.srt.mem_cache.common import ( maybe_cache_unfinished_req, release_kv_cache, @@ -2519,23 +2523,63 @@ def init_req_max_new_tokens(self, req): ) max_new_tokens = min(max_new_tokens, self.max_new_tokens_limit) - # Keep this bound consistent with PrefillAdder's admission budget: - # ceil_page(input_len) + max_new_tokens + page_size must be strictly - # smaller than max_total_num_tokens. Otherwise a request can be accepted - # into the waiting queue but can never be scheduled, blocking the queue - # and eventually making health checks fail. + # Keep this bound consistent with PrefillAdder's admission budget. paged_input_len = -(-input_len // self.page_size) * self.page_size - req.sampling_params.max_new_tokens = max( + max_new_tokens = max( 0, min( max_new_tokens, self.max_req_len - input_len - 1, + ), + ) + allocator = self.token_to_kv_pool_allocator + if ( + isinstance(allocator, UnifiedSWATokenToKVPoolAllocator) + and allocator.supports_asymmetric_reservation + ): + + def fits_shared_pool(candidate: int) -> bool: + full_tokens = paged_input_len + candidate + self.page_size + swa_tokens = estimate_swa_kv_tokens( + input_len, + candidate, + sliding_window_size=self.sliding_window_size, + page_size=self.page_size, + allocation_limit=self.chunked_prefill_size, + ) + return allocator.can_reserve( + full_tokens, + swa_tokens, + empty_pool=True, + require_token_slack=True, + ) + + if not fits_shared_pool(0): + req.set_finish_with_abort( + "Request prompt exceeds the unified FULL/SWA KV byte budget: " + f"input_len={input_len}." + ) + max_new_tokens = 0 + elif not fits_shared_pool(max_new_tokens): + lo, hi = 0, max_new_tokens + while lo < hi: + mid = (lo + hi + 1) // 2 + if fits_shared_pool(mid): + lo = mid + else: + hi = mid - 1 + max_new_tokens = lo + else: + # Non-unified pools keep their fixed token partition. + max_new_tokens = min( + max_new_tokens, self.max_total_num_tokens * get_parallel().attn_dcp_size - paged_input_len - self.page_size - 1, - ), - ) + ) + + req.sampling_params.max_new_tokens = max(0, max_new_tokens) # Clipping above can push max_new_tokens below min_new_tokens, which # would suppress EOS for the whole generation. Restore the invariant. if req.sampling_params.min_new_tokens > req.sampling_params.max_new_tokens: diff --git a/python/sglang/srt/managers/scheduler_components/invariant_checker.py b/python/sglang/srt/managers/scheduler_components/invariant_checker.py index a7340ca20346..d232b6d9f0fb 100644 --- a/python/sglang/srt/managers/scheduler_components/invariant_checker.py +++ b/python/sglang/srt/managers/scheduler_components/invariant_checker.py @@ -24,6 +24,7 @@ from sglang.srt.mem_cache.allocator.swa import is_swa_req_ring from sglang.srt.mem_cache.allocator.unified_hybrid_swa import ( UnifiedMambaSWATokenToKVPoolAllocator, + UnifiedSWATokenToKVPoolAllocator, ) from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache from sglang.srt.mem_cache.memory_pool import ReqToTokenPool @@ -94,12 +95,22 @@ def _check_pool_invariant( return leak, msg def _check_full_pool(self, ps: PoolStats, uncached: int = 0) -> Tuple[bool, str]: - if self.is_hybrid_swa and not self.full_tokens_per_layer: + allocator = self.token_to_kv_pool_allocator + is_unified_mamba_swa = isinstance( + allocator, UnifiedMambaSWATokenToKVPoolAllocator + ) + full_capacity = ( + allocator.current_full_capacity + if isinstance(allocator, UnifiedSWATokenToKVPoolAllocator) + and allocator.supports_asymmetric_reservation + else self.full_tokens_per_layer + ) + if self.is_hybrid_swa and not full_capacity: return False, "" if self.is_hybrid_swa: protected = self.tree_cache.full_protected_size() session_held = self.pool_stats_observer.session_held_full_tokens() - total = self.full_tokens_per_layer + total = full_capacity elif self.is_hybrid_ssm: # Branch on cache type for the protected accessor (MambaRadixCache # splits full/mamba; ChunkCache only has the single protected_size). @@ -119,7 +130,6 @@ def _check_full_pool(self, ps: PoolStats, uncached: int = 0) -> Tuple[bool, str] session_held = self.pool_stats_observer.session_held_tokens() total = self.max_total_num_tokens full_evictable_size = ps.full_evictable_size - allocator = self.token_to_kv_pool_allocator if get_parallel().dcp_enabled and allocator.page_size > 1: # DCP stores logical tokens in widened physical pages. Prefix cache # counters are logical-token based, while the allocator frees whole @@ -129,14 +139,14 @@ def _check_full_pool(self, ps: PoolStats, uncached: int = 0) -> Tuple[bool, str] // allocator.page_size * allocator.page_size ) - full_available = ps.full_available_size - if isinstance(allocator, UnifiedMambaSWATokenToKVPoolAllocator): - # Pair the static per-layer total with the conserve view, never the - # byte-coordinated one -- see `conserve_full_available_size`. - full_available = allocator.conserve_full_available_size() + full_available_size = ( + allocator.conserve_full_available_size() + if is_unified_mamba_swa + else ps.full_available_size + ) leak, msg = self._check_pool_invariant( "full", - full_available, + full_available_size, full_evictable_size, protected, session_held, @@ -162,18 +172,24 @@ def _check_swa_pool(self, ps: PoolStats, uncached: int = 0) -> Tuple[bool, str]: f"evictable={ps.swa_evictable_size}, " f"total={self.swa_tokens_per_layer}" ) - swa_available = ps.swa_available_size if isinstance(allocator, UnifiedMambaSWATokenToKVPoolAllocator): - # Tri-pool: same floating-boundary phantom as the full pool -- use the - # slot-conservation view, not the byte-coordinated min (see _check_full_pool). - swa_available = allocator.conserve_swa_available_size() + available = allocator.conserve_swa_available_size() + total = self.swa_tokens_per_layer + else: + available = ps.swa_available_size + total = ( + allocator.current_swa_capacity + if isinstance(allocator, UnifiedSWATokenToKVPoolAllocator) + and allocator.supports_asymmetric_reservation + else self.swa_tokens_per_layer + ) return self._check_pool_invariant( "swa", - swa_available, + available, ps.swa_evictable_size, self.tree_cache.swa_protected_size(), self.pool_stats_observer.session_held_swa_tokens(), - self.swa_tokens_per_layer, + total, uncached, ) diff --git a/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py b/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py index 9e1b962bac21..1b8577b3b94a 100644 --- a/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py +++ b/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py @@ -14,6 +14,7 @@ from sglang.srt.mem_cache.allocator.swa import is_swa_req_ring from sglang.srt.mem_cache.allocator.unified_hybrid_swa import ( UnifiedMambaSWATokenToKVPoolAllocator, + UnifiedSWATokenToKVPoolAllocator, ) if TYPE_CHECKING: @@ -289,12 +290,17 @@ def _get_mamba_token_info(self): ) def _get_swa_token_info(self) -> PoolStats: - # `*_num_used` is `static_cap - (available + evictable)`, so the - # available term must match the static cap's denomination: the conserve - # view, never the byte-coordinated one (see - # `conserve_full_available_size`). Measured ~25-90x inflated otherwise. allocator = self.token_to_kv_pool_allocator + full_capacity = self.full_tokens_per_layer + swa_capacity = self.swa_tokens_per_layer + if ( + isinstance(allocator, UnifiedSWATokenToKVPoolAllocator) + and allocator.supports_asymmetric_reservation + ): + full_capacity = allocator.current_full_capacity + swa_capacity = allocator.current_swa_capacity if isinstance(allocator, UnifiedMambaSWATokenToKVPoolAllocator): + # The tri-pool reports static capacities paired with conserve views. full_available_size = allocator.conserve_full_available_size() swa_available_size = allocator.conserve_swa_available_size() else: @@ -306,12 +312,8 @@ def _get_swa_token_info(self) -> PoolStats: # prefixes still report swa_evictable; counting it drives usage negative. if is_swa_req_ring(self.token_to_kv_pool_allocator): swa_evictable_size = 0 - full_num_used = self.full_tokens_per_layer - ( - full_available_size + full_evictable_size - ) - swa_num_used = self.swa_tokens_per_layer - ( - swa_available_size + swa_evictable_size - ) + full_num_used = full_capacity - (full_available_size + full_evictable_size) + swa_num_used = swa_capacity - (swa_available_size + swa_evictable_size) # FIXME(hisparse): host-backup transiently over-releases the device pool # counter, producing negative full_num_used / swa_num_used. We clamp to 0 # to keep token_usage / leak checks sane, but the underlying accounting @@ -319,13 +321,18 @@ def _get_swa_token_info(self) -> PoolStats: if self.enable_hisparse: full_num_used = max(0, full_num_used) swa_num_used = max(0, swa_num_used) - if not self.full_tokens_per_layer: + if not full_capacity: full_num_used = 0 full_available_size = 0 full_token_usage = 0.0 else: - full_token_usage = full_num_used / self.full_tokens_per_layer - swa_token_usage = swa_num_used / self.swa_tokens_per_layer + full_token_usage = full_num_used / full_capacity + if not swa_capacity: + swa_num_used = 0 + swa_available_size = 0 + swa_token_usage = 0.0 + else: + swa_token_usage = swa_num_used / swa_capacity return PoolStats( is_hybrid_swa=True, diff --git a/python/sglang/srt/mem_cache/allocator/base.py b/python/sglang/srt/mem_cache/allocator/base.py index 68ecb5e18cec..7e1b34db465e 100644 --- a/python/sglang/srt/mem_cache/allocator/base.py +++ b/python/sglang/srt/mem_cache/allocator/base.py @@ -97,6 +97,16 @@ def verify_byte_accounting(self) -> list: violation strings, empty when healthy. Static pools have no byte model.""" return [] + @property + def current_full_capacity(self) -> int: + """Capacity represented by full_available_size plus live allocations.""" + return self.size_full + + @property + def current_swa_capacity(self) -> int: + """SWA counterpart of current_full_capacity for hybrid allocators.""" + return self.size_swa + def mamba_full_cache_donor(self) -> MambaFullCacheDonor | None: """Return the shared-pool donor capability, if this allocator has one.""" return None diff --git a/python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py b/python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py index b29a09824207..b824ed062777 100644 --- a/python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py +++ b/python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py @@ -17,6 +17,7 @@ from __future__ import annotations import logging +import math from typing import Callable, List, Optional, Sequence, Tuple import torch @@ -49,6 +50,8 @@ class UnifiedSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator): `available_size()` (joint bytes, in TOKENS) is the only safe alloc pre-check. """ + supports_asymmetric_reservation = True + # Parent's `size` property has no setter but base init does `self.size = size`; # the no-op setter below absorbs that write. @property @@ -65,26 +68,36 @@ def __init__( unified_buffer: UnifiedKVPool, kvcache, # UnifiedSWAKVPool device: str, - full_max_total_num_tokens: int, - swa_max_total_num_tokens: int, + full_max_total_num_tokens: Optional[int] = None, + swa_max_total_num_tokens: Optional[int] = None, page_size: int = 1, need_sort: bool = False, forward_stream: Optional[torch.cuda.Stream] = None, lazy_compaction: bool = False, ): - # Set _size_full / _size_swa BEFORE base init (read during it). STATIC - # partition caps -- the slot-conservation value the leak invariant expects. - self._size_full = full_max_total_num_tokens - self._size_swa = swa_max_total_num_tokens - self._full_max_total_num_tokens = full_max_total_num_tokens - self._swa_max_total_num_tokens = swa_max_total_num_tokens + if (full_max_total_num_tokens is None) != (swa_max_total_num_tokens is None): + raise ValueError( + "full_max_total_num_tokens and swa_max_total_num_tokens must " + "either both be set or both be omitted" + ) + legacy_capacities = full_max_total_num_tokens is not None + self._size_full = ( + int(full_max_total_num_tokens) + if legacy_capacities + else unified_buffer.max_slots("full") - 1 + ) + self._size_swa = ( + int(swa_max_total_num_tokens) + if legacy_capacities + else unified_buffer.max_slots("swa") - 1 + ) self.page_size = page_size # The parent is inherited only for the isinstance contract: skip its # static-partition sub-pool allocation, which the unified pool replaces. BaseTokenToKVPoolAllocator.__init__( self, - size=full_max_total_num_tokens, + size=self._size_full, page_size=page_size, dtype=unified_buffer.mha_spec("full").store_dtype, device=device, @@ -120,6 +133,16 @@ def __init__( ) self._wire_peers() + self._empty_shared_gap_bytes = self.full_attn_allocator._current_gap_bytes() + if not legacy_capacities: + self._size_full = self.full_attn_allocator.available_size() + self._size_swa = min( + self.swa_attn_allocator.available_size(), + len(self.full_attn_allocator.free_virtual_ids) * page_size, + ) + self._full_max_total_num_tokens = self._size_full + self._swa_max_total_num_tokens = self._size_swa + # Epoch-keyed memo for the joint capacity view (any chain member's # mutation invalidates -- see `MultiEndedAllocator._chain_capacity_epoch`). self._joint_avail_memo_epoch: Optional[int] = None @@ -143,7 +166,7 @@ def __init__( "[unified-memory-pool] UnifiedSWATokenToKVPoolAllocator ready: " "full max_slots=%d (min_slot_index=%d, entry_bytes=%d), " "swa max_slots=%d (min_slot_index=%d, entry_bytes=%d), " - "static caps full=%d swa=%d, joint available=%d", + "max capacity full=%d swa=%d, joint available=%d", self.full_attn_allocator.max_slots, self.full_attn_allocator.min_slot_index, self.full_attn_allocator.entry_bytes, @@ -270,8 +293,250 @@ def _ask_float_for_room(self, need_tokens: int) -> None: """No float in a two-END chain -- nothing can slide.""" return None - # `size_full` / `size_swa` are inherited and read the static caps; reporting - # `max_slots - 1` here would be ~= full_max + swa_max and over-promise. + # `size_full` / `size_swa` bound each side independently; current capacities + # also account for the peer's live byte usage. + + @property + def current_full_capacity(self) -> int: + return self.full_available_size() + self.full_attn_allocator.allocated_count() + + @property + def current_swa_capacity(self) -> int: + return self.swa_available_size() + self.swa_attn_allocator.allocated_count() + + def reclaim_plan( + self, + full_tokens: int | float, + swa_tokens: int | float, + *, + full_evictable_tokens: int = 0, + swa_evictable_tokens: int = 0, + empty_pool: bool = False, + ) -> Optional[Tuple[int, int]]: + """Return cumulative FULL/SWA eviction targets, or None if impossible. + + The tree evicts FULL before SWA. Minimize required SWA reclaim with all + evictable FULL available, then trim excess FULL reclaim. FULL eviction's + actual SWA cascade is counted by the tree's shared eviction tracker. + """ + if full_tokens < 0 or swa_tokens < 0: + return None + + page_size = self.page_size + full_pages = (math.ceil(full_tokens) + page_size - 1) // page_size + swa_pages = (math.ceil(swa_tokens) + page_size - 1) // page_size + if swa_pages > full_pages: + return None + + compacted = empty_pool or not self.lazy_compaction or self._compaction_allowed() + + def fits(full_reclaim_pages: int, swa_reclaim_pages: int) -> bool: + return self._fits_page_demand( + full_pages, + swa_pages, + full_reclaim_pages=full_reclaim_pages, + swa_reclaim_pages=swa_reclaim_pages, + compacted=compacted, + empty_pool=empty_pool, + ) + + if empty_pool: + return (0, 0) if fits(0, 0) else None + if fits(0, 0): + return (0, 0) + + max_full_pages = min( + self.full_attn_allocator.allocated_count() // page_size, + max(0, int(full_evictable_tokens)) // page_size, + ) + max_swa_pages = min( + self.swa_attn_allocator.allocated_count() // page_size, + max(0, int(swa_evictable_tokens)) // page_size, + ) + if not fits(max_full_pages, max_swa_pages): + return None + + def first_fit(high: int, predicate: Callable[[int], bool]) -> int: + low = 0 + while low < high: + mid = (low + high) // 2 + if predicate(mid): + high = mid + else: + low = mid + 1 + return low + + swa_reclaim_pages = first_fit( + max_swa_pages, lambda value: fits(max_full_pages, value) + ) + full_reclaim_pages = first_fit( + max_full_pages, lambda value: fits(value, swa_reclaim_pages) + ) + return ( + full_reclaim_pages * page_size, + swa_reclaim_pages * page_size, + ) + + def can_reserve( + self, + full_tokens: int | float, + swa_tokens: int | float, + *, + full_evictable_tokens: int = 0, + swa_evictable_tokens: int = 0, + empty_pool: bool = False, + require_token_slack: bool = False, + ) -> bool: + """Check pending FULL/SWA demand against the shared byte envelope. + + Scheduler admission keeps the historical one-token strict slack at an + empty-pool boundary. Live admission checks the state reachable after + reclaiming the currently evictable pages from both sides. + """ + if full_tokens < 0 or swa_tokens < 0: + return False + if not self.supports_asymmetric_reservation: + if ( + full_tokens != swa_tokens + or full_evictable_tokens + or swa_evictable_tokens + or empty_pool + ): + return False + return full_tokens <= self.available_size() + if require_token_slack and ( + full_tokens >= self.size_full or swa_tokens >= self.size_swa + ): + return False + + return ( + self.reclaim_plan( + full_tokens, + swa_tokens, + full_evictable_tokens=full_evictable_tokens, + swa_evictable_tokens=swa_evictable_tokens, + empty_pool=empty_pool, + ) + is not None + ) + + def _compaction_allowed(self) -> bool: + return all( + allocator.disagg_move_gate is None or allocator.disagg_move_gate() + for allocator in (self.full_attn_allocator, self.swa_attn_allocator) + ) + + def _fits_page_demand( + self, + num_full_pages: int, + num_swa_pages: int, + *, + full_reclaim_pages: int = 0, + swa_reclaim_pages: int = 0, + compacted: bool, + empty_pool: bool = False, + ) -> bool: + """Check one FULL/SWA page demand against a single allocator snapshot.""" + if min(num_full_pages, num_swa_pages) < 0: + return False + + fa, sa = self.full_attn_allocator, self.swa_attn_allocator + if empty_pool: + full_live_pages = swa_live_pages = 0 + full_reclaim_pages = swa_reclaim_pages = 0 + else: + full_live_pages = fa.allocated_count() // self.page_size + swa_live_pages = sa.allocated_count() // self.page_size + full_reclaim_pages = min(full_live_pages, max(0, int(full_reclaim_pages))) + swa_reclaim_pages = min(swa_live_pages, max(0, int(swa_reclaim_pages))) + + full_live_pages -= full_reclaim_pages + swa_live_pages -= swa_reclaim_pages + full_total_pages = full_live_pages + num_full_pages + swa_total_pages = swa_live_pages + num_swa_pages + virtual_page_capacity = fa.num_virtual_ids - fa.min_page_index + full_page_capacity = min( + virtual_page_capacity, + fa.num_pages - fa.min_page_index, + ) + swa_page_capacity = min( + virtual_page_capacity, + sa.num_pages - sa.min_page_index, + ) + if full_total_pages > full_page_capacity or swa_total_pages > swa_page_capacity: + return False + + if compacted: + full_virtual_room = virtual_page_capacity - full_live_pages + full_holes = swa_holes = 0 + full_index_room = full_page_capacity - full_live_pages + swa_index_room = swa_page_capacity - swa_live_pages + gap_bytes = ( + self._empty_shared_gap_bytes + - full_live_pages * fa.entry_bytes_per_page + - swa_live_pages * sa.entry_bytes_per_page + ) + else: + full_virtual_room = len(fa.free_virtual_ids) + full_reclaim_pages + full_holes = len(fa._free_phys_pages) + full_reclaim_pages + swa_holes = len(sa._free_phys_pages) + swa_reclaim_pages + full_index_room = fa.num_pages - fa.min_page_index - fa._allocated_pages() + swa_index_room = sa.num_pages - sa.min_page_index - sa._allocated_pages() + gap_bytes = fa._current_gap_bytes() + + if num_full_pages > full_virtual_room: + return False + if num_full_pages > full_holes + full_index_room: + return False + if num_swa_pages > swa_holes + swa_index_room: + return False + full_extensions = max(0, num_full_pages - full_holes) + swa_extensions = max(0, num_swa_pages - swa_holes) + return ( + full_extensions * fa.entry_bytes_per_page + + swa_extensions * sa.entry_bytes_per_page + <= max(0, gap_bytes) + ) + + def ensure_capacity(self, full_tokens: int, swa_tokens: int) -> bool: + """Gate one allocation and compact both sides on shortfall.""" + if full_tokens < 0 or swa_tokens < 0: + return False + if full_tokens == 0 and swa_tokens == 0: + return True + if not self.supports_asymmetric_reservation: + if full_tokens != swa_tokens: + return False + need_tokens = int(full_tokens) + if need_tokens <= self.available_size(): + return True + return _relieve_for_alloc(self, need_tokens) + page_size = self.page_size + num_full_pages = (int(full_tokens) + page_size - 1) // page_size + num_swa_pages = (int(swa_tokens) + page_size - 1) // page_size + if num_swa_pages > num_full_pages: + return False + if self._fits_page_demand( + num_full_pages, + num_swa_pages, + compacted=False, + ): + return True + if not self.lazy_compaction or not self._compaction_allowed(): + return False + if not self._fits_page_demand( + num_full_pages, + num_swa_pages, + compacted=True, + ): + return False + self.full_attn_allocator._flush(urgent=True) + self.swa_attn_allocator._flush(urgent=True) + return self._fits_page_demand( + num_full_pages, + num_swa_pages, + compacted=False, + ) @property def draft_virtual_id_space(self) -> int: @@ -396,11 +661,8 @@ def swa_v2p_page_table(self) -> torch.Tensor: def alloc(self, need_size: int) -> Optional[torch.Tensor]: with record_function("UnifiedSWAAlloc.alloc"): - # Joint pre-check. Both sides are mutual peers (each side's compaction - # opens gap for the other), so flush BOTH on shortfall. - if need_size > self.available_size(): - if not _relieve_for_alloc(self, need_size): - return None + if not self.ensure_capacity(need_size, need_size): + return None # Snapshot the virtual PAGES full will consume, to bind them on swa too. num_pages = need_size // self.page_size fa = self.full_attn_allocator @@ -437,9 +699,8 @@ def _extend_in_virtual_space( prefix_lens=prefix_lens_cpu, ) need_tokens = num_new_pages * self.page_size - if need_tokens > self.available_size(): - if not _relieve_for_alloc(self, need_tokens): - return None + if not self.ensure_capacity(need_tokens, need_tokens): + return None # Snapshot the virtual PAGES the kernel will consume; clone so swa keeps # its view after the slice is consumed. @@ -555,9 +816,8 @@ def alloc_decode( seq_lens=seq_lens_cpu, page_size=self.page_size, decode=True ) need_tokens = num_new_pages * self.page_size - if need_tokens > self.available_size(): - if not _relieve_for_alloc(self, need_tokens): - return None + if not self.ensure_capacity(need_tokens, need_tokens): + return None fa = self.full_attn_allocator new_virtual_pages = fa.free_virtual_ids[:num_new_pages].clone() @@ -811,6 +1071,8 @@ class UnifiedMambaSWATokenToKVPoolAllocator(UnifiedSWATokenToKVPoolAllocator): `UnifiedMambaSlotAllocator`. """ + supports_asymmetric_reservation = False + def __init__( self, *, diff --git a/python/sglang/srt/mem_cache/allocator/unified_sub_pool.py b/python/sglang/srt/mem_cache/allocator/unified_sub_pool.py index 81f040476f1d..303df87a59cd 100644 --- a/python/sglang/srt/mem_cache/allocator/unified_sub_pool.py +++ b/python/sglang/srt/mem_cache/allocator/unified_sub_pool.py @@ -337,17 +337,27 @@ def __init__( self.page_size = page_size * dcp_size self.num_pages = max_slots // self.pool_page_size # `min_page_index` = ceil(min_slot_index / pool_page_size), keeping the - # reserved sink floor covered (see `_reserved_floor_bytes`). + # reserved-sink invariant (min_page_index * entry_bytes_per_page >= entry_max). + assert virtual_num_pages is None or not is_id_owner, ( + "only a non-owner allocator may use another pool's virtual-id space" + ) + # An ID owner has one virtual page per physical page. A non-owner may + # deliberately have a different physical capacity while indexing the + # owner's virtual page space (the unified SWA allocator does this). + self.num_virtual_ids = ( + virtual_num_pages if virtual_num_pages is not None else self.num_pages + ) + if is_id_owner: + assert self.num_virtual_ids == self.num_pages + assert self.num_virtual_ids > 0, "virtual page count must be positive" self.min_page_index = ( self.min_slot_index + self.pool_page_size - 1 ) // self.pool_page_size self.entry_bytes_per_page = self.entry_bytes * self.pool_page_size - # v2p is indexed by VIRTUAL page id, p2v by PHYSICAL page id. A non-owner - # consumes the owner's ids, so the two counts are unrelated. - self.num_virtual_ids = ( - self.num_pages if virtual_num_pages is None else virtual_num_pages - ) + # v2p is indexed by VIRTUAL page id, p2v by PHYSICAL page id. A + # non-owner consumes the owner's ids, so its v2p spans the owner's + # count; the two are unrelated and either can be the larger. # Page 0 is the padding anchor; the trailing row is the -1 sentinel. self.virtual_to_physical = torch.full( (self.num_virtual_ids + 1,), @@ -554,6 +564,7 @@ def allocator_state_str(self) -> str: f"is_id_owner={self.is_id_owner}, page_size={self.page_size}, " f"min_page_index={self.min_page_index}, " f"num_pages={self.num_pages}, " + f"num_virtual_ids={self.num_virtual_ids}, " f"watermark_physical={self.watermark_physical}, " f"allocated_pages={self._allocated_pages()}" ) diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index 9757fe57a238..289c6f19b4e3 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/python/sglang/srt/mem_cache/common.py @@ -161,15 +161,36 @@ def evict_from_tree_cache(tree_cache: BasePrefixCache | None, num_tokens: int): if isinstance(allocator, SWATokenToKVPoolAllocator): # Hybrid allocator - full_available_size = allocator.full_available_size() - swa_available_size = allocator.swa_available_size() + from sglang.srt.mem_cache.allocator.unified_hybrid_swa import ( + UnifiedSWATokenToKVPoolAllocator, + ) - 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) + is_unified_allocator = ( + isinstance(allocator, UnifiedSWATokenToKVPoolAllocator) + and allocator.supports_asymmetric_reservation + ) + if is_unified_allocator: + reclaim_plan = allocator.reclaim_plan( + num_tokens, + num_tokens, + full_evictable_tokens=tree_cache.full_evictable_size(), + swa_evictable_tokens=tree_cache.swa_evictable_size(), + ) + if reclaim_plan is None: + return + full_num_tokens, swa_num_tokens = reclaim_plan + else: + full_num_tokens = max(0, num_tokens - allocator.full_available_size()) + swa_num_tokens = max(0, num_tokens - allocator.swa_available_size()) + + if full_num_tokens or swa_num_tokens: tree_cache.evict_for_alloc( EvictParams(num_tokens=full_num_tokens, swa_num_tokens=swa_num_tokens) ) + if is_unified_allocator: + # The reclaim plan may rely on compacting cross-side holes. Make + # that projected state real before callers inspect availability. + allocator.ensure_capacity(num_tokens, num_tokens) else: # Standard allocator: evict only the shortfall (mirrors the SWA arm) available_size = allocator.available_size() diff --git a/python/sglang/srt/mem_cache/kv_cache_configurator.py b/python/sglang/srt/mem_cache/kv_cache_configurator.py index 7fedaaf26dac..9b0daf6af1e9 100644 --- a/python/sglang/srt/mem_cache/kv_cache_configurator.py +++ b/python/sglang/srt/mem_cache/kv_cache_configurator.py @@ -240,6 +240,7 @@ class _PoolSizes(msgspec.Struct, frozen=True, kw_only=True): c128_state_pool_size: int c4_state_dtype: Optional[torch.dtype] c128_state_dtype: Optional[torch.dtype] + unified_memory_pool_bytes: Optional[int] = None unified_total_bytes: Optional[int] = None @@ -423,6 +424,7 @@ def _derive_pool_sizes(self, *, config: MemoryPoolConfig) -> _PoolSizes: max_running_requests=max_running_requests, full_max_total_num_tokens=full_max_total_num_tokens, swa_max_total_num_tokens=swa_max_total_num_tokens, + unified_memory_pool_bytes=config.unified_memory_pool_bytes, c4_max_total_num_tokens=c4_max_total_num_tokens, c128_max_total_num_tokens=c128_max_total_num_tokens, c4_state_pool_size=c4_state_pool_size, @@ -469,6 +471,7 @@ def _init_pools( max_num_reqs=sizes.max_running_requests, full_max_total_num_tokens=sizes.full_max_total_num_tokens, swa_max_total_num_tokens=sizes.swa_max_total_num_tokens, + unified_memory_pool_bytes=sizes.unified_memory_pool_bytes, unified_total_bytes=sizes.unified_total_bytes, ) else: @@ -821,8 +824,9 @@ def _init_unified_swa_pools( self, *, max_num_reqs: int, - full_max_total_num_tokens: Optional[int], - swa_max_total_num_tokens: Optional[int], + full_max_total_num_tokens: Optional[int] = None, + swa_max_total_num_tokens: Optional[int] = None, + unified_memory_pool_bytes: Optional[int] = None, unified_total_bytes: Optional[int] = None, ) -> UnifiedPoolBundle: """Build the unified-pool stack for a hybrid-SWA model (Triton): one byte @@ -899,6 +903,10 @@ def _init_unified_swa_pools( if self.layer_info.start_layer <= i < self.layer_info.end_layer ] + total_bytes = unified_memory_pool_bytes + if total_bytes is None and self.spec_algorithm.is_none(): + total_bytes = unified_total_bytes + bundle = init_unified_swa_pools( device=self.device, kv_cache_dtype=self.kv_cache_dtype, @@ -915,6 +923,7 @@ def _init_unified_swa_pools( full_attention_layer_ids=full_attention_layer_ids, full_max_total_num_tokens=full_max_total_num_tokens, swa_max_total_num_tokens=swa_max_total_num_tokens, + total_bytes=total_bytes, enable_memory_saver=get_exec().features.enable_memory_saver, need_sort=get_disagg().disaggregation_mode in ("decode", "prefill"), # Overlap mode: same wait_stream(forward_stream) rationale as @@ -922,9 +931,6 @@ def _init_unified_swa_pools( forward_stream=self.forward_stream, # Lazy compaction: default ON, with env var escape hatch for rollback / A/B. lazy_compaction=_should_enable_lazy_compaction(), - # Draft workers keep the token-count byte sum (spec is asserted - # off under unified; belt only). - unified_total_bytes=(None if self.is_draft_worker else unified_total_bytes), # bs=1 feasibility floor inputs. `model_context_len` bounds the # sliding window term only -- the full-attention side is not # charged, see `_check_bs1_feasibility_floor`. diff --git a/python/sglang/srt/mem_cache/unified_memory_pool.py b/python/sglang/srt/mem_cache/unified_memory_pool.py index cdddc623153e..149edbe28017 100644 --- a/python/sglang/srt/mem_cache/unified_memory_pool.py +++ b/python/sglang/srt/mem_cache/unified_memory_pool.py @@ -1713,8 +1713,9 @@ def init_unified_swa_pools( end_layer: int, swa_attention_layer_ids: List[int], full_attention_layer_ids: List[int], - full_max_total_num_tokens: int, - swa_max_total_num_tokens: int, + full_max_total_num_tokens: Optional[int] = None, + swa_max_total_num_tokens: Optional[int] = None, + total_bytes: Optional[int] = None, enable_memory_saver: bool, need_sort: bool, forward_stream: Optional[torch.cuda.Stream] = None, @@ -1758,15 +1759,25 @@ def init_unified_swa_pools( store_dtype=store_dtype, grow_direction="up", ) - if unified_total_bytes is not None: - # PROFILED byte budget, sized from directly: the re-sum's floor losses - # stay out of the buffer, and the token counts remain boot labels. + if total_bytes is not None and unified_total_bytes is not None: + raise ValueError("total_bytes and unified_total_bytes are mutually exclusive") + if total_bytes is None: total_bytes = unified_total_bytes - else: + + legacy_allocator_capacities = {} + if total_bytes is None: + if full_max_total_num_tokens is None or swa_max_total_num_tokens is None: + raise ValueError( + "total_bytes or both legacy full/SWA capacities must be provided" + ) total_bytes = ( full_max_total_num_tokens * full_spec.entry_bytes() + swa_max_total_num_tokens * swa_spec.entry_bytes() ) + legacy_allocator_capacities = { + "full_max_total_num_tokens": full_max_total_num_tokens, + "swa_max_total_num_tokens": swa_max_total_num_tokens, + } if model_context_len is not None: # bs=1 floor: ONE sliding window of swa KV (+ a page of slack for the # page-granular walk) + the slot-0 sink. The full side is not charged @@ -1785,6 +1796,8 @@ def init_unified_swa_pools( ], factory="init_unified_swa_pools", ) + if total_bytes <= 0: + raise ValueError(f"total_bytes must be positive, got {total_bytes}") shared_pool = UnifiedKVPool( total_bytes=total_bytes, sub_pool_specs=[full_spec, swa_spec], @@ -1805,12 +1818,11 @@ def init_unified_swa_pools( unified_buffer=shared_pool, kvcache=token_to_kv_pool, device=device, - full_max_total_num_tokens=full_max_total_num_tokens, - swa_max_total_num_tokens=swa_max_total_num_tokens, page_size=page_size, need_sort=need_sort, forward_stream=forward_stream, lazy_compaction=lazy_compaction, + **legacy_allocator_capacities, ) logger.info( @@ -1841,12 +1853,12 @@ def init_unified_swa_pools( page_size, ) logger.info( - "[unified-memory-pool] total_bytes=%d (=%.2f GB), full_max_total_num_tokens=%d, " - "swa_max_total_num_tokens=%d, joint_available=%d slots", + "[unified-memory-pool] total_bytes=%d (=%.2f GB), " + "full_capacity=%d, swa_capacity=%d, joint_available=%d slots", total_bytes, total_bytes / GB, - full_max_total_num_tokens, - swa_max_total_num_tokens, + allocator.size_full, + allocator.size_swa, allocator.available_size(), ) logger.info( diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 14ad4fb4afe6..d20eb3e38ec4 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -1370,7 +1370,14 @@ def unload_lora_adapter(self, lora_ref: LoRARef): def effective_max_total_num_tokens(self): """Return the max token pool size considering hybrid swa settings.""" if self.is_hybrid_swa: - capacity = self.full_max_total_num_tokens or self.swa_max_total_num_tokens + if get_memory().enable_unified_memory: + capacity = self.token_to_kv_pool_allocator.size_full + if get_schedule().max_total_tokens is not None: + capacity = min(capacity, get_schedule().max_total_tokens) + else: + capacity = ( + self.full_max_total_num_tokens or self.swa_max_total_num_tokens + ) else: capacity = self.max_total_num_tokens if (req_to_token_pool := getattr(self, "req_to_token_pool", None)) is not None: diff --git a/python/sglang/srt/model_executor/pool_configurator.py b/python/sglang/srt/model_executor/pool_configurator.py index 8ec857a6ae53..132ecf358772 100644 --- a/python/sglang/srt/model_executor/pool_configurator.py +++ b/python/sglang/srt/model_executor/pool_configurator.py @@ -71,6 +71,7 @@ class MemoryPoolConfig: max_running_requests: Optional[int] = None full_max_total_num_tokens: Optional[int] = None swa_max_total_num_tokens: Optional[int] = None + unified_memory_pool_bytes: Optional[int] = None # DSV4 compressed-attention pool sizes (target only; draft workers leave at 0). c4_max_total_num_tokens: int = 0 @@ -698,6 +699,12 @@ def _draft_pool_bytes_per_token(self) -> int: + self._draft_cell_size ) + def _unified_pool_bytes(self, full_tokens: int, swa_tokens: int) -> int: + return ( + full_tokens * self._full_per_token * self._full_layers_num + + swa_tokens * self._swa_per_token * self._swa_layers_num + ) + def _max_unified_full_tokens( self, available_bytes: int, @@ -707,7 +714,6 @@ def _max_unified_full_tokens( """Find the largest page-aligned full capacity whose allocations fit.""" draft_bytes_per_token = self._draft_pool_bytes_per_token() target_full_bytes_per_token = self._full_per_token * self._full_layers_num - target_swa_bytes_per_token = self._swa_per_token * self._swa_layers_num assert target_full_bytes_per_token > 0 def allocation_bytes(full_pages: int) -> int: @@ -719,10 +725,7 @@ def allocation_bytes(full_pages: int) -> int: // page_size * page_size ) - target_bytes = ( - full_tokens * target_full_bytes_per_token - + swa_tokens * target_swa_bytes_per_token - ) + target_bytes = self._unified_pool_bytes(full_tokens, swa_tokens) virtual_span = max(target_bytes // target_full_bytes_per_token - 1, 0) draft_tokens = ceil_align(virtual_span, page_size) + page_size return target_bytes + draft_tokens * draft_bytes_per_token @@ -760,9 +763,10 @@ def align_page_size(x: int) -> int: full_tokens = align_page_size(max_total_num_tokens) swa_tokens = align_page_size(int(full_tokens * self._swa_full_tokens_ratio)) - self.validate_swa_pool_size( - swa_tokens, self._sliding_window_size, self._page_size - ) + if not self._enable_unified_memory: + self.validate_swa_pool_size( + swa_tokens, self._sliding_window_size, self._page_size + ) logger.info( f"Use sliding window memory pool. " @@ -773,16 +777,17 @@ def align_page_size(x: int) -> int: max_total_num_tokens=full_tokens, full_max_total_num_tokens=full_tokens, swa_max_total_num_tokens=swa_tokens, + unified_memory_pool_bytes=( + self._unified_pool_bytes(full_tokens, swa_tokens) + if self._enable_unified_memory + else None + ), ) def calculate_pool_sizes( self, available_bytes: int, page_size: int ) -> MemoryPoolConfig: - if ( - self._enable_unified_memory - and self._full_layers_num > 0 - and self._draft_pool_bytes_per_token() > 0 - ): + if self._enable_unified_memory and self._full_layers_num > 0: max_total_num_tokens = self._max_unified_full_tokens( available_bytes, page_size ) @@ -873,7 +878,7 @@ def calculate_pool_sizes( * self._swa_per_token * (self._swa_layers_num + self._draft_swa_layers_num) ) - if self._enable_unified_memory and self._draft_pool_bytes_per_token() > 0: + if self._enable_unified_memory: full_tokens = self._max_unified_full_tokens( available_bytes, page_size, fixed_swa_tokens=swa_tokens ) @@ -898,6 +903,11 @@ def calculate_pool_sizes( max_total_num_tokens=full_tokens, full_max_total_num_tokens=full_tokens, swa_max_total_num_tokens=swa_tokens, + unified_memory_pool_bytes=( + self._unified_pool_bytes(full_tokens, swa_tokens) + if self._enable_unified_memory + else None + ), ) def calculate_pool_sizes_from_max_tokens( @@ -910,6 +920,13 @@ def calculate_pool_sizes_from_max_tokens( max_total_num_tokens=full_tokens, full_max_total_num_tokens=full_tokens, swa_max_total_num_tokens=min(swa_tokens, max_total_num_tokens), + unified_memory_pool_bytes=( + self._unified_pool_bytes( + full_tokens, min(swa_tokens, max_total_num_tokens) + ) + if self._enable_unified_memory + else None + ), ) diff --git a/test/registered/unit/managers/test_scheduler_init_req_max_new_tokens.py b/test/registered/unit/managers/test_scheduler_init_req_max_new_tokens.py index 29b1402e507f..431dd85b744e 100644 --- a/test/registered/unit/managers/test_scheduler_init_req_max_new_tokens.py +++ b/test/registered/unit/managers/test_scheduler_init_req_max_new_tokens.py @@ -56,6 +56,7 @@ def _new_scheduler( scheduler.max_total_num_tokens = max_total_num_tokens scheduler.page_size = page_size scheduler.max_new_tokens_limit = envs.SGLANG_MAX_NEW_TOKENS_LIMIT.get() + scheduler.token_to_kv_pool_allocator = None return scheduler def _new_req(self, max_new_tokens, input_len: int = 8, min_new_tokens: int = 0): diff --git a/test/registered/unit/mem_cache/test_hisparse_max_token_pool_size.py b/test/registered/unit/mem_cache/test_hisparse_max_token_pool_size.py index 4a5a974834f3..dd4f1726d4aa 100644 --- a/test/registered/unit/mem_cache/test_hisparse_max_token_pool_size.py +++ b/test/registered/unit/mem_cache/test_hisparse_max_token_pool_size.py @@ -15,6 +15,7 @@ from sglang.srt.disaggregation.decode import DecodePreallocQueue from sglang.srt.model_executor.model_runner import ModelRunner +from sglang.srt.runtime_context import get_context from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -77,6 +78,9 @@ def test_non_hisparse_uses_effective_max_total_num_tokens(self): self.assertEqual(instance.max_token_pool_size, 1024) def test_non_hisparse_hybrid_swa_prefers_full_max(self): + self.enterContext( + get_context().override_server_args(enable_unified_memory=False) + ) instance = _make_model_runner( enable_hisparse=False, token_to_kv_pool_allocator=SimpleNamespace(), diff --git a/test/registered/unit/mem_cache/test_multi_ended_allocator.py b/test/registered/unit/mem_cache/test_multi_ended_allocator.py index 4c02d2d101f1..2baf5cceb07c 100644 --- a/test/registered/unit/mem_cache/test_multi_ended_allocator.py +++ b/test/registered/unit/mem_cache/test_multi_ended_allocator.py @@ -29,6 +29,7 @@ import torch +from sglang.srt.managers.schedule_policy import estimate_swa_kv_tokens from sglang.srt.mem_cache.allocator.unified_hybrid_swa import ( UnifiedSWATokenToKVPoolAllocator, ) @@ -558,6 +559,7 @@ def _build( swa_layer_num=2, head_num=2, head_dim=4, + page_size=1, ): full_spec = MHASubPoolSpec( name="full", @@ -584,6 +586,7 @@ def _build( sub_pool_specs=[full_spec, swa_spec], device=_DEV, enable_memory_saver=False, + page_size=page_size, ) kvcache = _FakeUnifiedSWAKVPool(pool) allocator = UnifiedSWATokenToKVPoolAllocator( @@ -592,11 +595,104 @@ def _build( device=_DEV, full_max_total_num_tokens=n_full_slots, swa_max_total_num_tokens=n_swa_slots, + page_size=page_size, need_sort=False, forward_stream=None, ) return pool, allocator, kvcache + def test_empty_pool_reservation_matches_packed_byte_boundary(self): + page_size = 4 + _, allocator, _ = self._build( + n_full_slots=40, + n_swa_slots=24, + full_layer_num=4, + swa_layer_num=2, + page_size=page_size, + ) + full_allocator = allocator.full_attn_allocator + swa_allocator = allocator.swa_attn_allocator + swa_pages = 2 + full_pages = ( + allocator._empty_shared_gap_bytes + - swa_pages * swa_allocator.entry_bytes_per_page + ) // full_allocator.entry_bytes_per_page + packed_bytes = ( + full_pages * full_allocator.entry_bytes_per_page + + swa_pages * swa_allocator.entry_bytes_per_page + ) + + self.assertLessEqual( + full_pages + 1, full_allocator.num_pages - full_allocator.min_page_index + ) + self.assertLessEqual( + swa_pages, swa_allocator.num_pages - swa_allocator.min_page_index + ) + self.assertLessEqual(packed_bytes, allocator._empty_shared_gap_bytes) + self.assertGreater( + packed_bytes + full_allocator.entry_bytes_per_page, + allocator._empty_shared_gap_bytes, + ) + self.assertTrue( + allocator.can_reserve( + full_pages * page_size, + swa_pages * page_size, + empty_pool=True, + ) + ) + self.assertFalse( + allocator.can_reserve( + full_pages * page_size + 1, + swa_pages * page_size, + empty_pool=True, + ) + ) + + extend_tokens = 32 + max_new_tokens = 0 + reservation_full_tokens = extend_tokens + max_new_tokens + page_size + reservation_swa_tokens = estimate_swa_kv_tokens( + extend_tokens, + max_new_tokens, + sliding_window_size=16, + page_size=page_size, + include_allocated_tail=False, + ) + reservation_swa_with_tail = estimate_swa_kv_tokens( + extend_tokens, + max_new_tokens, + sliding_window_size=16, + page_size=page_size, + ) + reservation_bytes = ( + reservation_full_tokens // page_size + ) * full_allocator.entry_bytes_per_page + ( + reservation_swa_tokens // page_size + ) * swa_allocator.entry_bytes_per_page + reservation_bytes_with_tail = ( + reservation_full_tokens // page_size + ) * full_allocator.entry_bytes_per_page + ( + reservation_swa_with_tail // page_size + ) * swa_allocator.entry_bytes_per_page + self.assertLessEqual(reservation_bytes, allocator._empty_shared_gap_bytes) + self.assertGreater( + reservation_bytes_with_tail, allocator._empty_shared_gap_bytes + ) + self.assertTrue( + allocator.can_reserve( + reservation_full_tokens, + reservation_swa_tokens, + empty_pool=True, + ) + ) + self.assertFalse( + allocator.can_reserve( + reservation_full_tokens, + reservation_swa_with_tail, + empty_pool=True, + ) + ) + def _alloc(self, allocator, kvcache, n): """Allocate N virtual ids; stamp the data marker on both sub-pools.""" v = allocator.alloc(n) diff --git a/test/registered/unit/model_executor/test_pool_configurator.py b/test/registered/unit/model_executor/test_pool_configurator.py index 5a6b532925da..13963e6a6861 100644 --- a/test/registered/unit/model_executor/test_pool_configurator.py +++ b/test/registered/unit/model_executor/test_pool_configurator.py @@ -83,6 +83,7 @@ def _make_model_runner( disaggregation_mode="null", max_running_requests=None, disaggregation_decode_extra_slots=0, + enable_unified_memory=False, kv_lora_rank=512, qk_rope_head_dim=64, swa_kv_lora_rank=128, @@ -150,6 +151,7 @@ def _make_model_runner( disaggregation_mode=disaggregation_mode, max_running_requests=max_running_requests, disaggregation_decode_extra_slots=disaggregation_decode_extra_slots, + enable_unified_memory=enable_unified_memory, enable_hisparse=False, enable_hierarchical_cache=False, enable_dsa_cache_layer_split=False, @@ -299,7 +301,14 @@ def test_dsa_mla_cell_size_uses_backend_kv_layout( class TestHybridSWAConfigurator(CustomTestCase): """Hybrid SWA: full/swa split, ratio, memory invariant.""" - def _make_swa_runner(self, full_layers=16, swa_layers=16, ratio=0.5, page_size=1): + def _make_swa_runner( + self, + full_layers=16, + swa_layers=16, + ratio=0.5, + page_size=1, + enable_unified_memory=False, + ): return _make_model_runner( self, is_hybrid_swa=True, @@ -308,6 +317,7 @@ def _make_swa_runner(self, full_layers=16, swa_layers=16, ratio=0.5, page_size=1 swa_num_kv_heads=4, page_size=page_size, swa_full_tokens_ratio=ratio, + enable_unified_memory=enable_unified_memory, ) def _run(self, available_bytes, **kwargs): @@ -329,6 +339,68 @@ def test_memory_utilization(self): self.assertLessEqual(used, available) self.assertGreater(used, available * 0.99) + def test_unified_capacity_is_maximal_with_draft_pool(self): + page_size = 8 + full_layers = 2 + swa_layers = 1 + draft_layers = 2 + draft_swa_layers = 1 + ratio = 0.5 + mr = _make_model_runner( + self, + is_hybrid_swa=True, + full_attention_layer_ids=list(range(full_layers)), + swa_attention_layer_ids=list(range(full_layers, full_layers + swa_layers)), + swa_num_kv_heads=4, + swa_full_tokens_ratio=ratio, + page_size=page_size, + enable_unified_memory=True, + speculative_algorithm="EAGLE", + ) + mr.spec_algorithm.is_eagle.return_value = True + mr.spec_algorithm.is_none.return_value = False + mr.spec_aux_config.eagle_draft_num_layers = draft_layers + mr.spec_aux_config.eagle_draft_swa_num_layers = draft_swa_layers + + full_bytes_per_token = _full_per_token(mr) + swa_bytes_per_token = _swa_per_token(mr) + target_full_bytes_per_token = full_bytes_per_token * full_layers + draft_bytes_per_token = ( + full_bytes_per_token * (draft_layers - draft_swa_layers) + + swa_bytes_per_token * draft_swa_layers + ) + + def allocation_bytes(full_tokens, *, include_reserved_draft_page=True): + swa_tokens = int(full_tokens * ratio) // page_size * page_size + target_bytes = ( + full_tokens * target_full_bytes_per_token + + swa_tokens * swa_bytes_per_token * swa_layers + ) + virtual_span = max(target_bytes // target_full_bytes_per_token - 1, 0) + draft_tokens = (virtual_span + page_size - 1) // page_size * page_size + if include_reserved_draft_page: + draft_tokens += page_size + return target_bytes + draft_tokens * draft_bytes_per_token + + expected_full_tokens = 10 * page_size + available = allocation_bytes( + expected_full_tokens + page_size, + include_reserved_draft_page=False, + ) + with mock_cpu_env(): + from sglang.srt.model_executor.pool_configurator import ( + create_memory_pool_configurator, + ) + + cfg = create_memory_pool_configurator(mr) + config = cfg.calculate_pool_sizes(available, page_size) + + full_tokens = config.full_max_total_num_tokens + self.assertEqual(full_tokens % page_size, 0) + self.assertEqual(full_tokens, expected_full_tokens) + self.assertLessEqual(allocation_bytes(full_tokens), available) + self.assertGreater(allocation_bytes(full_tokens + page_size), available) + @patch( "sglang.srt.mem_cache.kv_cache_configurator.calculate_mla_kv_cache_dim", return_value=576, From a7084b373d9c847caad3037bc2426aedbc5586c4 Mon Sep 17 00:00:00 2001 From: yhzhuang Date: Thu, 10 Sep 2026 15:47:12 -0700 Subject: [PATCH 41/55] Keep HiSparse fixture compatible with Python 3.10 --- .../unit/mem_cache/test_hisparse_max_token_pool_size.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/test/registered/unit/mem_cache/test_hisparse_max_token_pool_size.py b/test/registered/unit/mem_cache/test_hisparse_max_token_pool_size.py index dd4f1726d4aa..4c590ce02bf1 100644 --- a/test/registered/unit/mem_cache/test_hisparse_max_token_pool_size.py +++ b/test/registered/unit/mem_cache/test_hisparse_max_token_pool_size.py @@ -78,9 +78,6 @@ def test_non_hisparse_uses_effective_max_total_num_tokens(self): self.assertEqual(instance.max_token_pool_size, 1024) def test_non_hisparse_hybrid_swa_prefers_full_max(self): - self.enterContext( - get_context().override_server_args(enable_unified_memory=False) - ) instance = _make_model_runner( enable_hisparse=False, token_to_kv_pool_allocator=SimpleNamespace(), @@ -89,8 +86,9 @@ def test_non_hisparse_hybrid_swa_prefers_full_max(self): full_max_total_num_tokens=3000, swa_max_total_num_tokens=500, ) - self.assertEqual(instance.max_token_pool_size, 3000) - self.assertEqual(instance.effective_max_total_num_tokens, 3000) + with get_context().override_server_args(enable_unified_memory=False): + self.assertEqual(instance.max_token_pool_size, 3000) + self.assertEqual(instance.effective_max_total_num_tokens, 3000) def _make_prealloc_queue( From b505c110976c9ff0ce1305ca016cbced749012ae Mon Sep 17 00:00:00 2001 From: yhzhuang Date: Thu, 10 Sep 2026 16:04:23 -0700 Subject: [PATCH 42/55] Remove unrelated UMP diff churn --- python/sglang/srt/arg_groups/kv_cache_hook.py | 9 +++-- python/sglang/srt/disaggregation/decode.py | 2 + python/sglang/srt/managers/schedule_policy.py | 6 +-- .../scheduler_components/invariant_checker.py | 38 +++++++++---------- python/sglang/srt/mem_cache/allocator/base.py | 10 ----- .../mem_cache/allocator/unified_sub_pool.py | 24 +++++------- 6 files changed, 36 insertions(+), 53 deletions(-) diff --git a/python/sglang/srt/arg_groups/kv_cache_hook.py b/python/sglang/srt/arg_groups/kv_cache_hook.py index b5ba4eac39c4..4d473c79ef23 100644 --- a/python/sglang/srt/arg_groups/kv_cache_hook.py +++ b/python/sglang/srt/arg_groups/kv_cache_hook.py @@ -200,10 +200,6 @@ def handle_unified_memory_pool(server_args: Any) -> None: if not cfg.enable_unified_memory: return if cfg.disaggregation_mode != "null": - assert cfg.pp_size == 1, ( - "--enable-unified-memory with PD disaggregation does not support " - "pipeline parallelism (--pp-size > 1)." - ) # Constraints of the whole-envelope transfer; see # UnifiedMLATokenToKVPool.get_contiguous_buf_infos. assert cfg.disaggregation_transfer_backend == "mooncake", ( @@ -211,6 +207,11 @@ def handle_unified_memory_pool(server_args: Any) -> None: "the mooncake transfer backend; got " f"{cfg.disaggregation_transfer_backend!r}." ) + assert cfg.pp_size == 1, ( + "--enable-unified-memory with PD disaggregation does not support " + "pipeline parallelism (whole-envelope transfer has no per-layer " + "entries to subset)." + ) assert not envs.SGLANG_DISABLE_LAZY_COMPACTION.get(), ( "--enable-unified-memory with PD disaggregation requires lazy " "compaction; unset SGLANG_DISABLE_LAZY_COMPACTION." diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 98ce8dadac44..c1cf16029d7a 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -454,6 +454,7 @@ def _reclaim_swa_tail_capacity( EvictParams(swa_num_tokens=required - available) ) available = self.token_to_kv_pool_allocator.swa_available_size() + if available < required: return ( f"SWA eviction insufficient: needed={required}, " @@ -855,6 +856,7 @@ def resume_retracted_reqs( break if uses_swa_tail_prealloc and swa_required > swa_allocatable_tokens: break + resumed_reqs.append(req) indices_to_remove.add(i) req.is_retracted = False diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index 85810debb4ac..a0da8bfbc70c 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -902,8 +902,7 @@ def _swa_chunk_cap( ) if cap <= 0: return 0 - cap = cap // self.page_size * self.page_size - return cap + return cap // self.page_size * self.page_size assert max_chunk_tokens is not None cap = max_chunk_tokens // self.page_size * self.page_size @@ -1481,6 +1480,7 @@ def add_one_req( # this returns 0, so the debit sites below reuse the value. mamba_gap_reserve = self._mamba_gap_budget_for_req(req) total_tokens += mamba_gap_reserve + # adjusting the input_tokens based on host_hit_length and page_size swa_extend_input_len = cand_extend_input_len - req.host_hit_length real_input_tokens = self.ceil_paged_tokens(swa_extend_input_len) @@ -1510,7 +1510,7 @@ def add_one_req( return AddReqResult.OTHER with self._lock_node(req.last_node): - # Capacity may decrease after the lock acquisition. + # self.rem_total_tokens may decrease after the lock acquisition if not self.is_unified_swa and total_tokens >= self.rem_total_tokens: return AddReqResult.NO_TOKEN diff --git a/python/sglang/srt/managers/scheduler_components/invariant_checker.py b/python/sglang/srt/managers/scheduler_components/invariant_checker.py index d232b6d9f0fb..a3a95d0128d8 100644 --- a/python/sglang/srt/managers/scheduler_components/invariant_checker.py +++ b/python/sglang/srt/managers/scheduler_components/invariant_checker.py @@ -96,9 +96,6 @@ def _check_pool_invariant( def _check_full_pool(self, ps: PoolStats, uncached: int = 0) -> Tuple[bool, str]: allocator = self.token_to_kv_pool_allocator - is_unified_mamba_swa = isinstance( - allocator, UnifiedMambaSWATokenToKVPoolAllocator - ) full_capacity = ( allocator.current_full_capacity if isinstance(allocator, UnifiedSWATokenToKVPoolAllocator) @@ -139,14 +136,14 @@ def _check_full_pool(self, ps: PoolStats, uncached: int = 0) -> Tuple[bool, str] // allocator.page_size * allocator.page_size ) - full_available_size = ( - allocator.conserve_full_available_size() - if is_unified_mamba_swa - else ps.full_available_size - ) + full_available = ps.full_available_size + if isinstance(allocator, UnifiedMambaSWATokenToKVPoolAllocator): + # Pair the static per-layer total with the conserve view, never the + # byte-coordinated one -- see `conserve_full_available_size`. + full_available = allocator.conserve_full_available_size() leak, msg = self._check_pool_invariant( "full", - full_available_size, + full_available, full_evictable_size, protected, session_held, @@ -172,24 +169,23 @@ def _check_swa_pool(self, ps: PoolStats, uncached: int = 0) -> Tuple[bool, str]: f"evictable={ps.swa_evictable_size}, " f"total={self.swa_tokens_per_layer}" ) + swa_available = ps.swa_available_size if isinstance(allocator, UnifiedMambaSWATokenToKVPoolAllocator): - available = allocator.conserve_swa_available_size() - total = self.swa_tokens_per_layer - else: - available = ps.swa_available_size - total = ( - allocator.current_swa_capacity - if isinstance(allocator, UnifiedSWATokenToKVPoolAllocator) - and allocator.supports_asymmetric_reservation - else self.swa_tokens_per_layer - ) + # Tri-pool: same floating-boundary phantom as the full pool -- use the + # slot-conservation view, not the byte-coordinated min (see _check_full_pool). + swa_available = allocator.conserve_swa_available_size() return self._check_pool_invariant( "swa", - available, + swa_available, ps.swa_evictable_size, self.tree_cache.swa_protected_size(), self.pool_stats_observer.session_held_swa_tokens(), - total, + ( + allocator.current_swa_capacity + if isinstance(allocator, UnifiedSWATokenToKVPoolAllocator) + and allocator.supports_asymmetric_reservation + else self.swa_tokens_per_layer + ), uncached, ) diff --git a/python/sglang/srt/mem_cache/allocator/base.py b/python/sglang/srt/mem_cache/allocator/base.py index 7e1b34db465e..68ecb5e18cec 100644 --- a/python/sglang/srt/mem_cache/allocator/base.py +++ b/python/sglang/srt/mem_cache/allocator/base.py @@ -97,16 +97,6 @@ def verify_byte_accounting(self) -> list: violation strings, empty when healthy. Static pools have no byte model.""" return [] - @property - def current_full_capacity(self) -> int: - """Capacity represented by full_available_size plus live allocations.""" - return self.size_full - - @property - def current_swa_capacity(self) -> int: - """SWA counterpart of current_full_capacity for hybrid allocators.""" - return self.size_swa - def mamba_full_cache_donor(self) -> MambaFullCacheDonor | None: """Return the shared-pool donor capability, if this allocator has one.""" return None diff --git a/python/sglang/srt/mem_cache/allocator/unified_sub_pool.py b/python/sglang/srt/mem_cache/allocator/unified_sub_pool.py index 303df87a59cd..99ebe16ba414 100644 --- a/python/sglang/srt/mem_cache/allocator/unified_sub_pool.py +++ b/python/sglang/srt/mem_cache/allocator/unified_sub_pool.py @@ -337,27 +337,21 @@ def __init__( self.page_size = page_size * dcp_size self.num_pages = max_slots // self.pool_page_size # `min_page_index` = ceil(min_slot_index / pool_page_size), keeping the - # reserved-sink invariant (min_page_index * entry_bytes_per_page >= entry_max). + # reserved sink floor covered (see `_reserved_floor_bytes`). + self.min_page_index = ( + self.min_slot_index + self.pool_page_size - 1 + ) // self.pool_page_size + self.entry_bytes_per_page = self.entry_bytes * self.pool_page_size + + # v2p is indexed by VIRTUAL page id, p2v by PHYSICAL page id. A non-owner + # consumes the owner's ids, so the two counts are unrelated. assert virtual_num_pages is None or not is_id_owner, ( "only a non-owner allocator may use another pool's virtual-id space" ) - # An ID owner has one virtual page per physical page. A non-owner may - # deliberately have a different physical capacity while indexing the - # owner's virtual page space (the unified SWA allocator does this). self.num_virtual_ids = ( - virtual_num_pages if virtual_num_pages is not None else self.num_pages + self.num_pages if virtual_num_pages is None else virtual_num_pages ) - if is_id_owner: - assert self.num_virtual_ids == self.num_pages assert self.num_virtual_ids > 0, "virtual page count must be positive" - self.min_page_index = ( - self.min_slot_index + self.pool_page_size - 1 - ) // self.pool_page_size - self.entry_bytes_per_page = self.entry_bytes * self.pool_page_size - - # v2p is indexed by VIRTUAL page id, p2v by PHYSICAL page id. A - # non-owner consumes the owner's ids, so its v2p spans the owner's - # count; the two are unrelated and either can be the larger. # Page 0 is the padding anchor; the trailing row is the -1 sentinel. self.virtual_to_physical = torch.full( (self.num_virtual_ids + 1,), From 99ee7808c4dac4a37b40cae6af257d0c46d12fc1 Mon Sep 17 00:00:00 2001 From: yhzhuang Date: Sat, 12 Sep 2026 09:10:24 -0700 Subject: [PATCH 43/55] Fix unified pool admission rounding --- python/sglang/srt/managers/schedule_policy.py | 6 +-- python/sglang/srt/managers/scheduler.py | 2 +- .../test_scheduler_init_req_max_new_tokens.py | 34 ++++++++++++ .../mem_cache/test_multi_ended_allocator.py | 52 +++++++++++++++++++ 4 files changed, 88 insertions(+), 6 deletions(-) diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index a0da8bfbc70c..0270e00ea1ef 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -981,11 +981,7 @@ def _check_swa_admission( full_ever_tokens = unified_swa_tokens = None if self.is_unified_swa: full_input_tokens = len(req.full_untruncated_fill_ids) - full_ever_tokens = ( - self.ceil_paged_tokens(full_input_tokens) - + max_new_tokens - + self.page_size - ) + full_ever_tokens = full_input_tokens + max_new_tokens + self.page_size unified_swa_tokens = self._swa_budget_for_req( full_input_tokens, max_new_tokens ) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 4e849d6b98e4..a40632973471 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -2539,7 +2539,7 @@ def init_req_max_new_tokens(self, req): ): def fits_shared_pool(candidate: int) -> bool: - full_tokens = paged_input_len + candidate + self.page_size + full_tokens = input_len + candidate + self.page_size swa_tokens = estimate_swa_kv_tokens( input_len, candidate, diff --git a/test/registered/unit/managers/test_scheduler_init_req_max_new_tokens.py b/test/registered/unit/managers/test_scheduler_init_req_max_new_tokens.py index 431dd85b744e..43f182778e26 100644 --- a/test/registered/unit/managers/test_scheduler_init_req_max_new_tokens.py +++ b/test/registered/unit/managers/test_scheduler_init_req_max_new_tokens.py @@ -2,8 +2,11 @@ import unittest from types import SimpleNamespace +import torch + from sglang.srt.environ import envs from sglang.srt.managers.scheduler import Scheduler +from sglang.srt.mem_cache.unified_memory_pool import init_unified_swa_pools from sglang.srt.runtime_context import get_parallel from sglang.test.ci.ci_register import register_cpu_ci @@ -181,6 +184,37 @@ def test_admission_rules_sweep(self): ) self._init_and_check(scheduler, req) + def test_unified_budget_rounds_prompt_and_decode_together(self): + bundle = init_unified_swa_pools( + device="cpu", + kv_cache_dtype=torch.float16, + head_num=1, + head_dim=4, + v_head_dim=4, + swa_head_num=1, + swa_head_dim=4, + swa_v_head_dim=4, + page_size=4, + start_layer=0, + end_layer=2, + swa_attention_layer_ids=[1], + full_attention_layer_ids=[0], + total_bytes=384, + enable_memory_saver=False, + need_sort=False, + lazy_compaction=True, + ) + scheduler = self._new_scheduler(page_size=4) + scheduler.token_to_kv_pool_allocator = bundle.token_to_kv_pool_allocator + scheduler.sliding_window_size = 4 + scheduler.chunked_prefill_size = 4 + scheduler.max_new_tokens_limit = None + for prompt_len in (4, 5, 6, 7): + with self.subTest(prompt_len=prompt_len): + req = self._new_req(max_new_tokens=1, input_len=prompt_len) + scheduler.init_req_max_new_tokens(req) + self.assertEqual(req.sampling_params.max_new_tokens, 1) + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/mem_cache/test_multi_ended_allocator.py b/test/registered/unit/mem_cache/test_multi_ended_allocator.py index 2baf5cceb07c..bac0687f9f1d 100644 --- a/test/registered/unit/mem_cache/test_multi_ended_allocator.py +++ b/test/registered/unit/mem_cache/test_multi_ended_allocator.py @@ -601,6 +601,58 @@ def _build( ) return pool, allocator, kvcache + def test_reclaim_plan_matches_exhaustive_page_targets(self): + page_size = 4 + _, allocator, _ = self._build( + n_full_slots=40, n_swa_slots=24, page_size=page_size + ) + allocator.lazy_compaction = True + for sub_pool in (allocator.full_attn_allocator, allocator.swa_attn_allocator): + sub_pool.lazy_compaction = True + sub_pool.disagg_move_gate = lambda: False + live = allocator.alloc(16) + self.assertIsNotNone(live) + allocator.free(live[4:8]) + allocator.free_swa(live[8:12]) + + for compacted in (False, True): + for sub_pool in ( + allocator.full_attn_allocator, + allocator.swa_attn_allocator, + ): + sub_pool.disagg_move_gate = lambda: compacted + for full_evictable, swa_evictable in ((0, 0), (7, 5), (12, 8), (100, 100)): + max_full = min(12, full_evictable) // page_size + max_swa = min(8, swa_evictable) // page_size + for full_pages in range(9): + for swa_pages in range(full_pages + 1): + feasible = [ + (full * page_size, swa * page_size) + for swa in range(max_swa + 1) + for full in range(max_full + 1) + if allocator._fits_page_demand( + full_pages, + swa_pages, + full_reclaim_pages=full, + swa_reclaim_pages=swa, + compacted=compacted, + ) + ] + with self.subTest( + compacted=compacted, + evictable=(full_evictable, swa_evictable), + pages=(full_pages, swa_pages), + ): + self.assertEqual( + allocator.reclaim_plan( + full_pages * page_size, + swa_pages * page_size, + full_evictable_tokens=full_evictable, + swa_evictable_tokens=swa_evictable, + ), + feasible[0] if feasible else None, + ) + def test_empty_pool_reservation_matches_packed_byte_boundary(self): page_size = 4 _, allocator, _ = self._build( From 8fd4aa44a2a0c9f17fffe9c619187a0c72a2709d Mon Sep 17 00:00:00 2001 From: yhzhuang Date: Sat, 12 Sep 2026 21:38:57 -0700 Subject: [PATCH 44/55] Refactor unified SWA budget handling --- python/sglang/srt/managers/schedule_policy.py | 258 ++++++++---------- python/sglang/srt/managers/scheduler.py | 78 +++--- .../scheduler_components/invariant_checker.py | 35 +-- .../pool_stats_observer.py | 15 +- .../mem_cache/allocator/unified_hybrid_swa.py | 61 +++-- python/sglang/srt/mem_cache/common.py | 27 +- .../srt/model_executor/pool_configurator.py | 27 +- 7 files changed, 237 insertions(+), 264 deletions(-) diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index 0270e00ea1ef..54bece2ae4c5 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -54,7 +54,7 @@ ) from sglang.srt.mem_cache.allocator.unified_hybrid_swa import ( UnifiedMambaSWATokenToKVPoolAllocator, - UnifiedSWATokenToKVPoolAllocator, + supports_swa_byte_budget, ) from sglang.srt.mem_cache.allocator.unified_mamba import ( UnifiedMambaTokenToKVPoolAllocator, @@ -639,12 +639,7 @@ def __init__( self.token_to_kv_pool_allocator, PureSWATokenToKVPoolAllocator ) self.is_hybrid_ssm_cache = self.tree_cache.supports_mamba() - self.is_unified_swa = ( - isinstance( - self.token_to_kv_pool_allocator, UnifiedSWATokenToKVPoolAllocator - ) - and self.token_to_kv_pool_allocator.supports_asymmetric_reservation - ) + self.is_unified_swa = supports_swa_byte_budget(self.token_to_kv_pool_allocator) self.rem_swa_token_offset = ( num_mixed_decode_tokens if self.is_unified_swa else 0 @@ -855,8 +850,6 @@ def _unified_swa_reservation_fits( *, empty_pool: bool = False, ) -> bool: - if not self.is_unified_swa: - return True allocator = self.token_to_kv_pool_allocator if empty_pool: full_evictable_tokens = 0 @@ -875,36 +868,14 @@ def _unified_swa_reservation_fits( require_token_slack=empty_pool, ) - def _swa_chunk_cap( + def _unified_swa_chunk_cap( self, max_new_tokens: int, swa_host_hit_length: int = 0, *, - max_chunk_tokens: Optional[int] = None, + max_chunk_tokens: int, ) -> int: - """Largest page-aligned extend chunk the SWA pool can admit right now. - - Split pools reserve headroom below rem_swa_tokens; unified pools use - the joint byte budget. Returns 0 if not even one page fits. Only valid - when is_hybrid_swa is True. - - Escape hatch for a request whose budget can never pass the - _swa_budget_for_req gate (extend near/above the pool size, or a large - load-back charge): without shrinking its chunk it would be rejected - forever (head-of-line livelock). Shrinking is sound because past a - chunk boundary only the sliding window stays locked — the rest turns - evictable — so each pass's transient footprint fits the pool.""" - if not self.is_unified_swa: - # extend_input_len=0: this solves for the extend chunk itself, so - # the reserved headroom is the post-chunk decode window only. - cap = int(self.rem_swa_tokens) - self._swa_reserved_tokens( - 0, max_new_tokens, swa_host_hit_length - ) - if cap <= 0: - return 0 - return cap // self.page_size * self.page_size - - assert max_chunk_tokens is not None + """Largest page-aligned chunk that fits the joint FULL/SWA byte budget.""" cap = max_chunk_tokens // self.page_size * self.page_size lo, hi = 0, cap // self.page_size while lo < hi: @@ -920,14 +891,18 @@ def _swa_chunk_cap( hi = mid - 1 return lo * self.page_size + def _swa_chunk_cap(self, max_new_tokens: int, swa_host_hit_length: int = 0) -> int: + """Largest chunk fitting the fixed SWA partition after decode headroom.""" + cap = int(self.rem_swa_tokens) - self._swa_reserved_tokens( + 0, max_new_tokens, swa_host_hit_length + ) + return max(0, cap // self.page_size * self.page_size) + def _swa_req_never_fits( self, extend_input_len: int, max_new_tokens: int, swa_host_hit_length: int = 0, - *, - full_tokens: Optional[int] = None, - unified_swa_tokens: Optional[int] = None, ) -> bool: """True when a request's SWA budget exceeds the *entire* SWA pool, so it can never be admitted whole no matter how far the pool drains. @@ -942,72 +917,83 @@ def _swa_req_never_fits( swa_tokens = self._swa_budget_for_req( extend_input_len, max_new_tokens, swa_host_hit_length ) - if self.is_unified_swa: - assert full_tokens is not None - assert unified_swa_tokens is not None - return not self._unified_swa_reservation_fits( - full_tokens, unified_swa_tokens, empty_pool=True - ) if self._swa_req_ring: return swa_tokens > self.token_to_kv_pool_allocator.size_swa return swa_tokens >= self.token_to_kv_pool_allocator.size_swa - def _check_swa_admission( + def _check_unified_swa_admission( self, req: Req, *, extend_input_len: int, total_tokens: int, ) -> tuple[bool, Optional[int]]: - """Check SWA pressure and return the usable chunk limit.""" - if not self.is_unified_swa: - extend_input_len = self.ceil_paged_tokens(extend_input_len) max_new_tokens = self._swa_new_tokens(req) swa_needed = self._swa_budget_for_req( extend_input_len, max_new_tokens, swa_host_hit_length=req.swa_host_hit_length, ) - reservation_fits = ( - self._unified_swa_reservation_fits(total_tokens, swa_needed) - if self.is_unified_swa - else swa_needed <= self.rem_swa_tokens - if self._swa_req_ring - else swa_needed < self.rem_swa_tokens - ) - if reservation_fits: + if self._unified_swa_reservation_fits(total_tokens, swa_needed): return True, self.rem_chunk_tokens - full_ever_tokens = unified_swa_tokens = None - if self.is_unified_swa: - full_input_tokens = len(req.full_untruncated_fill_ids) - full_ever_tokens = full_input_tokens + max_new_tokens + self.page_size - unified_swa_tokens = self._swa_budget_for_req( - full_input_tokens, max_new_tokens - ) - if not self._swa_req_never_fits( - extend_input_len, - max_new_tokens, - req.swa_host_hit_length, - full_tokens=full_ever_tokens, - unified_swa_tokens=unified_swa_tokens, + # Only shrink requests that cannot fit even after the pool drains; + # transient pressure must wait instead of consuming decode headroom. + full_input_tokens = len(req.full_untruncated_fill_ids) + if self._unified_swa_reservation_fits( + full_input_tokens + max_new_tokens + self.page_size, + self._swa_budget_for_req(full_input_tokens, max_new_tokens), + empty_pool=True, ): return False, None - - max_chunk_tokens = self.rem_chunk_tokens - chunk_max_new = max_new_tokens - if self.is_unified_swa: - max_chunk_tokens = min(max_chunk_tokens or 0, max(0, extend_input_len - 1)) - chunk_max_new = 0 - swa_cap = self._swa_chunk_cap( - chunk_max_new, + swa_cap = self._unified_swa_chunk_cap( + 0, req.swa_host_hit_length, - max_chunk_tokens=max_chunk_tokens, + max_chunk_tokens=min( + self.rem_chunk_tokens or 0, max(0, extend_input_len - 1) + ), ) if self.rem_chunk_tokens is None or swa_cap <= 0: return False, None return True, min(self.rem_chunk_tokens, swa_cap) + def _check_swa_admission( + self, req: Req, *, extend_input_len: int + ) -> tuple[bool, Optional[int]]: + extend_input_len = self.ceil_paged_tokens(extend_input_len) + max_new_tokens = self._swa_new_tokens(req) + swa_needed = self._swa_budget_for_req( + extend_input_len, max_new_tokens, req.swa_host_hit_length + ) + # Ring-slot capacity is exact; the fixed token partition keeps slack. + if ( + swa_needed <= self.rem_swa_tokens + if self._swa_req_ring + else swa_needed < self.rem_swa_tokens + ): + return True, self.rem_chunk_tokens + if not self._swa_req_never_fits( + extend_input_len, max_new_tokens, req.swa_host_hit_length + ): + return False, None + swa_cap = self._swa_chunk_cap(max_new_tokens, req.swa_host_hit_length) + if self.rem_chunk_tokens is None or swa_cap <= 0: + return False, None + return True, min(self.rem_chunk_tokens, swa_cap) + + def _check_prefill_budget( + self, req: Req, *, extend_input_len: int, total_tokens: int + ) -> tuple[bool, Optional[int]]: + if self.is_unified_swa: + return self._check_unified_swa_admission( + req, extend_input_len=extend_input_len, total_tokens=total_tokens + ) + if total_tokens >= self.rem_total_tokens: + return False, None + if self.is_hybrid_swa: + return self._check_swa_admission(req, extend_input_len=extend_input_len) + return True, self.rem_chunk_tokens + def _mamba_gap_budget_for_req(self, req: Req) -> int: """Shared-gap reservation (full-token-equivalents) for a request's new mamba state. Charged only on the SHARED Mamba pool (`_mamba_slot_cost > 0`) @@ -1215,6 +1201,25 @@ def add_dllm_staging_req(self, req: Req): else AddReqResult.CONTINUE ) + def _unified_swa_chunk_limit( + self, req: Req, extend_input_len: int, max_chunk_tokens: int + ) -> Optional[int]: + candidate_tokens = min(extend_input_len, max_chunk_tokens) + candidate_finishes = candidate_tokens >= extend_input_len + max_new_tokens = self._swa_new_tokens(req) if candidate_finishes else 0 + if self._unified_swa_reservation_fits( + self.ceil_paged_tokens(candidate_tokens) + max_new_tokens + self.page_size, + self._swa_budget_for_req(candidate_tokens, max_new_tokens), + ): + return max_chunk_tokens + max_partial = ( + max(0, candidate_tokens - 1) if candidate_finishes else candidate_tokens + ) + joint_cap = self._unified_swa_chunk_cap(0, max_chunk_tokens=max_partial) + if joint_cap <= 0: + return None + return min(max_chunk_tokens, joint_cap) + def add_chunked_req(self, req: Req): if self.dllm_config is not None: _rem_tokens = self._get_dllm_remain_tokens() @@ -1223,17 +1228,11 @@ def add_chunked_req(self, req: Req): _rem_tokens = self.rem_chunk_tokens else: _rem_tokens = min(self.rem_chunk_tokens, int(self.rem_total_tokens)) - if ( - self.is_hybrid_swa - and not self.is_unified_swa - and not self._swa_req_ring - ): - # alloc_extend needs extend_num_tokens + page_size per request, - # so reserve one page here to avoid OOM. - # Ring mode skips it: rem_swa_tokens counts slots, not chunk tokens. - _rem_tokens = min( - _rem_tokens, int(self.rem_swa_tokens) - self.page_size - ) + if self.is_hybrid_swa and not self._swa_req_ring: + # alloc_extend needs one extra page; ring mode counts slots. + _rem_tokens = min( + _rem_tokens, int(self.rem_swa_tokens) - self.page_size + ) # The chunked_req must be added to the list; otherwise, it will cause a memory leak. # Therefore, in certain cases where _rem_tokens <= 0, it should be replaced with rem_chunk_tokens. if _rem_tokens <= 0: @@ -1256,22 +1255,11 @@ def add_chunked_req(self, req: Req): req.prefix_indices ) if self.is_unified_swa: - candidate_tokens = min(cand_extend_input_len, _rem_tokens) - candidate_finishes = candidate_tokens >= cand_extend_input_len - candidate_max_new = self._swa_new_tokens(req) if candidate_finishes else 0 - if not self._unified_swa_reservation_fits( - self.ceil_paged_tokens(candidate_tokens) - + candidate_max_new - + self.page_size, - self._swa_budget_for_req(candidate_tokens, candidate_max_new), - ): - max_partial = candidate_tokens - if candidate_finishes: - max_partial = max(0, candidate_tokens - 1) - joint_cap = self._swa_chunk_cap(0, max_chunk_tokens=max_partial) - if joint_cap <= 0: - return req - _rem_tokens = min(_rem_tokens, joint_cap) + _rem_tokens = self._unified_swa_chunk_limit( + req, cand_extend_input_len, _rem_tokens + ) + if _rem_tokens is None: + return req truncated = cand_extend_input_len > _rem_tokens new_len = min(cand_extend_input_len, _rem_tokens) req.set_extend_range(len(req.prefix_indices), len(req.prefix_indices) + new_len) @@ -1316,22 +1304,25 @@ def add_one_req_ignore_eos(self, req: Req): # Shared Mamba pool: fold the new mamba state's shared-gap cost into the # budget gate so admission can't over-commit (0 for baseline / non-Mamba). paged_input += self._mamba_gap_budget_for_req(req) - if self.is_hybrid_swa: + if self.is_unified_swa: max_new_tokens = self._swa_new_tokens(req) swa_needed = self._swa_budget_for_req(cand_extend_input_len, max_new_tokens) - if self.is_unified_swa: - if not self._unified_swa_reservation_fits( - paged_input + max_new_tokens + self.page_size, - swa_needed, - ): - return AddReqResult.NO_TOKEN - elif ( - paged_input > min(self.cur_rem_tokens, self.rem_total_tokens) - or swa_needed > self.rem_swa_tokens + if not self._unified_swa_reservation_fits( + paged_input + max_new_tokens + self.page_size, + swa_needed, + ): + return AddReqResult.NO_TOKEN + else: + if paged_input > min(self.cur_rem_tokens, self.rem_total_tokens): + return AddReqResult.NO_TOKEN + if ( + self.is_hybrid_swa + and self._swa_budget_for_req( + cand_extend_input_len, self._swa_new_tokens(req) + ) + > self.rem_swa_tokens ): return AddReqResult.NO_TOKEN - elif paged_input > min(self.cur_rem_tokens, self.rem_total_tokens): - return AddReqResult.NO_TOKEN def add_req_state(r, insert_sort=False): new_token_ratio = ( @@ -1482,19 +1473,14 @@ def add_one_req( real_input_tokens = self.ceil_paged_tokens(swa_extend_input_len) prefix_len = len(req.prefix_indices) - if not self.is_unified_swa and total_tokens >= self.rem_total_tokens: + can_admit, chunk_tokens_limit = self._check_prefill_budget( + req, + extend_input_len=swa_extend_input_len, + total_tokens=total_tokens, + ) + if not can_admit: return AddReqResult.NO_TOKEN - chunk_tokens_limit = self.rem_chunk_tokens - if self.is_hybrid_swa: - can_admit, chunk_tokens_limit = self._check_swa_admission( - req, - extend_input_len=swa_extend_input_len, - total_tokens=total_tokens, - ) - if not can_admit: - return AddReqResult.NO_TOKEN - if ( self.rem_chunk_tokens is None and len(self.can_run_list) != 0 @@ -1506,19 +1492,15 @@ def add_one_req( return AddReqResult.OTHER with self._lock_node(req.last_node): - # self.rem_total_tokens may decrease after the lock acquisition - if not self.is_unified_swa and total_tokens >= self.rem_total_tokens: + # Locking removes evictable capacity, so repeat the budget check. + can_admit, chunk_tokens_limit = self._check_prefill_budget( + req, + extend_input_len=swa_extend_input_len, + total_tokens=total_tokens, + ) + if not can_admit: return AddReqResult.NO_TOKEN - if self.is_hybrid_swa: - can_admit, chunk_tokens_limit = self._check_swa_admission( - req, - extend_input_len=swa_extend_input_len, - total_tokens=total_tokens, - ) - if not can_admit: - return AddReqResult.NO_TOKEN - # Negotiate only after every KV-budget gate (a NO_TOKEN rank must # report not-prefillable via finalize()) and before init_load_back # (a delay verdict must not start KV load-back). diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index a40632973471..2a67ff080fad 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -281,7 +281,7 @@ ) from sglang.srt.mem_cache import kv_cache_builder from sglang.srt.mem_cache.allocator.unified_hybrid_swa import ( - UnifiedSWATokenToKVPoolAllocator, + supports_swa_byte_budget, ) from sglang.srt.mem_cache.common import ( maybe_cache_unfinished_req, @@ -2524,7 +2524,6 @@ def init_req_max_new_tokens(self, req): max_new_tokens = min(max_new_tokens, self.max_new_tokens_limit) # Keep this bound consistent with PrefillAdder's admission budget. - paged_input_len = -(-input_len // self.page_size) * self.page_size max_new_tokens = max( 0, min( @@ -2532,45 +2531,11 @@ def init_req_max_new_tokens(self, req): self.max_req_len - input_len - 1, ), ) - allocator = self.token_to_kv_pool_allocator - if ( - isinstance(allocator, UnifiedSWATokenToKVPoolAllocator) - and allocator.supports_asymmetric_reservation - ): - - def fits_shared_pool(candidate: int) -> bool: - full_tokens = input_len + candidate + self.page_size - swa_tokens = estimate_swa_kv_tokens( - input_len, - candidate, - sliding_window_size=self.sliding_window_size, - page_size=self.page_size, - allocation_limit=self.chunked_prefill_size, - ) - return allocator.can_reserve( - full_tokens, - swa_tokens, - empty_pool=True, - require_token_slack=True, - ) - - if not fits_shared_pool(0): - req.set_finish_with_abort( - "Request prompt exceeds the unified FULL/SWA KV byte budget: " - f"input_len={input_len}." - ) - max_new_tokens = 0 - elif not fits_shared_pool(max_new_tokens): - lo, hi = 0, max_new_tokens - while lo < hi: - mid = (lo + hi + 1) // 2 - if fits_shared_pool(mid): - lo = mid - else: - hi = mid - 1 - max_new_tokens = lo + if supports_swa_byte_budget(self.token_to_kv_pool_allocator): + max_new_tokens = self._clip_unified_swa_max_new_tokens(req, max_new_tokens) else: # Non-unified pools keep their fixed token partition. + paged_input_len = -(-input_len // self.page_size) * self.page_size max_new_tokens = min( max_new_tokens, self.max_total_num_tokens * get_parallel().attn_dcp_size @@ -2585,6 +2550,41 @@ def fits_shared_pool(candidate: int) -> bool: if req.sampling_params.min_new_tokens > req.sampling_params.max_new_tokens: req.sampling_params.min_new_tokens = req.sampling_params.max_new_tokens + def _clip_unified_swa_max_new_tokens(self, req: Req, max_new_tokens: int) -> int: + input_len = len(req.origin_input_ids) + + def fits_shared_pool(candidate: int) -> bool: + swa_tokens = estimate_swa_kv_tokens( + input_len, + candidate, + sliding_window_size=self.sliding_window_size, + page_size=self.page_size, + allocation_limit=self.chunked_prefill_size, + ) + return self.token_to_kv_pool_allocator.can_reserve( + input_len + candidate + self.page_size, + swa_tokens, + empty_pool=True, + require_token_slack=True, + ) + + if not fits_shared_pool(0): + req.set_finish_with_abort( + "Request prompt exceeds the unified FULL/SWA KV byte budget: " + f"input_len={input_len}." + ) + return 0 + if fits_shared_pool(max_new_tokens): + return max_new_tokens + lo, hi = 0, max_new_tokens + while lo < hi: + mid = (lo + hi + 1) // 2 + if fits_shared_pool(mid): + lo = mid + else: + hi = mid - 1 + return lo + def _process_and_broadcast_mm_inputs( self, raw_mm_inputs, diff --git a/python/sglang/srt/managers/scheduler_components/invariant_checker.py b/python/sglang/srt/managers/scheduler_components/invariant_checker.py index a3a95d0128d8..30742013ad54 100644 --- a/python/sglang/srt/managers/scheduler_components/invariant_checker.py +++ b/python/sglang/srt/managers/scheduler_components/invariant_checker.py @@ -22,10 +22,6 @@ ) from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator from sglang.srt.mem_cache.allocator.swa import is_swa_req_ring -from sglang.srt.mem_cache.allocator.unified_hybrid_swa import ( - UnifiedMambaSWATokenToKVPoolAllocator, - UnifiedSWATokenToKVPoolAllocator, -) from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache from sglang.srt.mem_cache.memory_pool import ReqToTokenPool from sglang.srt.observability.scheduler_stage_metrics import ( @@ -96,18 +92,12 @@ def _check_pool_invariant( def _check_full_pool(self, ps: PoolStats, uncached: int = 0) -> Tuple[bool, str]: allocator = self.token_to_kv_pool_allocator - full_capacity = ( - allocator.current_full_capacity - if isinstance(allocator, UnifiedSWATokenToKVPoolAllocator) - and allocator.supports_asymmetric_reservation - else self.full_tokens_per_layer - ) - if self.is_hybrid_swa and not full_capacity: + if self.is_hybrid_swa and not ps.full_capacity: return False, "" if self.is_hybrid_swa: protected = self.tree_cache.full_protected_size() session_held = self.pool_stats_observer.session_held_full_tokens() - total = full_capacity + total = ps.full_capacity elif self.is_hybrid_ssm: # Branch on cache type for the protected accessor (MambaRadixCache # splits full/mamba; ChunkCache only has the single protected_size). @@ -136,14 +126,9 @@ def _check_full_pool(self, ps: PoolStats, uncached: int = 0) -> Tuple[bool, str] // allocator.page_size * allocator.page_size ) - full_available = ps.full_available_size - if isinstance(allocator, UnifiedMambaSWATokenToKVPoolAllocator): - # Pair the static per-layer total with the conserve view, never the - # byte-coordinated one -- see `conserve_full_available_size`. - full_available = allocator.conserve_full_available_size() leak, msg = self._check_pool_invariant( "full", - full_available, + ps.full_available_size, full_evictable_size, protected, session_held, @@ -169,23 +154,13 @@ def _check_swa_pool(self, ps: PoolStats, uncached: int = 0) -> Tuple[bool, str]: f"evictable={ps.swa_evictable_size}, " f"total={self.swa_tokens_per_layer}" ) - swa_available = ps.swa_available_size - if isinstance(allocator, UnifiedMambaSWATokenToKVPoolAllocator): - # Tri-pool: same floating-boundary phantom as the full pool -- use the - # slot-conservation view, not the byte-coordinated min (see _check_full_pool). - swa_available = allocator.conserve_swa_available_size() return self._check_pool_invariant( "swa", - swa_available, + ps.swa_available_size, ps.swa_evictable_size, self.tree_cache.swa_protected_size(), self.pool_stats_observer.session_held_swa_tokens(), - ( - allocator.current_swa_capacity - if isinstance(allocator, UnifiedSWATokenToKVPoolAllocator) - and allocator.supports_asymmetric_reservation - else self.swa_tokens_per_layer - ), + ps.swa_capacity, uncached, ) diff --git a/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py b/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py index 1b8577b3b94a..38c6a5cc2797 100644 --- a/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py +++ b/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py @@ -14,7 +14,7 @@ from sglang.srt.mem_cache.allocator.swa import is_swa_req_ring from sglang.srt.mem_cache.allocator.unified_hybrid_swa import ( UnifiedMambaSWATokenToKVPoolAllocator, - UnifiedSWATokenToKVPoolAllocator, + supports_swa_byte_budget, ) if TYPE_CHECKING: @@ -39,6 +39,8 @@ class PoolStats: is_hisparse: bool = False # For hybrid-swa pools + full_capacity: Optional[int] = None + swa_capacity: Optional[int] = None swa_num_used: Optional[int] = None swa_token_usage: Optional[float] = None swa_available_size: Optional[int] = None @@ -293,13 +295,12 @@ def _get_swa_token_info(self) -> PoolStats: allocator = self.token_to_kv_pool_allocator full_capacity = self.full_tokens_per_layer swa_capacity = self.swa_tokens_per_layer - if ( - isinstance(allocator, UnifiedSWATokenToKVPoolAllocator) - and allocator.supports_asymmetric_reservation - ): + if supports_swa_byte_budget(allocator): full_capacity = allocator.current_full_capacity swa_capacity = allocator.current_swa_capacity - if isinstance(allocator, UnifiedMambaSWATokenToKVPoolAllocator): + full_available_size = allocator.full_available_size() + swa_available_size = allocator.swa_available_size() + elif isinstance(allocator, UnifiedMambaSWATokenToKVPoolAllocator): # The tri-pool reports static capacities paired with conserve views. full_available_size = allocator.conserve_full_available_size() swa_available_size = allocator.conserve_swa_available_size() @@ -336,6 +337,8 @@ def _get_swa_token_info(self) -> PoolStats: return PoolStats( is_hybrid_swa=True, + full_capacity=full_capacity, + swa_capacity=swa_capacity, full_num_used=full_num_used, full_token_usage=full_token_usage, full_available_size=full_available_size, diff --git a/python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py b/python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py index b824ed062777..42658e9a7494 100644 --- a/python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py +++ b/python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py @@ -18,7 +18,7 @@ import logging import math -from typing import Callable, List, Optional, Sequence, Tuple +from typing import Callable, List, Optional, Sequence, Tuple, TypeGuard import torch from torch.profiler import record_function @@ -43,6 +43,16 @@ logger = logging.getLogger(__name__) +def supports_swa_byte_budget( + allocator: BaseTokenToKVPoolAllocator | None, +) -> TypeGuard[UnifiedSWATokenToKVPoolAllocator]: + """Whether FULL/SWA demand can be checked against a two-pool byte budget.""" + return ( + isinstance(allocator, UnifiedSWATokenToKVPoolAllocator) + and allocator.supports_asymmetric_reservation + ) + + class UnifiedSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator): """Composite allocator for the hybrid SWA pair (full + swa MHA sub-pools). @@ -50,6 +60,8 @@ class UnifiedSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator): `available_size()` (joint bytes, in TOKENS) is the only safe alloc pre-check. """ + # Unequal FULL/SWA reservations share one byte budget. The tri-pool opts + # out because its Mamba state competes for those same bytes. supports_asymmetric_reservation = True # Parent's `size` property has no setter but base init does `self.size = size`; @@ -395,15 +407,6 @@ def can_reserve( """ if full_tokens < 0 or swa_tokens < 0: return False - if not self.supports_asymmetric_reservation: - if ( - full_tokens != swa_tokens - or full_evictable_tokens - or swa_evictable_tokens - or empty_pool - ): - return False - return full_tokens <= self.available_size() if require_token_slack and ( full_tokens >= self.size_full or swa_tokens >= self.size_swa ): @@ -504,13 +507,6 @@ def ensure_capacity(self, full_tokens: int, swa_tokens: int) -> bool: return False if full_tokens == 0 and swa_tokens == 0: return True - if not self.supports_asymmetric_reservation: - if full_tokens != swa_tokens: - return False - need_tokens = int(full_tokens) - if need_tokens <= self.available_size(): - return True - return _relieve_for_alloc(self, need_tokens) page_size = self.page_size num_full_pages = (int(full_tokens) + page_size - 1) // page_size num_swa_pages = (int(swa_tokens) + page_size - 1) // page_size @@ -1149,6 +1145,37 @@ def _wire_peers(self) -> None: # -- capacity -- + def can_reserve( + self, + full_tokens: int | float, + swa_tokens: int | float, + *, + full_evictable_tokens: int = 0, + swa_evictable_tokens: int = 0, + empty_pool: bool = False, + require_token_slack: bool = False, + ) -> bool: + if ( + full_tokens < 0 + or swa_tokens < 0 + or full_tokens != swa_tokens + or full_evictable_tokens + or swa_evictable_tokens + or empty_pool + ): + return False + return full_tokens <= self.available_size() + + def ensure_capacity(self, full_tokens: int, swa_tokens: int) -> bool: + if full_tokens < 0 or swa_tokens < 0 or full_tokens != swa_tokens: + return False + if full_tokens == 0: + return True + need_tokens = int(full_tokens) + if need_tokens <= self.available_size(): + return True + return _relieve_for_alloc(self, need_tokens) + def _compute_available_size(self) -> int: """Joint TOKENS for `alloc(N)`: N costs N full pages AND N swa pages, drawn from DIFFERENT bands -- full extends only into the high band, the float into diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index 289c6f19b4e3..f267a56660fa 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/python/sglang/srt/mem_cache/common.py @@ -162,14 +162,10 @@ def evict_from_tree_cache(tree_cache: BasePrefixCache | None, num_tokens: int): if isinstance(allocator, SWATokenToKVPoolAllocator): # Hybrid allocator from sglang.srt.mem_cache.allocator.unified_hybrid_swa import ( - UnifiedSWATokenToKVPoolAllocator, + supports_swa_byte_budget, ) - is_unified_allocator = ( - isinstance(allocator, UnifiedSWATokenToKVPoolAllocator) - and allocator.supports_asymmetric_reservation - ) - if is_unified_allocator: + if supports_swa_byte_budget(allocator): reclaim_plan = allocator.reclaim_plan( num_tokens, num_tokens, @@ -179,18 +175,23 @@ def evict_from_tree_cache(tree_cache: BasePrefixCache | None, num_tokens: int): if reclaim_plan is None: return full_num_tokens, swa_num_tokens = reclaim_plan - else: - full_num_tokens = max(0, num_tokens - allocator.full_available_size()) - swa_num_tokens = max(0, num_tokens - allocator.swa_available_size()) + if full_num_tokens or swa_num_tokens: + tree_cache.evict_for_alloc( + EvictParams( + num_tokens=full_num_tokens, swa_num_tokens=swa_num_tokens + ) + ) + # Realize any compaction assumed by the reclaim plan before callers + # inspect availability. + allocator.ensure_capacity(num_tokens, num_tokens) + return + full_num_tokens = max(0, num_tokens - allocator.full_available_size()) + swa_num_tokens = max(0, num_tokens - allocator.swa_available_size()) if full_num_tokens or swa_num_tokens: tree_cache.evict_for_alloc( EvictParams(num_tokens=full_num_tokens, swa_num_tokens=swa_num_tokens) ) - if is_unified_allocator: - # The reclaim plan may rely on compacting cross-side holes. Make - # that projected state real before callers inspect availability. - allocator.ensure_capacity(num_tokens, num_tokens) else: # Standard allocator: evict only the shortfall (mirrors the SWA arm) available_size = allocator.available_size() diff --git a/python/sglang/srt/model_executor/pool_configurator.py b/python/sglang/srt/model_executor/pool_configurator.py index 132ecf358772..1d54e801b815 100644 --- a/python/sglang/srt/model_executor/pool_configurator.py +++ b/python/sglang/srt/model_executor/pool_configurator.py @@ -773,6 +773,9 @@ def align_page_size(x: int) -> int: f"full_layer_tokens={full_tokens}, swa_layer_tokens={swa_tokens}" ) + return self._make_pool_config(full_tokens, swa_tokens) + + def _make_pool_config(self, full_tokens: int, swa_tokens: int) -> MemoryPoolConfig: return MemoryPoolConfig( max_total_num_tokens=full_tokens, full_max_total_num_tokens=full_tokens, @@ -899,16 +902,7 @@ def calculate_pool_sizes( f"Reduce --max-running-requests, lower SGLANG_SWA_EVICTION_INTERVAL, " f"or increase --mem-fraction-static." ) - return MemoryPoolConfig( - max_total_num_tokens=full_tokens, - full_max_total_num_tokens=full_tokens, - swa_max_total_num_tokens=swa_tokens, - unified_memory_pool_bytes=( - self._unified_pool_bytes(full_tokens, swa_tokens) - if self._enable_unified_memory - else None - ), - ) + return self._make_pool_config(full_tokens, swa_tokens) def calculate_pool_sizes_from_max_tokens( self, max_total_num_tokens: int, page_size: int @@ -916,17 +910,8 @@ def calculate_pool_sizes_from_max_tokens( # Constrained max_total goes to the full pool; SWA stays at its cap. swa_tokens = ceil_align(self._swa_cap, page_size) full_tokens = (max_total_num_tokens // page_size) * page_size - return MemoryPoolConfig( - max_total_num_tokens=full_tokens, - full_max_total_num_tokens=full_tokens, - swa_max_total_num_tokens=min(swa_tokens, max_total_num_tokens), - unified_memory_pool_bytes=( - self._unified_pool_bytes( - full_tokens, min(swa_tokens, max_total_num_tokens) - ) - if self._enable_unified_memory - else None - ), + return self._make_pool_config( + full_tokens, min(swa_tokens, max_total_num_tokens) ) From fa6c95ace4fe205475252d238802fa06f12124ae Mon Sep 17 00:00:00 2001 From: yhzhuang Date: Sat, 12 Sep 2026 22:51:55 -0700 Subject: [PATCH 45/55] Simplify shared SWA budget paths --- python/sglang/srt/managers/schedule_policy.py | 37 ++++++++----------- .../pool_stats_observer.py | 4 +- python/sglang/srt/mem_cache/common.py | 23 +++++------- .../sglang/srt/model_executor/model_runner.py | 15 +++----- .../unit/managers/test_prefill_adder.py | 2 +- 5 files changed, 33 insertions(+), 48 deletions(-) diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index 54bece2ae4c5..3ab80e4b717d 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -812,13 +812,13 @@ def _swa_budget_for_req( if self._swa_req_ring: # One ring slot per request, in the same unit as swa_available_size. return allocator.swa_ring_cost_tokens - if self.rem_chunk_tokens is not None: - alloc = min(extend_input_len, self.rem_chunk_tokens) - else: - alloc = extend_input_len - window = self.tree_cache.sliding_window_size - return max(alloc - window, 0) + self._swa_reserved_tokens( - extend_input_len, max_new_tokens, swa_host_hit_length + return estimate_swa_kv_tokens( + extend_input_len, + max_new_tokens, + sliding_window_size=self.tree_cache.sliding_window_size, + page_size=self.page_size, + allocation_limit=self.rem_chunk_tokens, + host_hit_length=swa_host_hit_length, ) def _swa_reserved_tokens( @@ -1304,25 +1304,20 @@ def add_one_req_ignore_eos(self, req: Req): # Shared Mamba pool: fold the new mamba state's shared-gap cost into the # budget gate so admission can't over-commit (0 for baseline / non-Mamba). paged_input += self._mamba_gap_budget_for_req(req) - if self.is_unified_swa: + if self.is_hybrid_swa: max_new_tokens = self._swa_new_tokens(req) swa_needed = self._swa_budget_for_req(cand_extend_input_len, max_new_tokens) - if not self._unified_swa_reservation_fits( + if self.is_unified_swa: + fits = self._unified_swa_reservation_fits( paged_input + max_new_tokens + self.page_size, swa_needed, - ): - return AddReqResult.NO_TOKEN + ) else: - if paged_input > min(self.cur_rem_tokens, self.rem_total_tokens): - return AddReqResult.NO_TOKEN - if ( - self.is_hybrid_swa - and self._swa_budget_for_req( - cand_extend_input_len, self._swa_new_tokens(req) - ) - > self.rem_swa_tokens - ): - return AddReqResult.NO_TOKEN + fits = paged_input <= min(self.cur_rem_tokens, self.rem_total_tokens) and ( + not self.is_hybrid_swa or swa_needed <= self.rem_swa_tokens + ) + if not fits: + return AddReqResult.NO_TOKEN def add_req_state(r, insert_sort=False): new_token_ratio = ( diff --git a/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py b/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py index 38c6a5cc2797..138663ef9127 100644 --- a/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py +++ b/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py @@ -298,9 +298,7 @@ def _get_swa_token_info(self) -> PoolStats: if supports_swa_byte_budget(allocator): full_capacity = allocator.current_full_capacity swa_capacity = allocator.current_swa_capacity - full_available_size = allocator.full_available_size() - swa_available_size = allocator.swa_available_size() - elif isinstance(allocator, UnifiedMambaSWATokenToKVPoolAllocator): + if isinstance(allocator, UnifiedMambaSWATokenToKVPoolAllocator): # The tri-pool reports static capacities paired with conserve views. full_available_size = allocator.conserve_full_available_size() swa_available_size = allocator.conserve_swa_available_size() diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index f267a56660fa..2e4012da5b21 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/python/sglang/srt/mem_cache/common.py @@ -165,7 +165,8 @@ def evict_from_tree_cache(tree_cache: BasePrefixCache | None, num_tokens: int): supports_swa_byte_budget, ) - if supports_swa_byte_budget(allocator): + is_unified_swa = supports_swa_byte_budget(allocator) + if is_unified_swa: reclaim_plan = allocator.reclaim_plan( num_tokens, num_tokens, @@ -175,23 +176,17 @@ def evict_from_tree_cache(tree_cache: BasePrefixCache | None, num_tokens: int): if reclaim_plan is None: return full_num_tokens, swa_num_tokens = reclaim_plan - if full_num_tokens or swa_num_tokens: - tree_cache.evict_for_alloc( - EvictParams( - num_tokens=full_num_tokens, swa_num_tokens=swa_num_tokens - ) - ) - # Realize any compaction assumed by the reclaim plan before callers - # inspect availability. - allocator.ensure_capacity(num_tokens, num_tokens) - return - - full_num_tokens = max(0, num_tokens - allocator.full_available_size()) - swa_num_tokens = max(0, num_tokens - allocator.swa_available_size()) + else: + full_num_tokens = max(0, num_tokens - allocator.full_available_size()) + swa_num_tokens = max(0, num_tokens - allocator.swa_available_size()) if full_num_tokens or swa_num_tokens: tree_cache.evict_for_alloc( EvictParams(num_tokens=full_num_tokens, swa_num_tokens=swa_num_tokens) ) + if is_unified_swa: + # Even a zero-reclaim plan can depend on compaction before callers + # inspect availability. + allocator.ensure_capacity(num_tokens, num_tokens) else: # Standard allocator: evict only the shortfall (mirrors the SWA arm) available_size = allocator.available_size() diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index d20eb3e38ec4..6395603fdde9 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -1369,15 +1369,12 @@ def unload_lora_adapter(self, lora_ref: LoRARef): @property def effective_max_total_num_tokens(self): """Return the max token pool size considering hybrid swa settings.""" - if self.is_hybrid_swa: - if get_memory().enable_unified_memory: - capacity = self.token_to_kv_pool_allocator.size_full - if get_schedule().max_total_tokens is not None: - capacity = min(capacity, get_schedule().max_total_tokens) - else: - capacity = ( - self.full_max_total_num_tokens or self.swa_max_total_num_tokens - ) + if self.is_hybrid_swa and get_memory().enable_unified_memory: + capacity = self.token_to_kv_pool_allocator.size_full + if get_schedule().max_total_tokens is not None: + capacity = min(capacity, get_schedule().max_total_tokens) + elif self.is_hybrid_swa: + capacity = self.full_max_total_num_tokens or self.swa_max_total_num_tokens else: capacity = self.max_total_num_tokens if (req_to_token_pool := getattr(self, "req_to_token_pool", None)) is not None: diff --git a/test/registered/unit/managers/test_prefill_adder.py b/test/registered/unit/managers/test_prefill_adder.py index 2fbe8e1eef37..84d9dcd8633e 100644 --- a/test/registered/unit/managers/test_prefill_adder.py +++ b/test/registered/unit/managers/test_prefill_adder.py @@ -653,7 +653,7 @@ def test_swa_admission_admits_short_cached_resume_at_two_window_pool(self): req.sampling_params = SimpleNamespace(max_new_tokens=40, ignore_eos=False) # Pre-fix: a constant sliding-window reservation rejects the resume. - with patch.object(adder, "_swa_reserved_tokens", return_value=WINDOW + PAGE): + with patch.object(adder, "_swa_budget_for_req", return_value=WINDOW + PAGE): self.assertIs( adder.add_one_req( req, has_chunked_req=False, truncation_align_size=None From d9c744726b41bbcfbf084a06a294ff6b1f1340f0 Mon Sep 17 00:00:00 2001 From: yhzhuang Date: Sun, 13 Sep 2026 00:57:32 -0700 Subject: [PATCH 46/55] Fix unified capturer indexing and asymmetric SWA reservation --- .../mem_cache/allocator/unified_hybrid_swa.py | 6 +- .../srt/mem_cache/kv_index_translator.py | 10 ++++ .../sglang/srt/model_executor/model_runner.py | 8 ++- python/sglang/srt/state_capturer/base.py | 12 +++- .../mem_cache/test_kv_index_translator.py | 56 ++++++++++++++++++- .../mem_cache/test_multi_ended_allocator.py | 21 ++++++- 6 files changed, 103 insertions(+), 10 deletions(-) diff --git a/python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py b/python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py index 42658e9a7494..bc4c6593baff 100644 --- a/python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py +++ b/python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py @@ -337,8 +337,8 @@ def reclaim_plan( page_size = self.page_size full_pages = (math.ceil(full_tokens) + page_size - 1) // page_size swa_pages = (math.ceil(swa_tokens) + page_size - 1) // page_size - if swa_pages > full_pages: - return None + # Restoring host SWA for device-resident FULL can require more new + # SWA pages than new FULL pages; only the shared budget constrains it. compacted = empty_pool or not self.lazy_compaction or self._compaction_allowed() @@ -510,8 +510,6 @@ def ensure_capacity(self, full_tokens: int, swa_tokens: int) -> bool: page_size = self.page_size num_full_pages = (int(full_tokens) + page_size - 1) // page_size num_swa_pages = (int(swa_tokens) + page_size - 1) // page_size - if num_swa_pages > num_full_pages: - return False if self._fits_page_demand( num_full_pages, num_swa_pages, diff --git a/python/sglang/srt/mem_cache/kv_index_translator.py b/python/sglang/srt/mem_cache/kv_index_translator.py index 2e1b970680fa..1ccb5b35077c 100644 --- a/python/sglang/srt/mem_cache/kv_index_translator.py +++ b/python/sglang/srt/mem_cache/kv_index_translator.py @@ -171,6 +171,16 @@ def __init__( ) self._index_table_memo: Optional[Tuple[weakref.ref, KVIndexTable]] = None + def capture_token_capacity(self, max_token_pool_size: int) -> int: + """Host capture rows are indexed by request-token IDs, not kernel IDs. + + Unified IDs span the whole virtual table even when admission is capped. + The logical page size also covers DCP-widened request-token IDs. + """ + if self.is_translating: + return self._full_v2p_table.numel() * self.page_size + return max_token_pool_size + self.page_size + # -- per-batch view -------------------------------------------------------- @property diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 6395603fdde9..021ee586d579 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -1114,7 +1114,9 @@ def init_routed_experts_capturer(self): RoutedExpertsCapturer.create( model=self.model, model_config=self.model_config, - num_tokens=self.max_token_pool_size + self.page_size, + num_tokens=self.kv_index_translator.capture_token_capacity( + self.max_token_pool_size + ), max_running_requests=self.max_running_requests, device=self.device, ) @@ -1124,7 +1126,9 @@ def init_indexer_capturer(self): set_global_indexer_capturer( create_indexer_capturer( model_config=self.model_config, - num_tokens=self.max_token_pool_size + self.page_size, + num_tokens=self.kv_index_translator.capture_token_capacity( + self.max_token_pool_size + ), max_running_requests=self.max_running_requests, device=self.device, ) diff --git a/python/sglang/srt/state_capturer/base.py b/python/sglang/srt/state_capturer/base.py index 6f6b8757841b..f451f7e6f5d2 100644 --- a/python/sglang/srt/state_capturer/base.py +++ b/python/sglang/srt/state_capturer/base.py @@ -201,13 +201,21 @@ def on_forward_end( slice_gpu = self._get_local_slice( forward_batch, can_run_graph, cuda_graph_batch ) + # get_topk reads req_to_token IDs; attention may have rebound the write + # loc to kernel-facing IDs, which are neither stable nor the same space. + out_cache_loc = forward_batch.out_cache_loc_virtual + if out_cache_loc is None: + out_cache_loc = forward_batch.out_cache_loc + else: + # Kernel batches may be padded; only real request tokens are stored. + slice_gpu = slice_gpu[: out_cache_loc.shape[0]] if no_copy_to_cpu: # Clone before the next overlapping forward reuses these buffers. return TopkCaptureOutput( - out_cache_loc=forward_batch.out_cache_loc.clone(), + out_cache_loc=out_cache_loc.clone(), topk=slice_gpu.clone(), host_cache=self.host_cache, ) - out_cache_loc_cpu = forward_batch.out_cache_loc.cpu() + out_cache_loc_cpu = out_cache_loc.cpu() self.host_cache.buffer[out_cache_loc_cpu] = slice_gpu.cpu() return None diff --git a/test/registered/unit/mem_cache/test_kv_index_translator.py b/test/registered/unit/mem_cache/test_kv_index_translator.py index a5e818e24143..fb497732be86 100644 --- a/test/registered/unit/mem_cache/test_kv_index_translator.py +++ b/test/registered/unit/mem_cache/test_kv_index_translator.py @@ -35,6 +35,8 @@ from sglang.srt.mem_cache.kv_index_translator import KVIndexTranslator, KVReadTables from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool from sglang.srt.mem_cache.unified_memory_pool import MHASubPoolSpec, UnifiedKVPool +from sglang.srt.state_capturer.base import BaseTopkCapturer +from sglang.test.test_utils import CustomTestCase _DEV = "cpu" _FULL_L = 2 @@ -135,6 +137,7 @@ def test_non_unified_returns_same_objects(self): device=_DEV, ) self.assertFalse(src.is_translating) + self.assertEqual(src.capture_token_capacity(17), 18) rows = torch.tensor([2, 0]) view = src.build_index_table( req_pool_indices=rows, seq_lens=torch.tensor([5, 3]) @@ -594,7 +597,7 @@ def test_dead_batch_never_matches(self): self.assertEqual(v2.ids.shape[0], 1) -class TestWriteLoc(unittest.TestCase): +class TestWriteLoc(CustomTestCase): """The two-phase write contract: `rebind_write_loc` rebinds the full side once at ForwardBatch construction, and the sliding-window write loc derives POINTWISE from the full-side values -- pads, slices, and fresh copies @@ -630,6 +633,57 @@ def test_rebind_translates_full_side_only(self): self.assertTrue(torch.equal(fb.out_cache_loc, want_full)) self.assertTrue(torch.equal(virt, keep)) + def test_topk_capture_round_trips_request_token_ids(self): + for ps in (1, 4, 64): + for translating in (False, True): + for overlap in (False, True): + with self.subTest( + page_size=ps, translating=translating, overlap=overlap + ): + src, _, _, _, virt, _, _ = self._built(ps=ps, n=3 * ps) + req_pool = SimpleNamespace(req_to_token=virt.clone()[None, :]) + fb = _FakeForwardBatch(out_cache_loc=virt.clone()) + if translating: + src.rebind_write_loc(fb) + fb.out_cache_loc = torch.cat( + [fb.out_cache_loc, virt.new_zeros(2)] + ) + else: + fb.out_cache_loc_virtual = None + + # Admission may be capped below IDs issued after reuse. + capacity = src.capture_token_capacity(ps) + self.assertGreater(capacity, int(virt.max())) + expected = ( + torch.arange(len(virt) * 4, dtype=torch.int32).reshape( + -1, 2, 2 + ) + + 1 + ) + cap = object.__new__(BaseTopkCapturer) + cap.topk_size = 2 + cap.device_cache = SimpleNamespace( + buffer=torch.cat([expected, expected.new_zeros(2, 2, 2)]) + ) + cap.host_cache = SimpleNamespace( + buffer=torch.zeros(capacity, 2, 2, dtype=torch.int32) + ) + result = cap.on_forward_end( + fb, False, None, no_copy_to_cpu=overlap + ) + if overlap: + fb.out_cache_loc.zero_() + if translating: + fb.out_cache_loc_virtual.zero_() + cap.device_cache.buffer.zero_() + result.map_device_tensors(lambda value: value.cpu()) + result.finalize() + self.assertTrue( + torch.equal( + cap.get_topk(0, len(virt) + 1, req_pool), expected + ) + ) + def test_swa_write_loc_round_trips_from_full_side(self): """Derived property: `field(full(t)) == swa(t)` for any virtual run t, across page sizes and multipliers.""" diff --git a/test/registered/unit/mem_cache/test_multi_ended_allocator.py b/test/registered/unit/mem_cache/test_multi_ended_allocator.py index bac0687f9f1d..6f677c94f64e 100644 --- a/test/registered/unit/mem_cache/test_multi_ended_allocator.py +++ b/test/registered/unit/mem_cache/test_multi_ended_allocator.py @@ -625,7 +625,7 @@ def test_reclaim_plan_matches_exhaustive_page_targets(self): max_full = min(12, full_evictable) // page_size max_swa = min(8, swa_evictable) // page_size for full_pages in range(9): - for swa_pages in range(full_pages + 1): + for swa_pages in range(9): feasible = [ (full * page_size, swa * page_size) for swa in range(max_swa + 1) @@ -653,6 +653,25 @@ def test_reclaim_plan_matches_exhaustive_page_targets(self): feasible[0] if feasible else None, ) + def test_restore_swa_without_allocating_more_full(self): + _, allocator, _ = self._build(page_size=4) + indices = allocator.alloc(8) + full_before = allocator.translate_kv_indices_for_transfer(indices).clone() + allocator.free_swa(indices) + + self.assertEqual(allocator.reclaim_plan(0, 8), (0, 0)) + self.assertTrue(allocator.can_reserve(0, 8)) + self.assertTrue(allocator.ensure_capacity(0, 8)) + allocator.swa_attn_allocator.alloc_with_virtual((indices // 4).unique()) + self.assertTrue( + torch.equal( + allocator.translate_kv_indices_for_transfer(indices), full_before + ) + ) + self.assertTrue( + bool((allocator.swa_attn_allocator.translate_kv_loc(indices) > 0).all()) + ) + def test_empty_pool_reservation_matches_packed_byte_boundary(self): page_size = 4 _, allocator, _ = self._build( From 6cf450c3d5de54af033a8bed67b0e24cf17ae26e Mon Sep 17 00:00:00 2001 From: yhzhuang Date: Sun, 13 Sep 2026 01:34:47 -0700 Subject: [PATCH 47/55] Consolidate unified allocation, admission and capacity policies --- python/sglang/srt/managers/schedule_policy.py | 115 ++++++++---------- .../pool_stats_observer.py | 19 ++- .../mem_cache/allocator/unified_hybrid_swa.py | 25 +++- python/sglang/srt/mem_cache/common.py | 24 +--- .../srt/mem_cache/kv_cache_configurator.py | 22 +++- .../srt/mem_cache/unified_memory_pool.py | 6 - .../sglang/srt/model_executor/model_runner.py | 12 +- .../test_hisparse_max_token_pool_size.py | 2 + .../mem_cache/test_multi_ended_allocator.py | 2 +- .../test_unified_byte_budget_sizing.py | 12 +- .../test_unified_radix_cache_unittest.py | 2 +- 11 files changed, 126 insertions(+), 115 deletions(-) diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index 3ab80e4b717d..20c54e3d2560 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -88,7 +88,6 @@ def estimate_swa_kv_tokens( page_size: int, allocation_limit: Optional[int] = None, host_hit_length: int = 0, - include_allocated_tail: bool = True, ) -> int: """Peak SWA reservation for one prefill/decode request.""" if sliding_window_size is None or sliding_window_size <= 0: @@ -99,9 +98,9 @@ def estimate_swa_kv_tokens( if allocation_limit is None else min(extend_input_len, allocation_limit) ) - allocated_tail = ( - max(allocated - sliding_window_size, 0) if include_allocated_tail else 0 - ) + allocated_tail = max(allocated - sliding_window_size, 0) + # A full window would double-charge short cached-prefix resumes; + # including extend keeps the reservation above the prefill allocation. reserved = ( allocated_tail + min(extend_input_len + max_new_tokens, sliding_window_size) @@ -821,18 +820,6 @@ def _swa_budget_for_req( host_hit_length=swa_host_hit_length, ) - def _swa_reserved_tokens( - self, extend_input_len: int, max_new_tokens: int, swa_host_hit_length: int = 0 - ) -> int: - return estimate_swa_kv_tokens( - extend_input_len, - max_new_tokens, - sliding_window_size=self.tree_cache.sliding_window_size, - page_size=self.page_size, - host_hit_length=swa_host_hit_length, - include_allocated_tail=False, - ) - def _swa_new_tokens(self, req: Req) -> int: """Tokens a request may still decode, for SWA headroom sizing. Mirrors the remaining-then-clip order of add_one_req's max_new: clip-then-subtract @@ -893,8 +880,14 @@ def _unified_swa_chunk_cap( def _swa_chunk_cap(self, max_new_tokens: int, swa_host_hit_length: int = 0) -> int: """Largest chunk fitting the fixed SWA partition after decode headroom.""" - cap = int(self.rem_swa_tokens) - self._swa_reserved_tokens( - 0, max_new_tokens, swa_host_hit_length + # Only the sliding window stays locked across chunks, so the next + # chunk can reuse the transient space outside its decode headroom. + cap = int(self.rem_swa_tokens) - estimate_swa_kv_tokens( + 0, + max_new_tokens, + sliding_window_size=self.tree_cache.sliding_window_size, + page_size=self.page_size, + host_hit_length=swa_host_hit_length, ) return max(0, cap // self.page_size * self.page_size) @@ -921,62 +914,52 @@ def _swa_req_never_fits( return swa_tokens > self.token_to_kv_pool_allocator.size_swa return swa_tokens >= self.token_to_kv_pool_allocator.size_swa - def _check_unified_swa_admission( + def _check_swa_admission( self, req: Req, *, extend_input_len: int, total_tokens: int, ) -> tuple[bool, Optional[int]]: + if not self.is_unified_swa: + extend_input_len = self.ceil_paged_tokens(extend_input_len) max_new_tokens = self._swa_new_tokens(req) swa_needed = self._swa_budget_for_req( extend_input_len, max_new_tokens, swa_host_hit_length=req.swa_host_hit_length, ) - if self._unified_swa_reservation_fits(total_tokens, swa_needed): - return True, self.rem_chunk_tokens - - # Only shrink requests that cannot fit even after the pool drains; - # transient pressure must wait instead of consuming decode headroom. - full_input_tokens = len(req.full_untruncated_fill_ids) - if self._unified_swa_reservation_fits( - full_input_tokens + max_new_tokens + self.page_size, - self._swa_budget_for_req(full_input_tokens, max_new_tokens), - empty_pool=True, - ): - return False, None - swa_cap = self._unified_swa_chunk_cap( - 0, - req.swa_host_hit_length, - max_chunk_tokens=min( - self.rem_chunk_tokens or 0, max(0, extend_input_len - 1) - ), - ) - if self.rem_chunk_tokens is None or swa_cap <= 0: - return False, None - return True, min(self.rem_chunk_tokens, swa_cap) - - def _check_swa_admission( - self, req: Req, *, extend_input_len: int - ) -> tuple[bool, Optional[int]]: - extend_input_len = self.ceil_paged_tokens(extend_input_len) - max_new_tokens = self._swa_new_tokens(req) - swa_needed = self._swa_budget_for_req( - extend_input_len, max_new_tokens, req.swa_host_hit_length - ) - # Ring-slot capacity is exact; the fixed token partition keeps slack. - if ( - swa_needed <= self.rem_swa_tokens - if self._swa_req_ring - else swa_needed < self.rem_swa_tokens - ): + swa_cap = 0 + # Only permanent shortfalls may shrink a chunk; transient pressure waits. + if self.is_unified_swa: + fits = self._unified_swa_reservation_fits(total_tokens, swa_needed) + if not fits: + full_input_tokens = len(req.full_untruncated_fill_ids) + if not self._unified_swa_reservation_fits( + full_input_tokens + max_new_tokens + self.page_size, + self._swa_budget_for_req(full_input_tokens, max_new_tokens), + empty_pool=True, + ): + swa_cap = self._unified_swa_chunk_cap( + 0, + req.swa_host_hit_length, + max_chunk_tokens=min( + self.rem_chunk_tokens or 0, max(0, extend_input_len - 1) + ), + ) + else: + # Ring-slot capacity is exact; the fixed token partition keeps slack. + fits = ( + swa_needed <= self.rem_swa_tokens + if self._swa_req_ring + else swa_needed < self.rem_swa_tokens + ) + if not fits and self._swa_req_never_fits( + extend_input_len, max_new_tokens, req.swa_host_hit_length + ): + swa_cap = self._swa_chunk_cap(max_new_tokens, req.swa_host_hit_length) + if fits: return True, self.rem_chunk_tokens - if not self._swa_req_never_fits( - extend_input_len, max_new_tokens, req.swa_host_hit_length - ): - return False, None - swa_cap = self._swa_chunk_cap(max_new_tokens, req.swa_host_hit_length) if self.rem_chunk_tokens is None or swa_cap <= 0: return False, None return True, min(self.rem_chunk_tokens, swa_cap) @@ -984,14 +967,12 @@ def _check_swa_admission( def _check_prefill_budget( self, req: Req, *, extend_input_len: int, total_tokens: int ) -> tuple[bool, Optional[int]]: - if self.is_unified_swa: - return self._check_unified_swa_admission( - req, extend_input_len=extend_input_len, total_tokens=total_tokens - ) - if total_tokens >= self.rem_total_tokens: + if not self.is_unified_swa and total_tokens >= self.rem_total_tokens: return False, None if self.is_hybrid_swa: - return self._check_swa_admission(req, extend_input_len=extend_input_len) + return self._check_swa_admission( + req, extend_input_len=extend_input_len, total_tokens=total_tokens + ) return True, self.rem_chunk_tokens def _mamba_gap_budget_for_req(self, req: Req) -> int: diff --git a/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py b/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py index 138663ef9127..823fe15b67cf 100644 --- a/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py +++ b/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py @@ -291,13 +291,13 @@ def _get_mamba_token_info(self): mamba_evictable_size=mamba_evictable_size, ) - def _get_swa_token_info(self) -> PoolStats: + def _swa_capacity_and_available(self) -> tuple[tuple[int, int], tuple[int, int]]: allocator = self.token_to_kv_pool_allocator - full_capacity = self.full_tokens_per_layer - swa_capacity = self.swa_tokens_per_layer if supports_swa_byte_budget(allocator): - full_capacity = allocator.current_full_capacity - swa_capacity = allocator.current_swa_capacity + return ( + (allocator.current_full_capacity, allocator.full_available_size()), + (allocator.current_swa_capacity, allocator.swa_available_size()), + ) if isinstance(allocator, UnifiedMambaSWATokenToKVPoolAllocator): # The tri-pool reports static capacities paired with conserve views. full_available_size = allocator.conserve_full_available_size() @@ -305,6 +305,15 @@ def _get_swa_token_info(self) -> PoolStats: else: full_available_size = allocator.full_available_size() swa_available_size = allocator.swa_available_size() + return ( + (self.full_tokens_per_layer, full_available_size), + (self.swa_tokens_per_layer, swa_available_size), + ) + + def _get_swa_token_info(self) -> PoolStats: + (full_capacity, full_available_size), (swa_capacity, swa_available_size) = ( + self._swa_capacity_and_available() + ) full_evictable_size = self.tree_cache.full_evictable_size() swa_evictable_size = self.tree_cache.swa_evictable_size() # Per-request SWA ring: released with the req slot, yet cached radix diff --git a/python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py b/python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py index bc4c6593baff..7248edf28b0f 100644 --- a/python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py +++ b/python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py @@ -524,14 +524,35 @@ def ensure_capacity(self, full_tokens: int, swa_tokens: int) -> bool: compacted=True, ): return False - self.full_attn_allocator._flush(urgent=True) - self.swa_attn_allocator._flush(urgent=True) + self.full_attn_allocator.flush_for_allocation() + self.swa_attn_allocator.flush_for_allocation() return self._fits_page_demand( num_full_pages, num_swa_pages, compacted=False, ) + def evict_to_free_tokens(self, tree_cache, num_tokens: int) -> None: + from sglang.srt.mem_cache.base_prefix_cache import EvictParams + + if tree_cache is None or tree_cache.is_chunk_cache(): + return + reclaim_plan = self.reclaim_plan( + num_tokens, + num_tokens, + full_evictable_tokens=tree_cache.full_evictable_size(), + swa_evictable_tokens=tree_cache.swa_evictable_size(), + ) + if reclaim_plan is None: + return + full_reclaim, swa_reclaim = reclaim_plan + if full_reclaim or swa_reclaim: + tree_cache.evict_for_alloc( + EvictParams(num_tokens=full_reclaim, swa_num_tokens=swa_reclaim) + ) + # A zero-reclaim plan can still depend on compaction before allocation. + self.ensure_capacity(num_tokens, num_tokens) + @property def draft_virtual_id_space(self) -> int: return self.full_attn_allocator.max_slots - 1 diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index 2e4012da5b21..5f73d202a795 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/python/sglang/srt/mem_cache/common.py @@ -165,28 +165,16 @@ def evict_from_tree_cache(tree_cache: BasePrefixCache | None, num_tokens: int): supports_swa_byte_budget, ) - is_unified_swa = supports_swa_byte_budget(allocator) - if is_unified_swa: - reclaim_plan = allocator.reclaim_plan( - num_tokens, - num_tokens, - full_evictable_tokens=tree_cache.full_evictable_size(), - swa_evictable_tokens=tree_cache.swa_evictable_size(), - ) - if reclaim_plan is None: - return - full_num_tokens, swa_num_tokens = reclaim_plan - else: - full_num_tokens = max(0, num_tokens - allocator.full_available_size()) - swa_num_tokens = max(0, num_tokens - allocator.swa_available_size()) + if supports_swa_byte_budget(allocator): + allocator.evict_to_free_tokens(tree_cache, num_tokens) + return + + full_num_tokens = max(0, num_tokens - allocator.full_available_size()) + swa_num_tokens = max(0, num_tokens - allocator.swa_available_size()) if full_num_tokens or swa_num_tokens: tree_cache.evict_for_alloc( EvictParams(num_tokens=full_num_tokens, swa_num_tokens=swa_num_tokens) ) - if is_unified_swa: - # Even a zero-reclaim plan can depend on compaction before callers - # inspect availability. - allocator.ensure_capacity(num_tokens, num_tokens) else: # Standard allocator: evict only the shortfall (mirrors the SWA arm) available_size = allocator.available_size() diff --git a/python/sglang/srt/mem_cache/kv_cache_configurator.py b/python/sglang/srt/mem_cache/kv_cache_configurator.py index 9b0daf6af1e9..908f43c9a754 100644 --- a/python/sglang/srt/mem_cache/kv_cache_configurator.py +++ b/python/sglang/srt/mem_cache/kv_cache_configurator.py @@ -290,6 +290,23 @@ def __post_init__(self) -> None: self.draft_model_idx in self.model_config.swa_attention_layer_ids ) + def hybrid_swa_token_capacity( + self, + *, + allocator: BaseTokenToKVPoolAllocator, + full_capacity: Optional[int], + swa_capacity: Optional[int], + ) -> int: + if get_memory().enable_unified_memory: + capacity = allocator.size_full + max_total_tokens = get_schedule().max_total_tokens + return ( + min(capacity, max_total_tokens) + if max_total_tokens is not None + else capacity + ) + return full_capacity or swa_capacity + def _build_fp4_quant_method(self, *, num_layers: int): if not is_float4_e2m1fn_x2(self.kv_cache_dtype): return None @@ -2369,9 +2386,8 @@ def config_from_budget( f"{config.max_total_num_tokens}" ) if max_tokens != config.max_total_num_tokens: - # Token-capped re-derivation: the profiled budget no longer - # applies; the recalced config's unified_total_bytes stays None - # and the factories fall back to the token-count byte sum. + # Re-derive the capped budget: SWA carries unified_memory_pool_bytes; + # Mamba factories fall back to token-count sizing without unified_total_bytes. config = configurator.calculate_pool_sizes_from_max_tokens( max_tokens, get_schedule().page_size ) diff --git a/python/sglang/srt/mem_cache/unified_memory_pool.py b/python/sglang/srt/mem_cache/unified_memory_pool.py index 149edbe28017..1cb9f2402575 100644 --- a/python/sglang/srt/mem_cache/unified_memory_pool.py +++ b/python/sglang/srt/mem_cache/unified_memory_pool.py @@ -1720,7 +1720,6 @@ def init_unified_swa_pools( need_sort: bool, forward_stream: Optional[torch.cuda.Stream] = None, lazy_compaction: bool = False, - unified_total_bytes: Optional[int] = None, model_context_len: Optional[int] = None, sliding_window_size: Optional[int] = None, ) -> UnifiedSWAPoolBundle: @@ -1759,11 +1758,6 @@ def init_unified_swa_pools( store_dtype=store_dtype, grow_direction="up", ) - if total_bytes is not None and unified_total_bytes is not None: - raise ValueError("total_bytes and unified_total_bytes are mutually exclusive") - if total_bytes is None: - total_bytes = unified_total_bytes - legacy_allocator_capacities = {} if total_bytes is None: if full_max_total_num_tokens is None or swa_max_total_num_tokens is None: diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 021ee586d579..31acf5902f59 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -1373,12 +1373,12 @@ def unload_lora_adapter(self, lora_ref: LoRARef): @property def effective_max_total_num_tokens(self): """Return the max token pool size considering hybrid swa settings.""" - if self.is_hybrid_swa and get_memory().enable_unified_memory: - capacity = self.token_to_kv_pool_allocator.size_full - if get_schedule().max_total_tokens is not None: - capacity = min(capacity, get_schedule().max_total_tokens) - elif self.is_hybrid_swa: - capacity = self.full_max_total_num_tokens or self.swa_max_total_num_tokens + if self.is_hybrid_swa: + capacity = self.kv_cache_configurator.hybrid_swa_token_capacity( + allocator=self.token_to_kv_pool_allocator, + full_capacity=self.full_max_total_num_tokens, + swa_capacity=self.swa_max_total_num_tokens, + ) else: capacity = self.max_total_num_tokens if (req_to_token_pool := getattr(self, "req_to_token_pool", None)) is not None: diff --git a/test/registered/unit/mem_cache/test_hisparse_max_token_pool_size.py b/test/registered/unit/mem_cache/test_hisparse_max_token_pool_size.py index 4c590ce02bf1..5a7222b3872a 100644 --- a/test/registered/unit/mem_cache/test_hisparse_max_token_pool_size.py +++ b/test/registered/unit/mem_cache/test_hisparse_max_token_pool_size.py @@ -14,6 +14,7 @@ from unittest.mock import MagicMock from sglang.srt.disaggregation.decode import DecodePreallocQueue +from sglang.srt.mem_cache.kv_cache_configurator import KVCacheConfigurator from sglang.srt.model_executor.model_runner import ModelRunner from sglang.srt.runtime_context import get_context from sglang.test.ci.ci_register import register_cpu_ci @@ -29,6 +30,7 @@ def _make_model_runner(**attrs): raise AttributeError on the internal `self.effective_max_total_num_tokens` read inside `max_token_pool_size`.""" instance = object.__new__(ModelRunner) + instance.kv_cache_configurator = object.__new__(KVCacheConfigurator) for name, value in attrs.items(): setattr(instance, name, value) return instance diff --git a/test/registered/unit/mem_cache/test_multi_ended_allocator.py b/test/registered/unit/mem_cache/test_multi_ended_allocator.py index 6f677c94f64e..aaed067bf844 100644 --- a/test/registered/unit/mem_cache/test_multi_ended_allocator.py +++ b/test/registered/unit/mem_cache/test_multi_ended_allocator.py @@ -727,7 +727,7 @@ def test_empty_pool_reservation_matches_packed_byte_boundary(self): max_new_tokens, sliding_window_size=16, page_size=page_size, - include_allocated_tail=False, + allocation_limit=16, ) reservation_swa_with_tail = estimate_swa_kv_tokens( extend_tokens, diff --git a/test/registered/unit/mem_cache/test_unified_byte_budget_sizing.py b/test/registered/unit/mem_cache/test_unified_byte_budget_sizing.py index f43ed5be8d77..d90ca5c45c17 100644 --- a/test/registered/unit/mem_cache/test_unified_byte_budget_sizing.py +++ b/test/registered/unit/mem_cache/test_unified_byte_budget_sizing.py @@ -13,9 +13,9 @@ # ============================================================================== """Byte-budget buffer sizing for the unified 2-pool factories. -With ``unified_total_bytes`` set the buffer is that many bytes exactly (the -mamba pair adds the state pool's bytes on top -- the budget is captured AFTER -the state carve-out); without it, sizing falls back to the token-count re-sum. +The SWA factory's ``total_bytes`` sets the exact buffer size. The Mamba pair's +``unified_total_bytes`` adds state bytes because its budget excludes that carve-out. +Without an explicit budget, sizing falls back to the token-count re-sum. Sizing from the ratio-derived token counts instead would re-introduce the configurator's rounding, which floors by the cell size and then page-aligns EACH side, losing up to about a page of tokens per side. @@ -91,7 +91,7 @@ def test_swa_factory_honors_the_budget_exactly(self): (64 + 32) * e + (e - 2), # almost one more entry ): with self.subTest(budget=budget): - bundle = _swa_factory(unified_total_bytes=budget) + bundle = _swa_factory(total_bytes=budget) self.assertEqual(bundle.unified_memory_pool.total_bytes, budget) def test_fallback_is_the_token_count_resum(self): @@ -166,7 +166,7 @@ def test_infeasible_budget_raises_before_construction(self): must fail loud instead of livelocking later.""" with self.assertRaises(RuntimeError) as ctx: _swa_factory( - unified_total_bytes=8 * _entry_bytes(), + total_bytes=8 * _entry_bytes(), model_context_len=4096, sliding_window_size=4096, ) @@ -184,7 +184,7 @@ def test_feasible_floor_inputs_are_not_rejected(self): ): with self.subTest(case=case): bundle = _swa_factory( - unified_total_bytes=200 * e, + total_bytes=200 * e, model_context_len=model_context_len, sliding_window_size=sliding_window_size, ) diff --git a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py index ef79f571ee96..791aa4065aa5 100644 --- a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py +++ b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py @@ -4167,7 +4167,7 @@ def test_buffer_only_anchor_lock_cap_clamped_by_context_headroom(self): def test_buffer_load_back_swa_window_charged_at_admission(self): """Admission contract: a request the SWA budget gate accepts must be - allocatable at batch time (_swa_reserved_tokens: "an admitted request + allocatable at batch time (estimate_swa_kv_tokens: "an admitted request cannot OOM"). Regression: buffer mode surfaced a staged prefetch as host_hit_length only, so the gate never charged the SWA window that consumption (init_load_back -> cc.load) allocates and the request From 0a80658c029fae8a0e9a9f5adc7742e76eb5b494 Mon Sep 17 00:00:00 2001 From: Cheng Wan Date: Sun, 13 Sep 2026 16:08:54 -0700 Subject: [PATCH 48/55] Refactor SWA admission into allocator-owned memory budgets --- python/sglang/srt/managers/schedule_policy.py | 393 ++--------------- python/sglang/srt/managers/scheduler.py | 61 +-- .../pool_stats_observer.py | 28 +- python/sglang/srt/mem_cache/allocator/base.py | 38 +- .../srt/mem_cache/allocator/hisparse.py | 14 + python/sglang/srt/mem_cache/allocator/swa.py | 35 ++ .../mem_cache/allocator/unified_hybrid_swa.py | 91 +++- python/sglang/srt/mem_cache/common.py | 35 +- python/sglang/srt/mem_cache/prefill_budget.py | 405 ++++++++++++++++++ .../unit/managers/test_prefill_adder.py | 176 ++++++-- .../test_scheduler_init_req_max_new_tokens.py | 12 +- .../mem_cache/test_multi_ended_allocator.py | 2 +- .../mem_cache/test_prefill_memory_budget.py | 224 ++++++++++ .../test_unified_radix_allocation_eviction.py | 6 +- 14 files changed, 1011 insertions(+), 509 deletions(-) create mode 100644 python/sglang/srt/mem_cache/prefill_budget.py create mode 100644 test/registered/unit/mem_cache/test_prefill_memory_budget.py diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index 20c54e3d2560..6978b369a75d 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -48,13 +48,10 @@ DeepSeekV4HiSparseTokenToKVPoolAllocator, ) from sglang.srt.mem_cache.allocator.swa import ( - PureSWATokenToKVPoolAllocator, SWATokenToKVPoolAllocator, - is_swa_req_ring, ) from sglang.srt.mem_cache.allocator.unified_hybrid_swa import ( UnifiedMambaSWATokenToKVPoolAllocator, - supports_swa_byte_budget, ) from sglang.srt.mem_cache.allocator.unified_mamba import ( UnifiedMambaTokenToKVPoolAllocator, @@ -80,37 +77,6 @@ ) -def estimate_swa_kv_tokens( - extend_input_len: int, - max_new_tokens: int, - *, - sliding_window_size: Optional[int], - page_size: int, - allocation_limit: Optional[int] = None, - host_hit_length: int = 0, -) -> int: - """Peak SWA reservation for one prefill/decode request.""" - if sliding_window_size is None or sliding_window_size <= 0: - reserved = extend_input_len + max_new_tokens + page_size - else: - allocated = ( - extend_input_len - if allocation_limit is None - else min(extend_input_len, allocation_limit) - ) - allocated_tail = max(allocated - sliding_window_size, 0) - # A full window would double-charge short cached-prefix resumes; - # including extend keeps the reservation above the prefill allocation. - reserved = ( - allocated_tail - + min(extend_input_len + max_new_tokens, sliding_window_size) - + page_size - ) - if host_hit_length > 0: - reserved += -(-host_hit_length // page_size) * page_size - return reserved - - # Threshold for in-batch prefix cache. # If a request has a matched prefix length (against existing cache) less than this value, # the scheduler runs the in-batch prefix caching check for this request. @@ -591,8 +557,6 @@ def __init__( self.prefill_tile_block_m = prefill_tile_block_m self.tree_cache = tree_cache self.token_to_kv_pool_allocator = token_to_kv_pool_allocator - # Per-request SWA ring: one fixed slot per request, not a token budget. - self._swa_req_ring = is_swa_req_ring(token_to_kv_pool_allocator) self.running_batch = running_batch self.new_token_ratio = new_token_ratio self.rem_input_tokens = rem_input_tokens - num_mixed_decode_tokens @@ -604,8 +568,9 @@ def __init__( if self.rem_chunk_tokens is not None: self.rem_chunk_tokens -= num_mixed_decode_tokens - self.rem_total_token_offset = num_mixed_decode_tokens - self.cur_rem_token_offset = num_mixed_decode_tokens + self.memory_budget = token_to_kv_pool_allocator.create_prefill_budget( + tree_cache, num_mixed_decode_tokens=num_mixed_decode_tokens + ) self.req_states = None self.can_run_list = [] @@ -621,7 +586,7 @@ def __init__( if running_batch is not None: # Estimate the offset in the remaining token space - self.rem_total_token_offset += sum( + self.memory_budget.total_offset += sum( [ self._get_running_request_total_token_offset(r) for r in running_batch.reqs @@ -634,16 +599,7 @@ def __init__( self.token_to_kv_pool_allocator, (SWATokenToKVPoolAllocator, DeepSeekV4HiSparseTokenToKVPoolAllocator), ) - self.is_all_swa = isinstance( - self.token_to_kv_pool_allocator, PureSWATokenToKVPoolAllocator - ) self.is_hybrid_ssm_cache = self.tree_cache.supports_mamba() - self.is_unified_swa = supports_swa_byte_budget(self.token_to_kv_pool_allocator) - - self.rem_swa_token_offset = ( - num_mixed_decode_tokens if self.is_unified_swa else 0 - ) - # A new state slot eats shared-gap bytes that `rem_total_tokens` counts # as free, so reserve per slot or admission over-commits. Gate on the # ALLOCATOR, not `is_hybrid_ssm_cache`: that is False for `ChunkCache`, @@ -730,95 +686,11 @@ def _get_running_request_total_token_offset(self, req: Req) -> int: @property def rem_total_tokens(self): - if self.is_all_swa: - available_and_evictable = ( - self.token_to_kv_pool_allocator.swa_available_size() - + self.tree_cache.swa_evictable_size() - ) - elif self.is_hybrid_swa: - available_and_evictable = ( - self.token_to_kv_pool_allocator.full_available_size() - + self.tree_cache.full_evictable_size() - ) - elif self.is_hybrid_ssm_cache: - available_and_evictable = ( - self.token_to_kv_pool_allocator.available_size() - + self.tree_cache.full_evictable_size() - ) - else: - available_and_evictable = ( - self.token_to_kv_pool_allocator.available_size() - + self.tree_cache.evictable_size() - ) - return available_and_evictable - self.rem_total_token_offset - - @property - def rem_swa_tokens(self): - allocator = self.token_to_kv_pool_allocator - if self._swa_req_ring: - # swa_available_size() already reports ring capacity; tree - # swa_evictable is in linear token units and frees no ring space. - return allocator.swa_available_size() - self.rem_swa_token_offset - return ( - allocator.swa_available_size() - + self.tree_cache.swa_evictable_size() - - self.rem_swa_token_offset - ) + return self.memory_budget.remaining_total @property def cur_rem_tokens(self): - if self.is_all_swa: - available_and_evictable = ( - self.token_to_kv_pool_allocator.swa_available_size() - + self.tree_cache.swa_evictable_size() - ) - elif self.is_hybrid_swa: - available_and_evictable = ( - self.token_to_kv_pool_allocator.full_available_size() - + self.tree_cache.full_evictable_size() - ) - elif self.is_hybrid_ssm_cache: - available_and_evictable = ( - self.token_to_kv_pool_allocator.available_size() - + self.tree_cache.full_evictable_size() - ) - else: - available_and_evictable = ( - self.token_to_kv_pool_allocator.available_size() - + self.tree_cache.evictable_size() - ) - - return available_and_evictable - self.cur_rem_token_offset - - def _swa_budget_for_req( - self, extend_input_len: int, max_new_tokens: int, swa_host_hit_length: int = 0 - ) -> int: - """SWA pool budget per request. Only valid when is_hybrid_swa is True. - - With chunked prefill + overlap scheduler, the peak SWA occupancy is: - chunk N (running, not yet in tree) + sliding window (locked in tree) - + chunk N+1 (new allocation) - Since chunk N and locked tokens are already excluded from - swa_available + swa_evictable, the budget only needs to cover the - chunk N+1 allocation plus decode headroom: - - budget = max(alloc - window, 0) + min(extend + max_new_tokens, window) + page - - where alloc = min(extend, rem_chunk); the min() cap keeps the two terms - from double-counting extend, so budget <= extend + max_new_tokens + page. - """ - allocator = self.token_to_kv_pool_allocator - if self._swa_req_ring: - # One ring slot per request, in the same unit as swa_available_size. - return allocator.swa_ring_cost_tokens - return estimate_swa_kv_tokens( - extend_input_len, - max_new_tokens, - sliding_window_size=self.tree_cache.sliding_window_size, - page_size=self.page_size, - allocation_limit=self.rem_chunk_tokens, - host_hit_length=swa_host_hit_length, - ) + return self.memory_budget.remaining_current def _swa_new_tokens(self, req: Req) -> int: """Tokens a request may still decode, for SWA headroom sizing. Mirrors @@ -830,150 +702,17 @@ def _swa_new_tokens(self, req: Req) -> int: CLIP_MAX_NEW_TOKENS, ) - def _unified_swa_reservation_fits( - self, - full_tokens: int, - swa_tokens: int, - *, - empty_pool: bool = False, - ) -> bool: - allocator = self.token_to_kv_pool_allocator - if empty_pool: - full_evictable_tokens = 0 - swa_evictable_tokens = 0 - else: - full_tokens += self.rem_total_token_offset - swa_tokens += self.rem_swa_token_offset - full_evictable_tokens = self.tree_cache.full_evictable_size() - swa_evictable_tokens = self.tree_cache.swa_evictable_size() - return allocator.can_reserve( - full_tokens, - swa_tokens, - full_evictable_tokens=full_evictable_tokens, - swa_evictable_tokens=swa_evictable_tokens, - empty_pool=empty_pool, - require_token_slack=empty_pool, - ) - - def _unified_swa_chunk_cap( - self, - max_new_tokens: int, - swa_host_hit_length: int = 0, - *, - max_chunk_tokens: int, - ) -> int: - """Largest page-aligned chunk that fits the joint FULL/SWA byte budget.""" - cap = max_chunk_tokens // self.page_size * self.page_size - lo, hi = 0, cap // self.page_size - while lo < hi: - mid = (lo + hi + 1) // 2 - chunk_tokens = mid * self.page_size - full_tokens = self.ceil_paged_tokens(chunk_tokens) + self.page_size - swa_tokens = self._swa_budget_for_req( - chunk_tokens, max_new_tokens, swa_host_hit_length - ) - if self._unified_swa_reservation_fits(full_tokens, swa_tokens): - lo = mid - else: - hi = mid - 1 - return lo * self.page_size - - def _swa_chunk_cap(self, max_new_tokens: int, swa_host_hit_length: int = 0) -> int: - """Largest chunk fitting the fixed SWA partition after decode headroom.""" - # Only the sliding window stays locked across chunks, so the next - # chunk can reuse the transient space outside its decode headroom. - cap = int(self.rem_swa_tokens) - estimate_swa_kv_tokens( - 0, - max_new_tokens, - sliding_window_size=self.tree_cache.sliding_window_size, - page_size=self.page_size, - host_hit_length=swa_host_hit_length, - ) - return max(0, cap // self.page_size * self.page_size) - - def _swa_req_never_fits( - self, - extend_input_len: int, - max_new_tokens: int, - swa_host_hit_length: int = 0, - ) -> bool: - """True when a request's SWA budget exceeds the *entire* SWA pool, so it - can never be admitted whole no matter how far the pool drains. - - This is the head-of-line livelock the _swa_chunk_cap escape hatch exists - for; the hatch must fire only in this case. A request that merely - exceeds *current* rem_swa (transient pressure) would fit once running - decodes free their windows, so it must wait — admitting it into the - decode headroom collapses the SWA evictable cushion and forces running - requests to retract (observed as a severe retraction/re-prefill storm on - hybrid-SWA models at high concurrency).""" - swa_tokens = self._swa_budget_for_req( - extend_input_len, max_new_tokens, swa_host_hit_length - ) - if self._swa_req_ring: - return swa_tokens > self.token_to_kv_pool_allocator.size_swa - return swa_tokens >= self.token_to_kv_pool_allocator.size_swa - - def _check_swa_admission( - self, - req: Req, - *, - extend_input_len: int, - total_tokens: int, - ) -> tuple[bool, Optional[int]]: - if not self.is_unified_swa: - extend_input_len = self.ceil_paged_tokens(extend_input_len) - max_new_tokens = self._swa_new_tokens(req) - swa_needed = self._swa_budget_for_req( - extend_input_len, - max_new_tokens, - swa_host_hit_length=req.swa_host_hit_length, - ) - swa_cap = 0 - # Only permanent shortfalls may shrink a chunk; transient pressure waits. - if self.is_unified_swa: - fits = self._unified_swa_reservation_fits(total_tokens, swa_needed) - if not fits: - full_input_tokens = len(req.full_untruncated_fill_ids) - if not self._unified_swa_reservation_fits( - full_input_tokens + max_new_tokens + self.page_size, - self._swa_budget_for_req(full_input_tokens, max_new_tokens), - empty_pool=True, - ): - swa_cap = self._unified_swa_chunk_cap( - 0, - req.swa_host_hit_length, - max_chunk_tokens=min( - self.rem_chunk_tokens or 0, max(0, extend_input_len - 1) - ), - ) - else: - # Ring-slot capacity is exact; the fixed token partition keeps slack. - fits = ( - swa_needed <= self.rem_swa_tokens - if self._swa_req_ring - else swa_needed < self.rem_swa_tokens - ) - if not fits and self._swa_req_never_fits( - extend_input_len, max_new_tokens, req.swa_host_hit_length - ): - swa_cap = self._swa_chunk_cap(max_new_tokens, req.swa_host_hit_length) - if fits: - return True, self.rem_chunk_tokens - if self.rem_chunk_tokens is None or swa_cap <= 0: - return False, None - return True, min(self.rem_chunk_tokens, swa_cap) - def _check_prefill_budget( self, req: Req, *, extend_input_len: int, total_tokens: int ) -> tuple[bool, Optional[int]]: - if not self.is_unified_swa and total_tokens >= self.rem_total_tokens: - return False, None - if self.is_hybrid_swa: - return self._check_swa_admission( - req, extend_input_len=extend_input_len, total_tokens=total_tokens - ) - return True, self.rem_chunk_tokens + return self.memory_budget.check_prefill( + extend_input_len=extend_input_len, + total_tokens=total_tokens, + max_new_tokens=self._swa_new_tokens(req), + input_tokens=len(req.full_untruncated_fill_ids), + swa_host_hit_length=req.swa_host_hit_length, + chunk_limit=self.rem_chunk_tokens, + ) def _mamba_gap_budget_for_req(self, req: Req) -> int: """Shared-gap reservation (full-token-equivalents) for a request's new @@ -994,12 +733,7 @@ def ceil_paged_tokens(self, tokens: int) -> int: return -(-tokens // self.page_size) * self.page_size def budget_state(self): - if self.is_unified_swa: - no_token = not self._unified_swa_reservation_fits(0, 0) - else: - no_token = self.rem_total_tokens <= 0 or self.cur_rem_tokens <= 0 - if not no_token and self.is_hybrid_swa: - no_token = self.rem_swa_tokens <= 0 + no_token = not self.memory_budget.has_capacity() # Gate new mamba slots separately: rem_total_tokens' full_evictable can't # cover a mamba slot, which needs mamba-recoverable bytes (see __init__). if not no_token and self.rem_mamba_slots is not None: @@ -1032,17 +766,12 @@ def _update_prefill_budget( raw_extend_input_len = extend_input_len extend_input_len = self.ceil_paged_tokens(extend_input_len) - # alloc_extend reserves an extra page_size per request to make sure the budget doesn't over-commit - page_overhead = self.page_size - # `mamba_gap_reserve` (shared Mamba pool only; 0 otherwise) charges the new - # mamba state's shared-gap cost to BOTH full budgets: the slot is allocated - # immediately (counts against `cur_rem`) and held for the request lifetime - # (counts against `rem_total`). See `_mamba_gap_budget_for_req`. - self.rem_total_token_offset += ( - extend_input_len + max_new_tokens + page_overhead + mamba_gap_reserve - ) - self.cur_rem_token_offset += ( - extend_input_len + page_overhead + mamba_gap_reserve + self.memory_budget.reserve( + extend_input_len, + max_new_tokens, + extra_tokens=mamba_gap_reserve, + chunk_limit=self.rem_chunk_tokens, + is_chunked_continuation=is_chunked_continuation, ) # The new mamba slot also consumes one mamba-recoverable slot (gated # separately so full_evictable can't cover it — see __init__). @@ -1050,14 +779,6 @@ def _update_prefill_budget( self.rem_mamba_slots -= 1 self.rem_input_tokens -= extend_input_len - if self.is_hybrid_swa: - # The ring slot is reserved once at first admission; charging it - # again on a continuation would double-count and over-throttle. - if not (self._swa_req_ring and is_chunked_continuation): - self.rem_swa_token_offset += self._swa_budget_for_req( - extend_input_len, max_new_tokens - ) - if self.dllm_config is not None: self.rem_dllm_tokens -= extend_input_len elif self.rem_chunk_tokens is not None: @@ -1182,44 +903,15 @@ def add_dllm_staging_req(self, req: Req): else AddReqResult.CONTINUE ) - def _unified_swa_chunk_limit( - self, req: Req, extend_input_len: int, max_chunk_tokens: int - ) -> Optional[int]: - candidate_tokens = min(extend_input_len, max_chunk_tokens) - candidate_finishes = candidate_tokens >= extend_input_len - max_new_tokens = self._swa_new_tokens(req) if candidate_finishes else 0 - if self._unified_swa_reservation_fits( - self.ceil_paged_tokens(candidate_tokens) + max_new_tokens + self.page_size, - self._swa_budget_for_req(candidate_tokens, max_new_tokens), - ): - return max_chunk_tokens - max_partial = ( - max(0, candidate_tokens - 1) if candidate_finishes else candidate_tokens - ) - joint_cap = self._unified_swa_chunk_cap(0, max_chunk_tokens=max_partial) - if joint_cap <= 0: - return None - return min(max_chunk_tokens, joint_cap) - def add_chunked_req(self, req: Req): if self.dllm_config is not None: _rem_tokens = self._get_dllm_remain_tokens() else: - if self.is_unified_swa: - _rem_tokens = self.rem_chunk_tokens - else: - _rem_tokens = min(self.rem_chunk_tokens, int(self.rem_total_tokens)) - if self.is_hybrid_swa and not self._swa_req_ring: - # alloc_extend needs one extra page; ring mode counts slots. - _rem_tokens = min( - _rem_tokens, int(self.rem_swa_tokens) - self.page_size - ) - # The chunked_req must be added to the list; otherwise, it will cause a memory leak. - # Therefore, in certain cases where _rem_tokens <= 0, it should be replaced with rem_chunk_tokens. - if _rem_tokens <= 0: - if self.is_hybrid_swa: - return req - _rem_tokens = self.rem_chunk_tokens + _rem_tokens = self.memory_budget.available_chunk_tokens( + self.rem_chunk_tokens + ) + if _rem_tokens is None: + return req # A mid-chunk rank prefills this pass regardless of the delayer # verdict, so report prefillable=True and ignore the result. @@ -1235,12 +927,13 @@ def add_chunked_req(self, req: Req): cand_extend_input_len = len(req.full_untruncated_fill_ids) - len( req.prefix_indices ) - if self.is_unified_swa: - _rem_tokens = self._unified_swa_chunk_limit( - req, cand_extend_input_len, _rem_tokens - ) - if _rem_tokens is None: - return req + _rem_tokens = self.memory_budget.fit_chunk( + extend_input_len=cand_extend_input_len, + max_new_tokens=self._swa_new_tokens(req), + chunk_limit=_rem_tokens, + ) + if _rem_tokens is None: + return req truncated = cand_extend_input_len > _rem_tokens new_len = min(cand_extend_input_len, _rem_tokens) req.set_extend_range(len(req.prefix_indices), len(req.prefix_indices) + new_len) @@ -1285,18 +978,12 @@ def add_one_req_ignore_eos(self, req: Req): # Shared Mamba pool: fold the new mamba state's shared-gap cost into the # budget gate so admission can't over-commit (0 for baseline / non-Mamba). paged_input += self._mamba_gap_budget_for_req(req) - if self.is_hybrid_swa: - max_new_tokens = self._swa_new_tokens(req) - swa_needed = self._swa_budget_for_req(cand_extend_input_len, max_new_tokens) - if self.is_unified_swa: - fits = self._unified_swa_reservation_fits( - paged_input + max_new_tokens + self.page_size, - swa_needed, - ) - else: - fits = paged_input <= min(self.cur_rem_tokens, self.rem_total_tokens) and ( - not self.is_hybrid_swa or swa_needed <= self.rem_swa_tokens - ) + fits = self.memory_budget.can_allocate_prefill( + paged_input=paged_input, + extend_input_len=cand_extend_input_len, + max_new_tokens=self._swa_new_tokens(req), + chunk_limit=self.rem_chunk_tokens, + ) if not fits: return AddReqResult.NO_TOKEN @@ -1666,7 +1353,7 @@ def preempt_to_schedule(self, req: Req) -> bool: release_counter = 0 for i, running_req in enumerate(self.running_batch.reqs): if running_req in preemptible_reqs: - self.rem_total_token_offset -= ( + self.memory_budget.total_offset -= ( self._get_running_request_total_token_offset(running_req) ) release_counter += 1 diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 2a67ff080fad..de2fbbe99cec 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -212,7 +212,6 @@ AddReqResult, PrefillAdder, SchedulePolicy, - estimate_swa_kv_tokens, ) from sglang.srt.managers.scheduler_components.batch_result_processor import ( SchedulerBatchResultProcessor, @@ -280,9 +279,6 @@ validate_input_length, ) from sglang.srt.mem_cache import kv_cache_builder -from sglang.srt.mem_cache.allocator.unified_hybrid_swa import ( - supports_swa_byte_budget, -) from sglang.srt.mem_cache.common import ( maybe_cache_unfinished_req, release_kv_cache, @@ -2531,18 +2527,18 @@ def init_req_max_new_tokens(self, req): self.max_req_len - input_len - 1, ), ) - if supports_swa_byte_budget(self.token_to_kv_pool_allocator): - max_new_tokens = self._clip_unified_swa_max_new_tokens(req, max_new_tokens) - else: - # Non-unified pools keep their fixed token partition. - paged_input_len = -(-input_len // self.page_size) * self.page_size - max_new_tokens = min( - max_new_tokens, - self.max_total_num_tokens * get_parallel().attn_dcp_size - - paged_input_len - - self.page_size - - 1, + max_new_tokens = self.token_to_kv_pool_allocator.max_new_tokens_for_memory( + input_len, + max_new_tokens, + token_capacity=self.max_total_num_tokens * get_parallel().attn_dcp_size, + sliding_window_size=self.sliding_window_size, + chunk_size=self.chunked_prefill_size, + ) + if max_new_tokens is None: + req.set_finish_with_abort( + f"Request prompt exceeds the KV memory budget: input_len={input_len}." ) + max_new_tokens = 0 req.sampling_params.max_new_tokens = max(0, max_new_tokens) # Clipping above can push max_new_tokens below min_new_tokens, which @@ -2550,41 +2546,6 @@ def init_req_max_new_tokens(self, req): if req.sampling_params.min_new_tokens > req.sampling_params.max_new_tokens: req.sampling_params.min_new_tokens = req.sampling_params.max_new_tokens - def _clip_unified_swa_max_new_tokens(self, req: Req, max_new_tokens: int) -> int: - input_len = len(req.origin_input_ids) - - def fits_shared_pool(candidate: int) -> bool: - swa_tokens = estimate_swa_kv_tokens( - input_len, - candidate, - sliding_window_size=self.sliding_window_size, - page_size=self.page_size, - allocation_limit=self.chunked_prefill_size, - ) - return self.token_to_kv_pool_allocator.can_reserve( - input_len + candidate + self.page_size, - swa_tokens, - empty_pool=True, - require_token_slack=True, - ) - - if not fits_shared_pool(0): - req.set_finish_with_abort( - "Request prompt exceeds the unified FULL/SWA KV byte budget: " - f"input_len={input_len}." - ) - return 0 - if fits_shared_pool(max_new_tokens): - return max_new_tokens - lo, hi = 0, max_new_tokens - while lo < hi: - mid = (lo + hi + 1) // 2 - if fits_shared_pool(mid): - lo = mid - else: - hi = mid - 1 - return lo - def _process_and_broadcast_mm_inputs( self, raw_mm_inputs, diff --git a/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py b/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py index 823fe15b67cf..f061506f1182 100644 --- a/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py +++ b/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py @@ -12,10 +12,6 @@ ) from sglang.srt.mem_cache.allocator.swa import is_swa_req_ring -from sglang.srt.mem_cache.allocator.unified_hybrid_swa import ( - UnifiedMambaSWATokenToKVPoolAllocator, - supports_swa_byte_budget, -) if TYPE_CHECKING: from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator @@ -291,28 +287,12 @@ def _get_mamba_token_info(self): mamba_evictable_size=mamba_evictable_size, ) - def _swa_capacity_and_available(self) -> tuple[tuple[int, int], tuple[int, int]]: - allocator = self.token_to_kv_pool_allocator - if supports_swa_byte_budget(allocator): - return ( - (allocator.current_full_capacity, allocator.full_available_size()), - (allocator.current_swa_capacity, allocator.swa_available_size()), - ) - if isinstance(allocator, UnifiedMambaSWATokenToKVPoolAllocator): - # The tri-pool reports static capacities paired with conserve views. - full_available_size = allocator.conserve_full_available_size() - swa_available_size = allocator.conserve_swa_available_size() - else: - full_available_size = allocator.full_available_size() - swa_available_size = allocator.swa_available_size() - return ( - (self.full_tokens_per_layer, full_available_size), - (self.swa_tokens_per_layer, swa_available_size), - ) - def _get_swa_token_info(self) -> PoolStats: (full_capacity, full_available_size), (swa_capacity, swa_available_size) = ( - self._swa_capacity_and_available() + self.token_to_kv_pool_allocator.swa_capacity_and_available( + full_capacity=self.full_tokens_per_layer, + swa_capacity=self.swa_tokens_per_layer, + ) ) full_evictable_size = self.tree_cache.full_evictable_size() swa_evictable_size = self.tree_cache.swa_evictable_size() diff --git a/python/sglang/srt/mem_cache/allocator/base.py b/python/sglang/srt/mem_cache/allocator/base.py index 68ecb5e18cec..639b1cbeafcf 100644 --- a/python/sglang/srt/mem_cache/allocator/base.py +++ b/python/sglang/srt/mem_cache/allocator/base.py @@ -69,12 +69,44 @@ def size_full(self): # The scheduler calls these unconditionally, with no allocator-type branches # on its side; byte-accounted composites override the token-count defaults. + def create_prefill_budget(self, tree_cache, *, num_mixed_decode_tokens=0): + from sglang.srt.mem_cache.prefill_budget import PrefillBudget + + return PrefillBudget( + self, tree_cache, num_mixed_decode_tokens=num_mixed_decode_tokens + ) + + def max_new_tokens_for_memory( + self, + input_tokens: int, + max_new_tokens: int, + *, + token_capacity: int, + sliding_window_size: int | None, + chunk_size: int | None, + ) -> int | None: + """Clip generation to the empty-pool budget; None means prompt cannot fit. + + token_capacity is the scheduler's configured capacity, including its + distributed token scaling. Shared pools use their physical byte layout. + """ + paged_input = -(-input_tokens // self.page_size) * self.page_size + return max( + 0, min(max_new_tokens, token_capacity - paged_input - self.page_size - 1) + ) + def evict_to_free_tokens(self, tree_cache, num_tokens: int) -> None: """Evict unlocked prefix-cache entries until this allocator can serve ``num_tokens`` or nothing evictable remains.""" - from sglang.srt.mem_cache.common import evict_from_tree_cache - - evict_from_tree_cache(tree_cache, num_tokens) + from sglang.srt.mem_cache.base_prefix_cache import EvictParams + from sglang.srt.mem_cache.common import _evict_until_allocatable + + if tree_cache is None or tree_cache.is_chunk_cache(): + return + shortfall = num_tokens - self.available_size() + if shortfall > 0: + tree_cache.evict_for_alloc(EvictParams(num_tokens=shortfall)) + _evict_until_allocatable(tree_cache, self, num_tokens) def check_decode_capacity( self, diff --git a/python/sglang/srt/mem_cache/allocator/hisparse.py b/python/sglang/srt/mem_cache/allocator/hisparse.py index 9dd096c3b317..e698a773d476 100644 --- a/python/sglang/srt/mem_cache/allocator/hisparse.py +++ b/python/sglang/srt/mem_cache/allocator/hisparse.py @@ -357,6 +357,20 @@ def translate_swa_indices_for_transfer( kv_indices ) + def create_prefill_budget(self, tree_cache, *, num_mixed_decode_tokens=0): + # Use this wrapper's capacity, which also includes the compressed pool. + from sglang.srt.mem_cache.prefill_budget import SWAPrefillBudget + + return SWAPrefillBudget( + self, tree_cache, num_mixed_decode_tokens=num_mixed_decode_tokens + ) + + def swa_capacity_and_available(self, *, full_capacity, swa_capacity): + return ( + (full_capacity, self.full_available_size()), + (swa_capacity, self.swa_available_size()), + ) + def full_available_size(self): return min( self.logical_attn_allocator.full_available_size(), diff --git a/python/sglang/srt/mem_cache/allocator/swa.py b/python/sglang/srt/mem_cache/allocator/swa.py index cfdf732fbb5f..7e573a7528a9 100644 --- a/python/sglang/srt/mem_cache/allocator/swa.py +++ b/python/sglang/srt/mem_cache/allocator/swa.py @@ -158,6 +158,31 @@ def available_size(self): self.swa_attn_allocator.available_size(), ) + def create_prefill_budget(self, tree_cache, *, num_mixed_decode_tokens=0): + from sglang.srt.mem_cache.prefill_budget import SWAPrefillBudget + + return SWAPrefillBudget( + self, tree_cache, num_mixed_decode_tokens=num_mixed_decode_tokens + ) + + def swa_capacity_and_available(self, *, full_capacity, swa_capacity): + return ( + (full_capacity, self.full_available_size()), + (swa_capacity, self.swa_available_size()), + ) + + def evict_to_free_tokens(self, tree_cache, num_tokens: int) -> None: + from sglang.srt.mem_cache.base_prefix_cache import EvictParams + + if tree_cache is None or tree_cache.is_chunk_cache(): + return + full_shortfall = max(0, num_tokens - self.full_available_size()) + swa_shortfall = max(0, num_tokens - self.swa_available_size()) + if full_shortfall or swa_shortfall: + tree_cache.evict_for_alloc( + EvictParams(num_tokens=full_shortfall, swa_num_tokens=swa_shortfall) + ) + def full_available_size(self): return self.full_attn_allocator.available_size() @@ -687,6 +712,16 @@ def __init__( def available_size(self): return self.swa_attn_allocator.available_size() + def create_prefill_budget(self, tree_cache, *, num_mixed_decode_tokens=0): + from sglang.srt.mem_cache.prefill_budget import SWAPrefillBudget + + return SWAPrefillBudget( + self, + tree_cache, + num_mixed_decode_tokens=num_mixed_decode_tokens, + all_swa=True, + ) + def full_available_size(self): return self.swa_attn_allocator.available_size() diff --git a/python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py b/python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py index 7248edf28b0f..c4d1c579ef39 100644 --- a/python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py +++ b/python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py @@ -18,7 +18,7 @@ import logging import math -from typing import Callable, List, Optional, Sequence, Tuple, TypeGuard +from typing import Callable, List, Optional, Sequence, Tuple import torch from torch.profiler import record_function @@ -43,16 +43,6 @@ logger = logging.getLogger(__name__) -def supports_swa_byte_budget( - allocator: BaseTokenToKVPoolAllocator | None, -) -> TypeGuard[UnifiedSWATokenToKVPoolAllocator]: - """Whether FULL/SWA demand can be checked against a two-pool byte budget.""" - return ( - isinstance(allocator, UnifiedSWATokenToKVPoolAllocator) - and allocator.supports_asymmetric_reservation - ) - - class UnifiedSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator): """Composite allocator for the hybrid SWA pair (full + swa MHA sub-pools). @@ -60,10 +50,6 @@ class UnifiedSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator): `available_size()` (joint bytes, in TOKENS) is the only safe alloc pre-check. """ - # Unequal FULL/SWA reservations share one byte budget. The tri-pool opts - # out because its Mamba state competes for those same bytes. - supports_asymmetric_reservation = True - # Parent's `size` property has no setter but base init does `self.size = size`; # the no-op setter below absorbs that write. @property @@ -532,6 +518,57 @@ def ensure_capacity(self, full_tokens: int, swa_tokens: int) -> bool: compacted=False, ) + def create_prefill_budget(self, tree_cache, *, num_mixed_decode_tokens=0): + from sglang.srt.mem_cache.prefill_budget import SharedSWAPrefillBudget + + return SharedSWAPrefillBudget( + self, tree_cache, num_mixed_decode_tokens=num_mixed_decode_tokens + ) + + def swa_capacity_and_available(self, *, full_capacity, swa_capacity): + return ( + (self.current_full_capacity, self.full_available_size()), + (self.current_swa_capacity, self.swa_available_size()), + ) + + def max_new_tokens_for_memory( + self, + input_tokens: int, + max_new_tokens: int, + *, + token_capacity: int, + sliding_window_size: int | None, + chunk_size: int | None, + ) -> int | None: + from sglang.srt.mem_cache.prefill_budget import estimate_swa_kv_tokens + + def fits(candidate): + return self.can_reserve( + input_tokens + candidate + self.page_size, + estimate_swa_kv_tokens( + input_tokens, + candidate, + sliding_window_size=sliding_window_size, + page_size=self.page_size, + allocation_limit=chunk_size, + ), + empty_pool=True, + require_token_slack=True, + ) + + if not fits(0): + return None + if fits(max_new_tokens): + return max_new_tokens + lo, hi = 0, max_new_tokens + while lo < hi: + mid = (lo + hi + 1) // 2 + if fits(mid): + lo = mid + else: + hi = mid - 1 + return lo + def evict_to_free_tokens(self, tree_cache, num_tokens: int) -> None: from sglang.srt.mem_cache.base_prefix_cache import EvictParams @@ -1086,8 +1123,6 @@ class UnifiedMambaSWATokenToKVPoolAllocator(UnifiedSWATokenToKVPoolAllocator): `UnifiedMambaSlotAllocator`. """ - supports_asymmetric_reservation = False - def __init__( self, *, @@ -1344,17 +1379,33 @@ def set_inflight_forward( super().set_inflight_forward(forward_done, out_cache_loc_virtual) self.mamba_allocator.set_inflight_forward(forward_done, None) + def create_prefill_budget(self, tree_cache, *, num_mixed_decode_tokens=0): + # Mamba competes for the shared gap too; retain the tri-pool's existing + # token and state-slot admission until it has a three-way reservation. + return SWATokenToKVPoolAllocator.create_prefill_budget( + self, tree_cache, num_mixed_decode_tokens=num_mixed_decode_tokens + ) + + def max_new_tokens_for_memory(self, *args, **kwargs): + return BaseTokenToKVPoolAllocator.max_new_tokens_for_memory( + self, *args, **kwargs + ) + + def swa_capacity_and_available(self, *, full_capacity, swa_capacity): + return ( + (full_capacity, self.conserve_full_available_size()), + (swa_capacity, self.conserve_swa_available_size()), + ) + def evict_to_free_tokens(self, tree_cache, num_tokens: int) -> None: """Joint-aware eviction: one tri-lifetime node frees bytes on several sides at once, so re-check the JOINT gate instead of the per-side shortfall.""" - from sglang.srt.mem_cache.common import evict_from_tree_cache - # Arbitrary retry bound; a round that frees nothing ends the loop anyway. for _ in range(4): before = self.available_size() if before >= num_tokens: return - evict_from_tree_cache(tree_cache, num_tokens) + SWATokenToKVPoolAllocator.evict_to_free_tokens(self, tree_cache, num_tokens) if self.available_size() <= before: return # no progress diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index 5f73d202a795..4810ad31bcab 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/python/sglang/srt/mem_cache/common.py @@ -11,7 +11,6 @@ ) from sglang.kernels.ops.memory.common import get_last_loc_kernel as get_last_loc_kernel from sglang.srt.mem_cache.allocator.page_interleave import page_interleave_shard_size -from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache, EvictParams from sglang.srt.mem_cache.hicache_storage import PoolTransfer from sglang.srt.mem_cache.memory_pool import HybridReqToTokenPool, ReqToTokenPool @@ -151,39 +150,11 @@ def maybe_cache_unfinished_req(req: Req, tree_cache: BasePrefixCache, **kwargs): def evict_from_tree_cache(tree_cache: BasePrefixCache | None, num_tokens: int): - if tree_cache is None: - return - - if tree_cache.is_chunk_cache(): - return - - allocator = tree_cache.token_to_kv_pool_allocator - - if isinstance(allocator, SWATokenToKVPoolAllocator): - # Hybrid allocator - from sglang.srt.mem_cache.allocator.unified_hybrid_swa import ( - supports_swa_byte_budget, + if tree_cache is not None and not tree_cache.is_chunk_cache(): + tree_cache.token_to_kv_pool_allocator.evict_to_free_tokens( + tree_cache, num_tokens ) - if supports_swa_byte_budget(allocator): - allocator.evict_to_free_tokens(tree_cache, num_tokens) - return - - full_num_tokens = max(0, num_tokens - allocator.full_available_size()) - swa_num_tokens = max(0, num_tokens - allocator.swa_available_size()) - if full_num_tokens or swa_num_tokens: - 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_for_alloc( - EvictParams(num_tokens=num_tokens - available_size) - ) - _evict_until_allocatable(tree_cache, allocator, num_tokens) - def _evict_until_allocatable( tree_cache: BasePrefixCache, allocator, num_tokens: int diff --git a/python/sglang/srt/mem_cache/prefill_budget.py b/python/sglang/srt/mem_cache/prefill_budget.py new file mode 100644 index 000000000000..0811eb763354 --- /dev/null +++ b/python/sglang/srt/mem_cache/prefill_budget.py @@ -0,0 +1,405 @@ +# Copyright 2026 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Memory reservations for one prefill pass. + +The scheduler supplies token demand and its chunk/decode limits. These objects +account for admitted but not yet allocated work and query live cache capacity: +locking a prefix or preempting a request must affect the next admission check. +They neither select requests nor mutate the prefix cache or allocator. +""" + +from typing import Optional + + +def estimate_swa_kv_tokens( + extend_input_len: int, + max_new_tokens: int, + *, + sliding_window_size: Optional[int], + page_size: int, + allocation_limit: Optional[int] = None, + host_hit_length: int = 0, +) -> int: + """Peak SWA reservation for one prefill/decode request.""" + if sliding_window_size is None or sliding_window_size <= 0: + reserved = extend_input_len + max_new_tokens + page_size + else: + allocated = ( + extend_input_len + if allocation_limit is None + else min(extend_input_len, allocation_limit) + ) + allocated_tail = max(allocated - sliding_window_size, 0) + # A full window would double-charge short cached-prefix resumes; + # including extend keeps the reservation above the prefill allocation. + reserved = ( + allocated_tail + + min(extend_input_len + max_new_tokens, sliding_window_size) + + page_size + ) + if host_hit_length > 0: + reserved += -(-host_hit_length // page_size) * page_size + return reserved + + +class PrefillBudget: + """Fixed token pool. Offsets include pending allocations and decode headroom.""" + + def __init__(self, allocator, tree_cache, *, num_mixed_decode_tokens: int = 0): + self.allocator = allocator + self.tree_cache = tree_cache + self.page_size = allocator.page_size + self.total_offset = num_mixed_decode_tokens + self.current_offset = num_mixed_decode_tokens + self.swa_offset = 0 + + def ceil_paged_tokens(self, tokens: int) -> int: + return -(-tokens // self.page_size) * self.page_size + + def _available_and_evictable(self): + evictable = ( + self.tree_cache.full_evictable_size() + if self.tree_cache.supports_mamba() + else self.tree_cache.evictable_size() + ) + return self.allocator.available_size() + evictable + + @property + def remaining_total(self): + return self._available_and_evictable() - self.total_offset + + @property + def remaining_current(self): + return self._available_and_evictable() - self.current_offset + + @property + def remaining_swa(self): + return 0 + + def has_capacity(self) -> bool: + return self.remaining_total > 0 and self.remaining_current > 0 + + def check_prefill( + self, + *, + extend_input_len: int, + total_tokens: int, + max_new_tokens: int, + input_tokens: int, + swa_host_hit_length: int, + chunk_limit: int | None, + ) -> tuple[bool, int | None]: + """Return admission feasibility and a memory bound on the requested chunk.""" + return ( + (True, chunk_limit) + if total_tokens < self.remaining_total + else (False, None) + ) + + def can_allocate_prefill( + self, + *, + paged_input: int, + extend_input_len: int, + max_new_tokens: int, + chunk_limit: int | None, + ) -> bool: + return paged_input <= min(self.remaining_current, self.remaining_total) + + def available_chunk_tokens(self, chunk_limit: int) -> int | None: + available = min(chunk_limit, int(self.remaining_total)) + # Single-pool continuation must make progress to release its KV. + return available if available > 0 else chunk_limit + + def fit_chunk( + self, + *, + extend_input_len: int, + max_new_tokens: int, + chunk_limit: int, + ) -> int | None: + return chunk_limit + + def reserve( + self, + extend_input_len: int, + max_new_tokens: int, + *, + extra_tokens: int = 0, + chunk_limit: int | None = None, + is_chunked_continuation: bool = False, + ) -> None: + extend_input_len = self.ceil_paged_tokens(extend_input_len) + immediate = extend_input_len + self.page_size + extra_tokens + self.total_offset += immediate + max_new_tokens + self.current_offset += immediate + + +class SWAPrefillBudget(PrefillBudget): + """Separate FULL/SWA partitions, including per-request SWA rings.""" + + def __init__(self, *args, all_swa=False, **kwargs): + super().__init__(*args, **kwargs) + self.all_swa = all_swa + self.req_ring = getattr(self.allocator, "swa_req_ring", False) is True + + def _available_and_evictable(self): + if self.all_swa: + return ( + self.allocator.swa_available_size() + + self.tree_cache.swa_evictable_size() + ) + return ( + self.allocator.full_available_size() + self.tree_cache.full_evictable_size() + ) + + @property + def remaining_swa(self): + evictable = 0 if self.req_ring else self.tree_cache.swa_evictable_size() + return self.allocator.swa_available_size() + evictable - self.swa_offset + + def swa_tokens( + self, + extend_input_len, + max_new_tokens, + *, + chunk_limit=None, + swa_host_hit_length=0, + ): + if self.req_ring: + return self.allocator.swa_ring_cost_tokens + return estimate_swa_kv_tokens( + extend_input_len, + max_new_tokens, + sliding_window_size=self.tree_cache.sliding_window_size, + page_size=self.page_size, + allocation_limit=chunk_limit, + host_hit_length=swa_host_hit_length, + ) + + def swa_never_fits(self, extend_input_len, max_new_tokens, **kwargs): + needed = self.swa_tokens(extend_input_len, max_new_tokens, **kwargs) + return ( + needed > self.allocator.size_swa + if self.req_ring + else needed >= self.allocator.size_swa + ) + + def _chunk_cap(self, max_new_tokens, swa_host_hit_length=0): + headroom = self.swa_tokens( + 0, max_new_tokens, swa_host_hit_length=swa_host_hit_length + ) + cap = int(self.remaining_swa) - headroom + return max(0, cap // self.page_size * self.page_size) + + def check_prefill( + self, + *, + extend_input_len: int, + total_tokens: int, + max_new_tokens: int, + input_tokens: int, + swa_host_hit_length: int, + chunk_limit: int | None, + ) -> tuple[bool, int | None]: + if total_tokens >= self.remaining_total: + return False, None + extend_input_len = self.ceil_paged_tokens(extend_input_len) + needed = self.swa_tokens( + extend_input_len, + max_new_tokens, + chunk_limit=chunk_limit, + swa_host_hit_length=swa_host_hit_length, + ) + fits = ( + needed <= self.remaining_swa + if self.req_ring + else needed < self.remaining_swa + ) + if fits: + return True, chunk_limit + # Only permanent shortfalls may shrink a chunk. Transient pressure waits + # so a new prefill does not consume running decodes' window headroom. + cap = 0 + if self.swa_never_fits( + extend_input_len, + max_new_tokens, + chunk_limit=chunk_limit, + swa_host_hit_length=swa_host_hit_length, + ): + cap = self._chunk_cap(max_new_tokens, swa_host_hit_length) + if chunk_limit is None or cap <= 0: + return False, None + return True, min(chunk_limit, cap) + + def has_capacity(self) -> bool: + return super().has_capacity() and self.remaining_swa > 0 + + def available_chunk_tokens(self, chunk_limit: int) -> int | None: + available = min(chunk_limit, int(self.remaining_total)) + if not self.req_ring: + available = min(available, int(self.remaining_swa) - self.page_size) + return available if available > 0 else None + + def can_allocate_prefill( + self, + *, + paged_input: int, + extend_input_len: int, + max_new_tokens: int, + chunk_limit: int | None, + ) -> bool: + return ( + super().can_allocate_prefill( + paged_input=paged_input, + extend_input_len=extend_input_len, + max_new_tokens=max_new_tokens, + chunk_limit=chunk_limit, + ) + and self.swa_tokens( + extend_input_len, max_new_tokens, chunk_limit=chunk_limit + ) + <= self.remaining_swa + ) + + def reserve( + self, + extend_input_len: int, + max_new_tokens: int, + *, + extra_tokens: int = 0, + chunk_limit: int | None = None, + is_chunked_continuation: bool = False, + ) -> None: + super().reserve(extend_input_len, max_new_tokens, extra_tokens=extra_tokens) + # A continuation already owns its ring slot. + if not (self.req_ring and is_chunked_continuation): + self.swa_offset += self.swa_tokens( + self.ceil_paged_tokens(extend_input_len), + max_new_tokens, + chunk_limit=chunk_limit, + ) + + +class SharedSWAPrefillBudget(SWAPrefillBudget): + """FULL and SWA reservations compete for the same physical byte budget.""" + + def __init__(self, *args, num_mixed_decode_tokens=0, **kwargs): + super().__init__( + *args, num_mixed_decode_tokens=num_mixed_decode_tokens, **kwargs + ) + self.swa_offset = num_mixed_decode_tokens + + def _fits(self, full_tokens, swa_tokens, *, empty_pool=False): + return self.allocator.can_reserve( + full_tokens + (0 if empty_pool else self.total_offset), + swa_tokens + (0 if empty_pool else self.swa_offset), + full_evictable_tokens=0 + if empty_pool + else self.tree_cache.full_evictable_size(), + swa_evictable_tokens=0 + if empty_pool + else self.tree_cache.swa_evictable_size(), + empty_pool=empty_pool, + require_token_slack=empty_pool, + ) + + def _joint_chunk_cap(self, *, max_chunk_tokens, chunk_limit, swa_host_hit_length=0): + lo, hi = 0, max(0, max_chunk_tokens) // self.page_size + while lo < hi: + mid = (lo + hi + 1) // 2 + tokens = mid * self.page_size + if self._fits( + tokens + self.page_size, + self.swa_tokens( + tokens, + 0, + chunk_limit=chunk_limit, + swa_host_hit_length=swa_host_hit_length, + ), + ): + lo = mid + else: + hi = mid - 1 + return lo * self.page_size + + def check_prefill( + self, + *, + extend_input_len: int, + total_tokens: int, + max_new_tokens: int, + input_tokens: int, + swa_host_hit_length: int, + chunk_limit: int | None, + ) -> tuple[bool, int | None]: + needed = self.swa_tokens( + extend_input_len, + max_new_tokens, + chunk_limit=chunk_limit, + swa_host_hit_length=swa_host_hit_length, + ) + if self._fits(total_tokens, needed): + return True, chunk_limit + if chunk_limit is None or self._fits( + input_tokens + max_new_tokens + self.page_size, + self.swa_tokens(input_tokens, max_new_tokens, chunk_limit=chunk_limit), + empty_pool=True, + ): + return False, None + cap = self._joint_chunk_cap( + max_chunk_tokens=min(chunk_limit, max(0, extend_input_len - 1)), + chunk_limit=chunk_limit, + swa_host_hit_length=swa_host_hit_length, + ) + return (True, min(chunk_limit, cap)) if cap > 0 else (False, None) + + def has_capacity(self) -> bool: + return self._fits(0, 0) + + def available_chunk_tokens(self, chunk_limit: int) -> int | None: + return chunk_limit + + def fit_chunk( + self, + *, + extend_input_len: int, + max_new_tokens: int, + chunk_limit: int, + ) -> int | None: + candidate = min(extend_input_len, chunk_limit) + finishes = candidate >= extend_input_len + headroom = max_new_tokens if finishes else 0 + if self._fits( + self.ceil_paged_tokens(candidate) + headroom + self.page_size, + self.swa_tokens(candidate, headroom, chunk_limit=chunk_limit), + ): + return chunk_limit + cap = self._joint_chunk_cap( + max_chunk_tokens=max(0, candidate - 1) if finishes else candidate, + chunk_limit=chunk_limit, + ) + return min(chunk_limit, cap) if cap > 0 else None + + def can_allocate_prefill( + self, + *, + paged_input: int, + extend_input_len: int, + max_new_tokens: int, + chunk_limit: int | None, + ) -> bool: + return self._fits( + paged_input + max_new_tokens + self.page_size, + self.swa_tokens(extend_input_len, max_new_tokens, chunk_limit=chunk_limit), + ) diff --git a/test/registered/unit/managers/test_prefill_adder.py b/test/registered/unit/managers/test_prefill_adder.py index 84d9dcd8633e..22aca576ff90 100644 --- a/test/registered/unit/managers/test_prefill_adder.py +++ b/test/registered/unit/managers/test_prefill_adder.py @@ -2,6 +2,8 @@ from types import SimpleNamespace from unittest.mock import MagicMock, patch +import torch + import sglang.srt.managers.schedule_policy as schedule_policy from sglang.srt.managers.schedule_batch import Req from sglang.srt.managers.schedule_policy import ( @@ -13,6 +15,12 @@ DecLockRefResult, IncLockRefResult, ) +from sglang.srt.mem_cache.prefill_budget import ( + PrefillBudget, + SWAPrefillBudget, + estimate_swa_kv_tokens, +) +from sglang.srt.mem_cache.unified_memory_pool import init_unified_swa_pools from sglang.srt.runtime_context import get_context from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler from sglang.srt.utils.common import Range @@ -71,6 +79,10 @@ def create_token_allocator( allocator.swa_available_size.return_value = swa_available_size allocator.available_size.return_value = available_size allocator.size_swa = size_swa + allocator.swa_req_ring = False + allocator.create_prefill_budget.side_effect = lambda tree_cache, **kwargs: ( + PrefillBudget(allocator, tree_cache, **kwargs) + ) return allocator def create_running_batch(self, reqs=None) -> MagicMock: @@ -100,6 +112,7 @@ def create_mock_req(self, rid, priority, max_new_tokens, output_len=0, wait_time req.time_stats = SimpleNamespace(wait_queue_entry_time=wait_time) req.retracted_stain = False req.host_hit_length = 0 + req.swa_host_hit_length = 0 req.storage_hit_length = 0 req.storage_hit_start = None req.host_hit_is_storage = False @@ -123,8 +136,96 @@ def create_adder(self, running_batch, **kwargs): priority_scheduling_preemption_threshold=0, ) defaults.update(kwargs) + defaults["token_to_kv_pool_allocator"].page_size = defaults["page_size"] return PrefillAdder(**defaults) + def create_shared_adder(self): + self.mock_tree_cache.supports_mamba.return_value = False + self.mock_tree_cache.sliding_window_size = 8 + self.mock_tree_cache.is_tree_cache.return_value = False + allocator = init_unified_swa_pools( + device="cpu", + kv_cache_dtype=torch.float16, + head_num=1, + head_dim=4, + v_head_dim=4, + swa_head_num=1, + swa_head_dim=4, + swa_v_head_dim=4, + page_size=4, + start_layer=0, + end_layer=2, + swa_attention_layer_ids=[1], + full_attention_layer_ids=[0], + total_bytes=1024, + enable_memory_saver=False, + need_sort=False, + lazy_compaction=True, + ).token_to_kv_pool_allocator + return self.create_adder( + self.create_running_batch(), + page_size=4, + rem_chunk_tokens=16, + token_to_kv_pool_allocator=allocator, + ) + + def create_shared_req(self, rid, max_new_tokens=4): + req = self.create_mock_req(rid, priority=0, max_new_tokens=max_new_tokens) + req.sampling_params.ignore_eos = False + req.swa_host_hit_length = 0 + req.last_node = MagicMock() + req.full_untruncated_fill_ids = list(range(12)) + req.set_extend_range = MagicMock( + side_effect=lambda start, end: setattr( + req, "extend_range", Range(start, end) + ) + ) + return req + + def test_shared_admission_reserves_all_pending_requests(self): + adder = self.create_shared_adder() + first, second = ( + self.create_shared_req("first"), + self.create_shared_req("second"), + ) + adder.add_one_req(first, has_chunked_req=False, truncation_align_size=None) + self.assertEqual(adder.can_run_list, [first]) + self.assertEqual( + adder.add_one_req( + second, has_chunked_req=False, truncation_align_size=None + ), + AddReqResult.NO_TOKEN, + ) + self.assertEqual(adder.can_run_list, [first]) + + def test_shared_admission_rechecks_after_prefix_lock(self): + adder = self.create_shared_adder() + self.assertIsNotNone(adder.token_to_kv_pool_allocator.alloc(24)) + self.mock_tree_cache.full_evictable_size.return_value = 24 + self.mock_tree_cache.swa_evictable_size.return_value = 24 + + def lock_prefix(_): + self.mock_tree_cache.full_evictable_size.return_value = 0 + self.mock_tree_cache.swa_evictable_size.return_value = 0 + return IncLockRefResult() + + self.mock_tree_cache.inc_lock_ref.side_effect = lock_prefix + req = self.create_shared_req("locked-prefix") + self.assertEqual( + adder.add_one_req(req, has_chunked_req=False, truncation_align_size=None), + AddReqResult.NO_TOKEN, + ) + self.mock_tree_cache.inc_lock_ref.assert_called_once() + self.assertEqual(adder.can_run_list, []) + + def test_shared_continuation_uses_memory_chunk_limit(self): + adder = self.create_shared_adder() + req = self.create_shared_req("continuation", max_new_tokens=80) + self.assertIs(adder.add_chunked_req(req), req) + self.assertEqual(req.extend_range.length, 8) + self.assertEqual(adder.memory_budget.total_offset, 12) + self.assertEqual(adder.memory_budget.swa_offset, 12) + def test_storage_prefetch_fulfillment_resolves_at_admission(self): adder = self.create_adder(self.create_running_batch()) req = self.create_mock_req("storage-hit", priority=0, max_new_tokens=1) @@ -184,7 +285,7 @@ def test_preempt_success_high_priority_values_first(self): running_batch = self.create_running_batch(running_reqs) adder = self.create_adder(running_batch) - self.assertEqual(adder.rem_total_token_offset, 225) + self.assertEqual(adder.memory_budget.total_offset, 225) self.mock_token_allocator.full_available_size.return_value = ( 225 # full occupation of GRam @@ -197,7 +298,9 @@ def test_preempt_success_high_priority_values_first(self): self.assertTrue(success) self.assertIn(running_reqs[0], adder.preempt_list) - self.assertEqual(adder.rem_total_token_offset, 175) # 50 + 75 + 100 - 50 = 175 + self.assertEqual( + adder.memory_budget.total_offset, 175 + ) # 50 + 75 + 100 - 50 = 175 running_batch.release_req.assert_called_once() def test_preempt_success_low_priority_values_first(self): @@ -214,7 +317,7 @@ def test_preempt_success_low_priority_values_first(self): running_batch = self.create_running_batch(running_reqs) adder = self.create_adder(running_batch) - self.assertEqual(adder.rem_total_token_offset, 225) + self.assertEqual(adder.memory_budget.total_offset, 225) self.mock_token_allocator.full_available_size.return_value = ( 225 # full occupation of GRam @@ -227,7 +330,9 @@ def test_preempt_success_low_priority_values_first(self): self.assertTrue(success) self.assertIn(running_reqs[2], adder.preempt_list) - self.assertEqual(adder.rem_total_token_offset, 125) # 50 + 75 + 100 - 100 = 125 + self.assertEqual( + adder.memory_budget.total_offset, 125 + ) # 50 + 75 + 100 - 100 = 125 running_batch.release_req.assert_called_once() def test_preempt_fail_low_priority_values_first(self): @@ -244,7 +349,7 @@ def test_preempt_fail_low_priority_values_first(self): running_batch = self.create_running_batch(running_reqs) adder = self.create_adder(running_batch) - self.assertEqual(adder.rem_total_token_offset, 225) + self.assertEqual(adder.memory_budget.total_offset, 225) self.mock_token_allocator.full_available_size.return_value = ( 225 # full occupation of GRam @@ -282,7 +387,7 @@ def test_preempt_fail_high_priority_values_first(self): running_batch = self.create_running_batch(running_reqs) adder = self.create_adder(running_batch) - self.assertEqual(adder.rem_total_token_offset, 225) + self.assertEqual(adder.memory_budget.total_offset, 225) self.mock_token_allocator.full_available_size.return_value = ( 225 # full occupation of GRam @@ -320,7 +425,7 @@ def test_preempt_skip_already_preempted_request(self): running_batch = self.create_running_batch(running_reqs) adder = self.create_adder(running_batch) - self.assertEqual(adder.rem_total_token_offset, 225) + self.assertEqual(adder.memory_budget.total_offset, 225) self.mock_token_allocator.full_available_size.return_value = 225 self.mock_token_allocator.available_size.return_value = 225 @@ -332,7 +437,7 @@ def test_preempt_skip_already_preempted_request(self): first_success = adder.preempt_to_schedule(first_req) self.assertTrue(first_success) self.assertIn(running_reqs[0], adder.preempt_list) - self.assertEqual(adder.rem_total_token_offset, 175) + self.assertEqual(adder.memory_budget.total_offset, 175) running_batch.release_req.assert_called_once() # Second call needs more tokens than currently free, so it would need to @@ -343,7 +448,7 @@ def test_preempt_skip_already_preempted_request(self): second_success = adder.preempt_to_schedule(second_req) self.assertFalse(second_success) - self.assertEqual(adder.rem_total_token_offset, 175) + self.assertEqual(adder.memory_budget.total_offset, 175) self.assertEqual(adder.preempt_list.count(running_reqs[0]), 1) running_batch.release_req.assert_called_once() @@ -363,7 +468,7 @@ def test_preempt_success_low_priority_values_first_exact_once(self): running_batch = self.create_running_batch(running_reqs) adder = self.create_adder(running_batch) - self.assertEqual(adder.rem_total_token_offset, 475) + self.assertEqual(adder.memory_budget.total_offset, 475) self.mock_token_allocator.full_available_size.return_value = ( 475 # full occupation of GRam @@ -376,7 +481,7 @@ def test_preempt_success_low_priority_values_first_exact_once(self): self.assertTrue(success) self.assertIn(running_reqs[2], adder.preempt_list) self.assertEqual( - adder.rem_total_token_offset, 375 + adder.memory_budget.total_offset, 375 ) # 50 + 75 + 100 + 125 + 125 - 100 = 375 running_batch.release_req.assert_called_once() @@ -396,7 +501,7 @@ def test_preempt_success_low_priority_values_first_exact_twice(self): running_batch = self.create_running_batch(running_reqs) adder = self.create_adder(running_batch) - self.assertEqual(adder.rem_total_token_offset, 475) + self.assertEqual(adder.memory_budget.total_offset, 475) self.mock_token_allocator.full_available_size.return_value = ( 475 # full occupation of GRam @@ -410,7 +515,7 @@ def test_preempt_success_low_priority_values_first_exact_twice(self): self.assertIn(running_reqs[2], adder.preempt_list) self.assertIn(running_reqs[3], adder.preempt_list) self.assertEqual( - adder.rem_total_token_offset, 250 + adder.memory_budget.total_offset, 250 ) # 50 + 75 + 100 + 125 + 125 - 100 - 125 = 250 self.assertEqual(running_batch.release_req.call_count, 2) @@ -432,8 +537,8 @@ def test_mixed_chunk_prefill_budgets(self): self.assertEqual(adder.rem_input_tokens, 192) # 200 - 8 self.assertEqual(adder.rem_chunk_tokens, 56) # 64 - 8 - self.assertEqual(adder.rem_total_token_offset, 408) # 8 + 8 * 50 - self.assertEqual(adder.cur_rem_token_offset, 8) + self.assertEqual(adder.memory_budget.total_offset, 408) # 8 + 8 * 50 + self.assertEqual(adder.memory_budget.current_offset, 8) self.assertEqual(adder.budget_state(), AddReqResult.CONTINUE) # Add a prefill that exactly consumes the chunk budget @@ -473,7 +578,7 @@ def test_mixed_chunk_prefill_budgets(self): self.assertEqual(adder2.rem_input_tokens, 195) # 200 - 5 self.assertEqual(adder2.rem_chunk_tokens, 59) # 64 - 5 - self.assertEqual(adder2.rem_total_token_offset, 255) # 5 + 5 * 50 + self.assertEqual(adder2.memory_budget.total_offset, 255) # 5 + 5 * 50 self.assertEqual(adder2.budget_state(), AddReqResult.CONTINUE) # Same prefill no longer exhausts the chunk budget @@ -538,6 +643,10 @@ def _build_hybrid_swa_chunked_req( rem_chunk_tokens=rem_chunk, ) adder.is_hybrid_swa = is_hybrid_swa + if is_hybrid_swa: + adder.memory_budget = SWAPrefillBudget( + self.mock_token_allocator, self.mock_tree_cache + ) req = self.create_mock_req("chunked", priority=0, max_new_tokens=128) req.prefix_indices = [] @@ -616,7 +725,16 @@ def test_swa_budget_for_req(self): page_size=page, rem_chunk_tokens=rem_chunk, ) - self.assertEqual(adder._swa_budget_for_req(extend, max_new), expected) + self.assertEqual( + estimate_swa_kv_tokens( + extend, + max_new, + sliding_window_size=window, + page_size=page, + allocation_limit=rem_chunk, + ), + expected, + ) def test_swa_admission_admits_short_cached_resume_at_two_window_pool(self): # Livelock regression (real incident). At an SWA pool ~= 2 sliding @@ -636,6 +754,9 @@ def test_swa_admission_admits_short_cached_resume_at_two_window_pool(self): self.mock_tree_cache.is_tree_cache.return_value = False adder = self.create_adder(self.create_running_batch(), page_size=PAGE) adder.is_hybrid_swa = True + adder.memory_budget = SWAPrefillBudget( + self.mock_token_allocator, self.mock_tree_cache + ) req = self.create_mock_req( "resume", priority=0, max_new_tokens=40, output_len=10 @@ -653,7 +774,9 @@ def test_swa_admission_admits_short_cached_resume_at_two_window_pool(self): req.sampling_params = SimpleNamespace(max_new_tokens=40, ignore_eos=False) # Pre-fix: a constant sliding-window reservation rejects the resume. - with patch.object(adder, "_swa_budget_for_req", return_value=WINDOW + PAGE): + with patch.object( + adder.memory_budget, "swa_tokens", return_value=WINDOW + PAGE + ): self.assertIs( adder.add_one_req( req, has_chunked_req=False, truncation_align_size=None @@ -905,6 +1028,11 @@ def create_swa_adder( ) -> PrefillAdder: self.mock_tree_cache.sliding_window_size = sliding_window self.mock_token_allocator = self.create_token_allocator(size_swa=size_swa) + self.mock_token_allocator.create_prefill_budget.side_effect = ( + lambda tree_cache, **kwargs: SWAPrefillBudget( + self.mock_token_allocator, tree_cache, **kwargs + ) + ) return self.create_adder( self.create_running_batch(), page_size=page_size, @@ -916,7 +1044,7 @@ def test_swa_never_fits_false_under_transient_pressure(self): # once running decodes drain -> must wait, not take the hatch. adder = self.create_swa_adder(size_swa=1024, sliding_window=128) self.assertFalse( - adder._swa_req_never_fits(extend_input_len=256, max_new_tokens=64) + adder.memory_budget.swa_never_fits(extend_input_len=256, max_new_tokens=64) ) def test_swa_never_fits_true_when_budget_exceeds_whole_pool(self): @@ -924,7 +1052,7 @@ def test_swa_never_fits_true_when_budget_exceeds_whole_pool(self): # pool: it can never fit however far the pool drains -> hatch. adder = self.create_swa_adder(size_swa=1024, sliding_window=128) self.assertTrue( - adder._swa_req_never_fits( + adder.memory_budget.swa_never_fits( extend_input_len=256, max_new_tokens=64, swa_host_hit_length=4096 ) ) @@ -934,14 +1062,14 @@ def test_swa_never_fits_is_gated_by_pool_capacity(self): # the budget against size_swa (guards against a wrong-accessor bug). req = dict(extend_input_len=256, max_new_tokens=64, swa_host_hit_length=600) self.assertTrue( - self.create_swa_adder(size_swa=512, sliding_window=128)._swa_req_never_fits( - **req - ) + self.create_swa_adder( + size_swa=512, sliding_window=128 + ).memory_budget.swa_never_fits(**req) ) self.assertFalse( self.create_swa_adder( size_swa=4096, sliding_window=128 - )._swa_req_never_fits(**req) + ).memory_budget.swa_never_fits(**req) ) diff --git a/test/registered/unit/managers/test_scheduler_init_req_max_new_tokens.py b/test/registered/unit/managers/test_scheduler_init_req_max_new_tokens.py index 43f182778e26..d651071e5d84 100644 --- a/test/registered/unit/managers/test_scheduler_init_req_max_new_tokens.py +++ b/test/registered/unit/managers/test_scheduler_init_req_max_new_tokens.py @@ -6,6 +6,7 @@ from sglang.srt.environ import envs from sglang.srt.managers.scheduler import Scheduler +from sglang.srt.mem_cache.allocator.token import TokenToKVPoolAllocator from sglang.srt.mem_cache.unified_memory_pool import init_unified_swa_pools from sglang.srt.runtime_context import get_parallel from sglang.test.ci.ci_register import register_cpu_ci @@ -59,7 +60,16 @@ def _new_scheduler( scheduler.max_total_num_tokens = max_total_num_tokens scheduler.page_size = page_size scheduler.max_new_tokens_limit = envs.SGLANG_MAX_NEW_TOKENS_LIMIT.get() - scheduler.token_to_kv_pool_allocator = None + scheduler.sliding_window_size = None + scheduler.chunked_prefill_size = None + scheduler.token_to_kv_pool_allocator = TokenToKVPoolAllocator( + size=max_total_num_tokens, + dtype=torch.int64, + device="cpu", + kvcache=None, + need_sort=False, + ) + scheduler.token_to_kv_pool_allocator.page_size = page_size return scheduler def _new_req(self, max_new_tokens, input_len: int = 8, min_new_tokens: int = 0): diff --git a/test/registered/unit/mem_cache/test_multi_ended_allocator.py b/test/registered/unit/mem_cache/test_multi_ended_allocator.py index aaed067bf844..615dfec53b75 100644 --- a/test/registered/unit/mem_cache/test_multi_ended_allocator.py +++ b/test/registered/unit/mem_cache/test_multi_ended_allocator.py @@ -29,7 +29,6 @@ import torch -from sglang.srt.managers.schedule_policy import estimate_swa_kv_tokens from sglang.srt.mem_cache.allocator.unified_hybrid_swa import ( UnifiedSWATokenToKVPoolAllocator, ) @@ -41,6 +40,7 @@ MultiEndedAllocator, ) from sglang.srt.mem_cache.base_prefix_cache import EvictParams +from sglang.srt.mem_cache.prefill_budget import estimate_swa_kv_tokens from sglang.srt.mem_cache.unified_cache.components import ComponentType from sglang.srt.mem_cache.unified_memory_pool import ( MambaSubPoolSpec, diff --git a/test/registered/unit/mem_cache/test_prefill_memory_budget.py b/test/registered/unit/mem_cache/test_prefill_memory_budget.py new file mode 100644 index 000000000000..3f08e40fdb5f --- /dev/null +++ b/test/registered/unit/mem_cache/test_prefill_memory_budget.py @@ -0,0 +1,224 @@ +"""CPU regressions for allocator-owned prefill admission and pending demand.""" + +import unittest +from types import SimpleNamespace +from unittest.mock import MagicMock + +import torch + +from sglang.srt.mem_cache.allocator.hisparse import ( + DeepSeekV4HiSparseTokenToKVPoolAllocator, +) +from sglang.srt.mem_cache.allocator.swa import ( + PureSWATokenToKVPoolAllocator, + SWATokenToKVPoolAllocator, +) +from sglang.srt.mem_cache.allocator.unified_hybrid_swa import ( + UnifiedMambaSWATokenToKVPoolAllocator, +) +from sglang.srt.mem_cache.common import evict_from_tree_cache +from sglang.srt.mem_cache.prefill_budget import SWAPrefillBudget +from sglang.srt.mem_cache.unified_memory_pool import init_unified_swa_pools +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + + +def _shared_allocator(): + return init_unified_swa_pools( + device="cpu", + kv_cache_dtype=torch.float16, + head_num=1, + head_dim=4, + v_head_dim=4, + swa_head_num=1, + swa_head_dim=4, + swa_v_head_dim=4, + page_size=4, + start_layer=0, + end_layer=2, + swa_attention_layer_ids=[1], + full_attention_layer_ids=[0], + total_bytes=1024, + enable_memory_saver=False, + need_sort=False, + lazy_compaction=True, + ).token_to_kv_pool_allocator + + +def _cache(): + return SimpleNamespace( + sliding_window_size=8, + full_evictable_size=lambda: 0, + swa_evictable_size=lambda: 0, + is_chunk_cache=lambda: False, + ) + + +class TestSharedPrefillMemoryBudget(unittest.TestCase): + def setUp(self): + self.allocator = _shared_allocator() + self.cache = _cache() + self.budget = self.allocator.create_prefill_budget(self.cache) + self.request = dict( + extend_input_len=12, + total_tokens=20, + max_new_tokens=4, + input_tokens=12, + swa_host_hit_length=0, + chunk_limit=16, + ) + + def test_pending_batch_cannot_spend_shared_bytes_twice(self): + self.assertEqual(self.budget.check_prefill(**self.request), (True, 16)) + self.budget.reserve(12, 4, chunk_limit=16) + # Each side separately has room, but their combined reservation does not. + self.assertGreater(self.budget.remaining_total, 20) + self.assertGreater(self.budget.remaining_swa, 16) + self.assertEqual(self.budget.check_prefill(**self.request), (False, None)) + self.assertEqual(self.allocator.full_attn_allocator.allocated_count(), 0) + self.assertEqual(self.allocator.swa_attn_allocator.allocated_count(), 0) + + def test_mixed_decode_reserves_both_sides(self): + budget = self.allocator.create_prefill_budget( + self.cache, num_mixed_decode_tokens=4 + ) + budget.reserve(12, 4, chunk_limit=16) + self.assertEqual( + (budget.total_offset, budget.current_offset, budget.swa_offset), + (24, 20, 20), + ) + self.assertEqual(budget.check_prefill(**self.request), (False, None)) + + def test_prefix_lock_changes_admission_without_rebuilding_budget(self): + self.assertIsNotNone(self.allocator.alloc(24)) + self.cache.full_evictable_size = lambda: 24 + self.cache.swa_evictable_size = lambda: 24 + self.assertEqual(self.budget.check_prefill(**self.request), (True, 16)) + # Locking the cached prefix removes its eviction credit. + self.cache.full_evictable_size = lambda: 0 + self.cache.swa_evictable_size = lambda: 0 + self.assertEqual(self.budget.check_prefill(**self.request), (False, None)) + + def test_final_chunk_reserves_decode_headroom(self): + limit = self.budget.fit_chunk( + extend_input_len=12, max_new_tokens=80, chunk_limit=16 + ) + self.assertEqual(limit, 8) + self.budget.reserve(limit, 0, chunk_limit=16, is_chunked_continuation=True) + self.assertIsNotNone(self.allocator.alloc(limit)) + + def test_host_swa_load_is_part_of_joint_demand(self): + self.assertEqual(self.budget.check_prefill(**self.request), (True, 16)) + request = {**self.request, "swa_host_hit_length": 32} + self.assertEqual(self.budget.check_prefill(**request), (False, None)) + + def test_prompt_clipping_uses_the_empty_pool(self): + kwargs = dict(token_capacity=1, sliding_window_size=8, chunk_size=16) + limit = self.allocator.max_new_tokens_for_memory(12, 80, **kwargs) + self.assertIsNotNone(limit) + self.assertGreater(limit, 0) + self.assertIsNotNone(self.allocator.alloc(24)) + self.assertEqual( + self.allocator.max_new_tokens_for_memory(12, 80, **kwargs), limit + ) + self.assertIsNone(self.allocator.max_new_tokens_for_memory(100, 0, **kwargs)) + + def test_shared_stats_pair_available_tokens_with_current_capacity(self): + self.assertIsNotNone(self.allocator.alloc(12)) + (full_capacity, full_free), (swa_capacity, swa_free) = ( + self.allocator.swa_capacity_and_available(full_capacity=1, swa_capacity=1) + ) + self.assertEqual(full_capacity - full_free, 12) + self.assertEqual(swa_capacity - swa_free, 12) + + def test_common_eviction_dispatches_joint_reclaim(self): + self.cache.token_to_kv_pool_allocator = self.allocator + self.allocator.evict_to_free_tokens = MagicMock() + evict_from_tree_cache(self.cache, 8) + self.allocator.evict_to_free_tokens.assert_called_once_with(self.cache, 8) + + +class TestFixedPrefillMemoryBudget(unittest.TestCase): + def _allocator(self, cls=SWATokenToKVPoolAllocator): + allocator = object.__new__(cls) + allocator.page_size = 4 + allocator._size_full = 128 + allocator._size_swa = 64 + allocator.full_available_size = lambda: 128 + allocator.swa_available_size = lambda: 64 + return allocator + + def test_ring_slot_reserved_once_and_evictable_tokens_give_no_credit(self): + allocator = self._allocator() + allocator._swa_req_ring = True + allocator._swa_ring_cost = 32 + cache = _cache() + cache.swa_evictable_size = lambda: 1000 + budget = allocator.create_prefill_budget(cache) + budget.reserve(12, 4, chunk_limit=16) + self.assertEqual(budget.remaining_swa, 32) + budget.reserve(12, 4, chunk_limit=16, is_chunked_continuation=True) + self.assertEqual(budget.remaining_swa, 32) + # The last exact ring slot remains admissible. + self.assertEqual( + budget.check_prefill( + extend_input_len=4, + total_tokens=12, + max_new_tokens=4, + input_tokens=4, + swa_host_hit_length=0, + chunk_limit=16, + ), + (True, 16), + ) + + def test_pure_swa_budget_reads_swa_capacity(self): + allocator = self._allocator(PureSWATokenToKVPoolAllocator) + allocator.full_available_size = lambda: 0 + budget = allocator.create_prefill_budget(_cache()) + self.assertEqual(budget.remaining_total, 64) + self.assertTrue(budget.has_capacity()) + + def test_hisparse_budget_reads_wrapper_capacity(self): + allocator = self._allocator(DeepSeekV4HiSparseTokenToKVPoolAllocator) + allocator.full_available_size = lambda: 12 + budget = allocator.create_prefill_budget(_cache()) + self.assertEqual(budget.remaining_total, 12) + self.assertEqual(budget.remaining_swa, 64) + + def test_tri_pool_keeps_fixed_admission_and_clipping(self): + allocator = self._allocator(UnifiedMambaSWATokenToKVPoolAllocator) + allocator.can_reserve = MagicMock( + side_effect=AssertionError("two-pool reservation") + ) + budget = allocator.create_prefill_budget(_cache()) + self.assertIs(type(budget), SWAPrefillBudget) + self.assertTrue(budget.has_capacity()) + self.assertEqual( + allocator.max_new_tokens_for_memory( + 5, + 100, + token_capacity=32, + sliding_window_size=8, + chunk_size=16, + ), + 19, + ) + + def test_tri_pool_eviction_does_not_reenter_common(self): + allocator = self._allocator(UnifiedMambaSWATokenToKVPoolAllocator) + allocator.available_size = lambda: 0 + allocator.full_available_size = lambda: 0 + allocator.swa_available_size = lambda: 0 + cache = _cache() + cache.token_to_kv_pool_allocator = allocator + cache.evict_for_alloc = MagicMock() + evict_from_tree_cache(cache, 8) + cache.evict_for_alloc.assert_called_once() + params = cache.evict_for_alloc.call_args.args[0] + self.assertEqual((params.num_tokens, params.swa_num_tokens), (8, 8)) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/mem_cache/test_unified_radix_allocation_eviction.py b/test/registered/unit/mem_cache/test_unified_radix_allocation_eviction.py index 68c8cf1262e5..5b99273aa8c7 100644 --- a/test/registered/unit/mem_cache/test_unified_radix_allocation_eviction.py +++ b/test/registered/unit/mem_cache/test_unified_radix_allocation_eviction.py @@ -554,7 +554,11 @@ def evict_leaf(_node_id, tracker): def test_common_helper_uses_allocation_aware_entry_point(self): tree_cache = MagicMock() tree_cache.is_chunk_cache.return_value = False - tree_cache.token_to_kv_pool_allocator.available_size.return_value = 30 + from sglang.srt.mem_cache.allocator.token import TokenToKVPoolAllocator + + allocator = object.__new__(TokenToKVPoolAllocator) + allocator.available_size = lambda: 30 + tree_cache.token_to_kv_pool_allocator = allocator evict_from_tree_cache(tree_cache, num_tokens=100) From 1f3b496c2395323c954d0f0db3e81d04e4d15697 Mon Sep 17 00:00:00 2001 From: yhzhuang Date: Sun, 13 Sep 2026 18:25:59 -0700 Subject: [PATCH 49/55] Fix unaligned unified SWA prefill admission --- python/sglang/srt/mem_cache/prefill_budget.py | 4 +- .../mem_cache/test_prefill_memory_budget.py | 82 ++++++++++++++++++- 2 files changed, 81 insertions(+), 5 deletions(-) diff --git a/python/sglang/srt/mem_cache/prefill_budget.py b/python/sglang/srt/mem_cache/prefill_budget.py index 0811eb763354..965daaa55f9d 100644 --- a/python/sglang/srt/mem_cache/prefill_budget.py +++ b/python/sglang/srt/mem_cache/prefill_budget.py @@ -381,7 +381,7 @@ def fit_chunk( finishes = candidate >= extend_input_len headroom = max_new_tokens if finishes else 0 if self._fits( - self.ceil_paged_tokens(candidate) + headroom + self.page_size, + candidate + headroom + self.page_size, self.swa_tokens(candidate, headroom, chunk_limit=chunk_limit), ): return chunk_limit @@ -400,6 +400,6 @@ def can_allocate_prefill( chunk_limit: int | None, ) -> bool: return self._fits( - paged_input + max_new_tokens + self.page_size, + extend_input_len + max_new_tokens + self.page_size, self.swa_tokens(extend_input_len, max_new_tokens, chunk_limit=chunk_limit), ) diff --git a/test/registered/unit/mem_cache/test_prefill_memory_budget.py b/test/registered/unit/mem_cache/test_prefill_memory_budget.py index 3f08e40fdb5f..2a73dfffc9e3 100644 --- a/test/registered/unit/mem_cache/test_prefill_memory_budget.py +++ b/test/registered/unit/mem_cache/test_prefill_memory_budget.py @@ -1,11 +1,15 @@ """CPU regressions for allocator-owned prefill admission and pending demand.""" import unittest +from array import array from types import SimpleNamespace from unittest.mock import MagicMock import torch +from sglang.srt.managers.schedule_batch import Req +from sglang.srt.managers.schedule_policy import PrefillAdder +from sglang.srt.managers.scheduler import Scheduler from sglang.srt.mem_cache.allocator.hisparse import ( DeepSeekV4HiSparseTokenToKVPoolAllocator, ) @@ -19,12 +23,14 @@ from sglang.srt.mem_cache.common import evict_from_tree_cache from sglang.srt.mem_cache.prefill_budget import SWAPrefillBudget from sglang.srt.mem_cache.unified_memory_pool import init_unified_swa_pools +from sglang.srt.runtime_context import get_parallel +from sglang.srt.sampling.sampling_params import SamplingParams from sglang.test.ci.ci_register import register_cpu_ci register_cpu_ci(est_time=5, suite="base-a-test-cpu") -def _shared_allocator(): +def _shared_allocator(*, page_size=4, total_bytes=1024): return init_unified_swa_pools( device="cpu", kv_cache_dtype=torch.float16, @@ -34,12 +40,12 @@ def _shared_allocator(): swa_head_num=1, swa_head_dim=4, swa_v_head_dim=4, - page_size=4, + page_size=page_size, start_layer=0, end_layer=2, swa_attention_layer_ids=[1], full_attention_layer_ids=[0], - total_bytes=1024, + total_bytes=total_bytes, enable_memory_saver=False, need_sort=False, lazy_compaction=True, @@ -139,6 +145,76 @@ def test_common_eviction_dispatches_joint_reclaim(self): self.allocator.evict_to_free_tokens.assert_called_once_with(self.cache, 8) +class TestSharedPrefillAdmission(unittest.TestCase): + def _new_admission(self, page_size, pool_pages, *, ignore_eos=False): + allocator = _shared_allocator( + page_size=page_size, total_bytes=pool_pages * page_size * 16 + ) + req = Req( + rid="unaligned-prompt", + origin_input_text=None, + origin_input_ids=array("q", [1] * (page_size + 1)), + sampling_params=SamplingParams(max_new_tokens=1, ignore_eos=ignore_eos), + ) + scheduler = Scheduler.__new__(Scheduler) + scheduler.max_req_len = 16 * page_size + scheduler.max_total_num_tokens = allocator.size_full + scheduler.page_size = page_size + scheduler.max_new_tokens_limit = None + scheduler.sliding_window_size = page_size + scheduler.chunked_prefill_size = page_size + scheduler.token_to_kv_pool_allocator = allocator + with get_parallel().override(attn_dcp_size=1): + scheduler.init_req_max_new_tokens(req) + self.assertEqual(req.sampling_params.max_new_tokens, 1) + req._refresh_fill_ids() + + cache = SimpleNamespace( + sliding_window_size=page_size, + disable=True, + full_evictable_size=lambda: 0, + swa_evictable_size=lambda: 0, + is_chunk_cache=lambda: True, + supports_mamba=lambda: False, + ) + adder = PrefillAdder( + page_size=page_size, + tree_cache=cache, + token_to_kv_pool_allocator=allocator, + running_batch=None, + new_token_ratio=1.0, + rem_input_tokens=16 * page_size, + rem_chunk_tokens=page_size, + ) + return allocator, req, adder + + def test_unaligned_final_chunk_makes_progress(self): + for page_size in (4, 64): + with self.subTest(page_size=page_size): + allocator, req, adder = self._new_admission(page_size, pool_pages=7) + req.prefix_indices = allocator.alloc(page_size) + self.assertIsNotNone(req.prefix_indices) + self.assertTrue(allocator.can_reserve(page_size + 2, page_size + 2)) + + self.assertIsNone(adder.add_chunked_req(req)) + self.assertEqual(adder.can_run_list, [req]) + self.assertEqual(req.extend_range.length, 1) + + def test_unaligned_ignore_eos_enters_empty_pool(self): + for page_size in (4, 64): + with self.subTest(page_size=page_size): + allocator, req, adder = self._new_admission( + page_size, pool_pages=6, ignore_eos=True + ) + self.assertEqual(len(req.prefix_indices), 0) + self.assertTrue(allocator.can_reserve(2 * page_size + 2, 2 * page_size)) + + adder.add_one_req( + req, has_chunked_req=False, truncation_align_size=None + ) + self.assertEqual(adder.can_run_list, [req]) + + class TestFixedPrefillMemoryBudget(unittest.TestCase): def _allocator(self, cls=SWATokenToKVPoolAllocator): allocator = object.__new__(cls) From f20860697232a2f565153b15fd242693a0a4257b Mon Sep 17 00:00:00 2001 From: yhzhuang Date: Sun, 13 Sep 2026 18:57:43 -0700 Subject: [PATCH 50/55] Refactor unified SWA allocators into sibling capacity policies Retain shared virtual-ID lifecycle in a common base, gate eager-pool post-capture planning, and preserve capacity realization results. Keep allocator-owned admission policy and existing tri-pool behavior. --- python/sglang/srt/arg_groups/overrides.py | 3 + python/sglang/srt/mem_cache/allocator/base.py | 7 +- .../mem_cache/allocator/unified_hybrid_swa.py | 904 +++++++++--------- python/sglang/srt/mem_cache/common.py | 6 +- .../srt/mem_cache/kv_cache_configurator.py | 8 +- .../srt/mem_cache/kv_index_translator.py | 6 +- .../mem_cache/unified_cache/components/swa.py | 4 +- .../unit/mem_cache/test_unified_tri_pool.py | 8 + 8 files changed, 492 insertions(+), 454 deletions(-) diff --git a/python/sglang/srt/arg_groups/overrides.py b/python/sglang/srt/arg_groups/overrides.py index 6d26836debea..1531ee1f3cd6 100644 --- a/python/sglang/srt/arg_groups/overrides.py +++ b/python/sglang/srt/arg_groups/overrides.py @@ -1808,6 +1808,9 @@ def post_capture_kv_sizing_planned(server_args: Any) -> bool: mla_enabled = use_mla_backend(server_args) if not envs.SGLANG_ENABLE_POST_CAPTURE_KV_SIZING.get(): return False + # Unified arenas are fully backed before capture and cannot resize afterward. + if cfg.enable_unified_memory: + return False if cfg.device != "cuda": return False if cfg.dcp_size != 1: diff --git a/python/sglang/srt/mem_cache/allocator/base.py b/python/sglang/srt/mem_cache/allocator/base.py index 639b1cbeafcf..0c4ec17ce821 100644 --- a/python/sglang/srt/mem_cache/allocator/base.py +++ b/python/sglang/srt/mem_cache/allocator/base.py @@ -95,9 +95,12 @@ def max_new_tokens_for_memory( 0, min(max_new_tokens, token_capacity - paged_input - self.page_size - 1) ) - def evict_to_free_tokens(self, tree_cache, num_tokens: int) -> None: + def evict_to_free_tokens(self, tree_cache, num_tokens: int) -> bool | None: """Evict unlocked prefix-cache entries until this allocator can serve - ``num_tokens`` or nothing evictable remains.""" + ``num_tokens`` or nothing evictable remains. + + Return whether capacity was realized, or None if it still needs checking. + """ from sglang.srt.mem_cache.base_prefix_cache import EvictParams from sglang.srt.mem_cache.common import _evict_until_allocatable diff --git a/python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py b/python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py index c4d1c579ef39..64d03f9aadfc 100644 --- a/python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py +++ b/python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py @@ -18,6 +18,7 @@ import logging import math +from abc import abstractmethod from typing import Callable, List, Optional, Sequence, Tuple import torch @@ -43,11 +44,10 @@ logger = logging.getLogger(__name__) -class UnifiedSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator): - """Composite allocator for the hybrid SWA pair (full + swa MHA sub-pools). +class UnifiedSWAAllocatorBase(SWATokenToKVPoolAllocator): + """Shared FULL/SWA virtual IDs, allocation lifecycle, and index translation. - One alloc(N) binds N pages on BOTH sides under the same virtual id, so - `available_size()` (joint bytes, in TOKENS) is the only safe alloc pre-check. + Concrete allocators define the two-ended or Mamba/SWA/FULL capacity policy. """ # Parent's `size` property has no setter but base init does `self.size = size`; @@ -178,18 +178,6 @@ def __init__( # -- construction hooks (the tri-pool subclass overrides both) -- - def _build_swa_attn_allocator(self, **kwargs) -> MultiEndedAllocator: - """The swa sub-allocator: an END pool in the 2-pool pair.""" - return MultiEndedAllocator( - sub_pool_name="swa", - is_id_owner=False, # non-owner; consumes virtuals minted by full - **kwargs, - ) - - def _wire_peers(self) -> None: - self.full_attn_allocator.bind_peer(self.swa_attn_allocator) - self.swa_attn_allocator.bind_peer(self.full_attn_allocator) - # -- capacity reporting (three-way split) -- def available_size(self) -> int: @@ -200,46 +188,6 @@ def available_size(self) -> int: self._joint_avail_memo_epoch = epoch return self._joint_avail_memo_tokens - def _compute_available_size(self) -> int: - """Joint byte budget in TOKENS: each composite alloc(1) consumes one - full-side AND one swa-side page under the same virtual id.""" - fa, sa = self.full_attn_allocator, self.swa_attn_allocator - e_f = fa.entry_bytes_per_page - e_s = sa.entry_bytes_per_page - # Direction-agnostic shared gap: the free byte band between the two pools. - if fa.grow_direction == "up": - gap_bytes = max(0, sa._byte_low_frontier() - fa._byte_high_frontier()) - else: - gap_bytes = max(0, fa._byte_low_frontier() - sa._byte_high_frontier()) - R_f = fa.num_pages - fa.min_page_index - fa._allocated_pages() - R_s = sa.num_pages - sa.min_page_index - sa._allocated_pages() - - if not self.lazy_compaction: - pages_by_bytes = gap_bytes // (e_f + e_s) - return min(pages_by_bytes, R_f, R_s) * self.page_size - - H_f = len(fa._free_phys_pages) - H_s = len(sa._free_phys_pages) - - K1 = min(H_f, H_s) # Phase 1: both drain - - # Phase 2: fewer-holes side extends; more-holes side keeps draining. - if H_f <= H_s: - e_phase2 = e_f - K_phase2_max = H_s - else: - e_phase2 = e_s - K_phase2_max = H_f - K2_room = K_phase2_max - K1 - K2 = min(K2_room, gap_bytes // e_phase2) if e_phase2 > 0 else K2_room - gap_bytes -= K2 * e_phase2 - - K3 = gap_bytes // (e_f + e_s) # Phase 3: both extend - - K_total = K1 + K2 + K3 - K_total = min(K_total, H_f + R_f, H_s + R_s) # index-space caps - return K_total * self.page_size - # Slot-conservation views for the leak invariant only; the byte-coordinated # value would flag spurious leaks. `allocated_count()` is in TOKENS. def _conserve_full_available_size(self) -> int: @@ -282,15 +230,6 @@ def schedulable_full_available_size(self) -> int: def schedulable_swa_available_size(self) -> int: return self.swa_attn_allocator.schedulable_available_size() - def _flush_targets(self): - """Flush ALL members, including ones that are not short themselves: a - one-sided hole is unusable, and compacting it yields SHARED gap.""" - return (self.full_attn_allocator, self.swa_attn_allocator) - - def _ask_float_for_room(self, need_tokens: int) -> None: - """No float in a two-END chain -- nothing can slide.""" - return None - # `size_full` / `size_swa` bound each side independently; current capacities # also account for the peer's live byte usage. @@ -302,390 +241,102 @@ def current_full_capacity(self) -> int: def current_swa_capacity(self) -> int: return self.swa_available_size() + self.swa_attn_allocator.allocated_count() - def reclaim_plan( - self, - full_tokens: int | float, - swa_tokens: int | float, - *, - full_evictable_tokens: int = 0, - swa_evictable_tokens: int = 0, - empty_pool: bool = False, - ) -> Optional[Tuple[int, int]]: - """Return cumulative FULL/SWA eviction targets, or None if impossible. + @property + def draft_virtual_id_space(self) -> int: + return self.full_attn_allocator.max_slots - 1 - The tree evicts FULL before SWA. Minimize required SWA reclaim with all - evictable FULL available, then trim excess FULL reclaim. FULL eviction's - actual SWA cascade is counted by the tree's shared eviction tracker. - """ - if full_tokens < 0 or swa_tokens < 0: - return None + def debug_print(self) -> str: + return ( + f"#full-available={self.full_attn_allocator.available_size()}, " + f"#swa-available={self.swa_attn_allocator.available_size()}, " + f"#joint-available={self.available_size()}" + ) - page_size = self.page_size - full_pages = (math.ceil(full_tokens) + page_size - 1) // page_size - swa_pages = (math.ceil(swa_tokens) + page_size - 1) // page_size - # Restoring host SWA for device-resident FULL can require more new - # SWA pages than new FULL pages; only the shared budget constrains it. + def get_kvcache(self): + return self._kvcache - compacted = empty_pool or not self.lazy_compaction or self._compaction_allowed() + def translate_kv_loc( + self, + loc: torch.Tensor, + *, + out: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """Full-layer read path: virtual TOKEN ids -> full-physical TOKEN ids. + ``out=`` writes in place, for cuda-graph buffer stability.""" + result = self.full_attn_allocator.translate_kv_loc(loc, out=out) + return result - def fits(full_reclaim_pages: int, swa_reclaim_pages: int) -> bool: - return self._fits_page_demand( - full_pages, - swa_pages, - full_reclaim_pages=full_reclaim_pages, - swa_reclaim_pages=swa_reclaim_pages, - compacted=compacted, - empty_pool=empty_pool, - ) + def translate_loc_from_full_to_swa( + self, + kv_indices: torch.Tensor, + *, + out: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """SWA-layer read path: virtual TOKEN ids -> swa kernel-facing ids.""" + return self.swa_attn_allocator.translate_kv_loc_for_kernel(kv_indices, out=out) - if empty_pool: - return (0, 0) if fits(0, 0) else None - if fits(0, 0): - return (0, 0) + @property + def kernel_page_multiplier(self) -> int: + return self.full_attn_allocator.kernel_page_multiplier - max_full_pages = min( - self.full_attn_allocator.allocated_count() // page_size, - max(0, int(full_evictable_tokens)) // page_size, - ) - max_swa_pages = min( - self.swa_attn_allocator.allocated_count() // page_size, - max(0, int(swa_evictable_tokens)) // page_size, - ) - if not fits(max_full_pages, max_swa_pages): - return None + @property + def full_v2p_page_table(self) -> torch.Tensor: + """Page-level virtual->physical table of the full sub-pool.""" + return self.full_attn_allocator.virtual_to_physical - def first_fit(high: int, predicate: Callable[[int], bool]) -> int: - low = 0 - while low < high: - mid = (low + high) // 2 - if predicate(mid): - high = mid - else: - low = mid + 1 - return low + @property + def full_p2v_page_table(self) -> torch.Tensor: + """Page-level physical->virtual table of the full sub-pool.""" + return self.full_attn_allocator.physical_to_virtual - swa_reclaim_pages = first_fit( - max_swa_pages, lambda value: fits(max_full_pages, value) - ) - full_reclaim_pages = first_fit( - max_full_pages, lambda value: fits(value, swa_reclaim_pages) - ) - return ( - full_reclaim_pages * page_size, - swa_reclaim_pages * page_size, - ) + def translate_kv_indices_for_transfer( + self, kv_indices: torch.Tensor + ) -> torch.Tensor: + """Virtual TOKEN ids -> full-sub-pool PHYSICAL token ids for the PD + transfer engine. - def can_reserve( - self, - full_tokens: int | float, - swa_tokens: int | float, - *, - full_evictable_tokens: int = 0, - swa_evictable_tokens: int = 0, - empty_pool: bool = False, - require_token_slack: bool = False, - ) -> bool: - """Check pending FULL/SWA demand against the shared byte envelope. + PHYSICAL, not kernel-facing: the transfer registers page ENVELOPES (see + `UnifiedMHATokenToKVPool.get_contiguous_buf_infos`). Without this + override the base identity would put VIRTUAL ids on the wire, which + address real bytes and so corrupt silently rather than fail. + """ + return self.full_attn_allocator.translate_kv_loc(kv_indices.to(torch.int64)) - Scheduler admission keeps the historical one-token strict slack at an - empty-pool boundary. Live admission checks the state reachable after - reclaiming the currently evictable pages from both sides. + def translate_swa_indices_for_transfer( + self, kv_indices: torch.Tensor + ) -> torch.Tensor: + """Virtual TOKEN ids -> swa-sub-pool PHYSICAL token ids. + + The SWA counterpart of the above. `translate_loc_from_full_to_swa` + cannot serve here: it returns KERNEL-FACING ids (the physical page + scaled by the sub-pool's per-page block count), which index the + per-layer views, whereas the SWA state component is registered as whole + page envelopes and addressed by physical page. """ - if full_tokens < 0 or swa_tokens < 0: - return False - if require_token_slack and ( - full_tokens >= self.size_full or swa_tokens >= self.size_swa - ): - return False + return self.swa_attn_allocator.translate_kv_loc(kv_indices.to(torch.int64)) - return ( - self.reclaim_plan( - full_tokens, - swa_tokens, - full_evictable_tokens=full_evictable_tokens, - swa_evictable_tokens=swa_evictable_tokens, - empty_pool=empty_pool, - ) - is not None - ) + def _move_gate_targets(self): + """Every member a compaction gate must cover. A subclass that adds an + end overrides THIS, and every gate widens with it.""" + return (self.full_attn_allocator, self.swa_attn_allocator) - def _compaction_allowed(self) -> bool: - return all( - allocator.disagg_move_gate is None or allocator.disagg_move_gate() - for allocator in (self.full_attn_allocator, self.swa_attn_allocator) + def set_disagg_move_gate(self, gate: Callable[[], bool]) -> None: + install_move_gate( + self._move_gate_targets(), + slot="disagg_move_gate", + gate=gate, + feature="PD disaggregation", + lazy_compaction=self.lazy_compaction, ) - def _fits_page_demand( + def translate_kv_loc_for_kernel( self, - num_full_pages: int, - num_swa_pages: int, + loc: torch.Tensor, *, - full_reclaim_pages: int = 0, - swa_reclaim_pages: int = 0, - compacted: bool, - empty_pool: bool = False, - ) -> bool: - """Check one FULL/SWA page demand against a single allocator snapshot.""" - if min(num_full_pages, num_swa_pages) < 0: - return False - - fa, sa = self.full_attn_allocator, self.swa_attn_allocator - if empty_pool: - full_live_pages = swa_live_pages = 0 - full_reclaim_pages = swa_reclaim_pages = 0 - else: - full_live_pages = fa.allocated_count() // self.page_size - swa_live_pages = sa.allocated_count() // self.page_size - full_reclaim_pages = min(full_live_pages, max(0, int(full_reclaim_pages))) - swa_reclaim_pages = min(swa_live_pages, max(0, int(swa_reclaim_pages))) - - full_live_pages -= full_reclaim_pages - swa_live_pages -= swa_reclaim_pages - full_total_pages = full_live_pages + num_full_pages - swa_total_pages = swa_live_pages + num_swa_pages - virtual_page_capacity = fa.num_virtual_ids - fa.min_page_index - full_page_capacity = min( - virtual_page_capacity, - fa.num_pages - fa.min_page_index, - ) - swa_page_capacity = min( - virtual_page_capacity, - sa.num_pages - sa.min_page_index, - ) - if full_total_pages > full_page_capacity or swa_total_pages > swa_page_capacity: - return False - - if compacted: - full_virtual_room = virtual_page_capacity - full_live_pages - full_holes = swa_holes = 0 - full_index_room = full_page_capacity - full_live_pages - swa_index_room = swa_page_capacity - swa_live_pages - gap_bytes = ( - self._empty_shared_gap_bytes - - full_live_pages * fa.entry_bytes_per_page - - swa_live_pages * sa.entry_bytes_per_page - ) - else: - full_virtual_room = len(fa.free_virtual_ids) + full_reclaim_pages - full_holes = len(fa._free_phys_pages) + full_reclaim_pages - swa_holes = len(sa._free_phys_pages) + swa_reclaim_pages - full_index_room = fa.num_pages - fa.min_page_index - fa._allocated_pages() - swa_index_room = sa.num_pages - sa.min_page_index - sa._allocated_pages() - gap_bytes = fa._current_gap_bytes() - - if num_full_pages > full_virtual_room: - return False - if num_full_pages > full_holes + full_index_room: - return False - if num_swa_pages > swa_holes + swa_index_room: - return False - full_extensions = max(0, num_full_pages - full_holes) - swa_extensions = max(0, num_swa_pages - swa_holes) - return ( - full_extensions * fa.entry_bytes_per_page - + swa_extensions * sa.entry_bytes_per_page - <= max(0, gap_bytes) - ) - - def ensure_capacity(self, full_tokens: int, swa_tokens: int) -> bool: - """Gate one allocation and compact both sides on shortfall.""" - if full_tokens < 0 or swa_tokens < 0: - return False - if full_tokens == 0 and swa_tokens == 0: - return True - page_size = self.page_size - num_full_pages = (int(full_tokens) + page_size - 1) // page_size - num_swa_pages = (int(swa_tokens) + page_size - 1) // page_size - if self._fits_page_demand( - num_full_pages, - num_swa_pages, - compacted=False, - ): - return True - if not self.lazy_compaction or not self._compaction_allowed(): - return False - if not self._fits_page_demand( - num_full_pages, - num_swa_pages, - compacted=True, - ): - return False - self.full_attn_allocator.flush_for_allocation() - self.swa_attn_allocator.flush_for_allocation() - return self._fits_page_demand( - num_full_pages, - num_swa_pages, - compacted=False, - ) - - def create_prefill_budget(self, tree_cache, *, num_mixed_decode_tokens=0): - from sglang.srt.mem_cache.prefill_budget import SharedSWAPrefillBudget - - return SharedSWAPrefillBudget( - self, tree_cache, num_mixed_decode_tokens=num_mixed_decode_tokens - ) - - def swa_capacity_and_available(self, *, full_capacity, swa_capacity): - return ( - (self.current_full_capacity, self.full_available_size()), - (self.current_swa_capacity, self.swa_available_size()), - ) - - def max_new_tokens_for_memory( - self, - input_tokens: int, - max_new_tokens: int, - *, - token_capacity: int, - sliding_window_size: int | None, - chunk_size: int | None, - ) -> int | None: - from sglang.srt.mem_cache.prefill_budget import estimate_swa_kv_tokens - - def fits(candidate): - return self.can_reserve( - input_tokens + candidate + self.page_size, - estimate_swa_kv_tokens( - input_tokens, - candidate, - sliding_window_size=sliding_window_size, - page_size=self.page_size, - allocation_limit=chunk_size, - ), - empty_pool=True, - require_token_slack=True, - ) - - if not fits(0): - return None - if fits(max_new_tokens): - return max_new_tokens - lo, hi = 0, max_new_tokens - while lo < hi: - mid = (lo + hi + 1) // 2 - if fits(mid): - lo = mid - else: - hi = mid - 1 - return lo - - def evict_to_free_tokens(self, tree_cache, num_tokens: int) -> None: - from sglang.srt.mem_cache.base_prefix_cache import EvictParams - - if tree_cache is None or tree_cache.is_chunk_cache(): - return - reclaim_plan = self.reclaim_plan( - num_tokens, - num_tokens, - full_evictable_tokens=tree_cache.full_evictable_size(), - swa_evictable_tokens=tree_cache.swa_evictable_size(), - ) - if reclaim_plan is None: - return - full_reclaim, swa_reclaim = reclaim_plan - if full_reclaim or swa_reclaim: - tree_cache.evict_for_alloc( - EvictParams(num_tokens=full_reclaim, swa_num_tokens=swa_reclaim) - ) - # A zero-reclaim plan can still depend on compaction before allocation. - self.ensure_capacity(num_tokens, num_tokens) - - @property - def draft_virtual_id_space(self) -> int: - return self.full_attn_allocator.max_slots - 1 - - def debug_print(self) -> str: - return ( - f"#full-available={self.full_attn_allocator.available_size()}, " - f"#swa-available={self.swa_attn_allocator.available_size()}, " - f"#joint-available={self.available_size()}" - ) - - def get_kvcache(self): - return self._kvcache - - def translate_kv_loc( - self, - loc: torch.Tensor, - *, - out: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - """Full-layer read path: virtual TOKEN ids -> full-physical TOKEN ids. - ``out=`` writes in place, for cuda-graph buffer stability.""" - result = self.full_attn_allocator.translate_kv_loc(loc, out=out) - return result - - def translate_loc_from_full_to_swa( - self, - kv_indices: torch.Tensor, - *, - out: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - """SWA-layer read path: virtual TOKEN ids -> swa kernel-facing ids.""" - return self.swa_attn_allocator.translate_kv_loc_for_kernel(kv_indices, out=out) - - @property - def kernel_page_multiplier(self) -> int: - return self.full_attn_allocator.kernel_page_multiplier - - @property - def full_v2p_page_table(self) -> torch.Tensor: - """Page-level virtual->physical table of the full sub-pool.""" - return self.full_attn_allocator.virtual_to_physical - - @property - def full_p2v_page_table(self) -> torch.Tensor: - """Page-level physical->virtual table of the full sub-pool.""" - return self.full_attn_allocator.physical_to_virtual - - def translate_kv_indices_for_transfer( - self, kv_indices: torch.Tensor - ) -> torch.Tensor: - """Virtual TOKEN ids -> full-sub-pool PHYSICAL token ids for the PD - transfer engine. - - PHYSICAL, not kernel-facing: the transfer registers page ENVELOPES (see - `UnifiedMHATokenToKVPool.get_contiguous_buf_infos`). Without this - override the base identity would put VIRTUAL ids on the wire, which - address real bytes and so corrupt silently rather than fail. - """ - return self.full_attn_allocator.translate_kv_loc(kv_indices.to(torch.int64)) - - def translate_swa_indices_for_transfer( - self, kv_indices: torch.Tensor - ) -> torch.Tensor: - """Virtual TOKEN ids -> swa-sub-pool PHYSICAL token ids. - - The SWA counterpart of the above. `translate_loc_from_full_to_swa` - cannot serve here: it returns KERNEL-FACING ids (the physical page - scaled by the sub-pool's per-page block count), which index the - per-layer views, whereas the SWA state component is registered as whole - page envelopes and addressed by physical page. - """ - return self.swa_attn_allocator.translate_kv_loc(kv_indices.to(torch.int64)) - - def _move_gate_targets(self): - """Every member a compaction gate must cover. A subclass that adds an - end overrides THIS, and every gate widens with it.""" - return (self.full_attn_allocator, self.swa_attn_allocator) - - def set_disagg_move_gate(self, gate: Callable[[], bool]) -> None: - install_move_gate( - self._move_gate_targets(), - slot="disagg_move_gate", - gate=gate, - feature="PD disaggregation", - lazy_compaction=self.lazy_compaction, - ) - - def translate_kv_loc_for_kernel( - self, - loc: torch.Tensor, - *, - out: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - """Full-pool virtual TOKEN ids -> kernel-facing ids.""" - return self.full_attn_allocator.translate_kv_loc_for_kernel(loc, out=out) + out: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """Full-pool virtual TOKEN ids -> kernel-facing ids.""" + return self.full_attn_allocator.translate_kv_loc_for_kernel(loc, out=out) def translate_write_loc_for_kernel( self, @@ -1039,14 +690,6 @@ def _release_page_reps(self, pieces: Sequence[torch.Tensor]) -> None: self.full_attn_allocator.clear_inverse_history() self.swa_attn_allocator.clear_inverse_history() - def verify_byte_accounting(self) -> List[str]: - return ( - _chain_byte_accounting_violations( - _end_pair_chain(self.full_attn_allocator, self.swa_attn_allocator) - ) - + self._joint_capacity_memo_violations() - ) - def _joint_capacity_memo_violations(self) -> List[str]: """Idle-time twin of `MultiEndedAllocator._capacity_memo_violations` for the composite joint view. Empty == healthy.""" @@ -1093,6 +736,385 @@ def set_inflight_forward( forward_done, out_cache_loc_virtual ) + @abstractmethod + def _build_swa_attn_allocator(self, **kwargs) -> MultiEndedAllocator: ... + + @abstractmethod + def _wire_peers(self) -> None: ... + + @abstractmethod + def _compute_available_size(self) -> int: ... + + @abstractmethod + def ensure_capacity(self, full_tokens: int, swa_tokens: int) -> bool: ... + + @abstractmethod + def _flush_targets(self): ... + + @abstractmethod + def _ask_float_for_room(self, need_tokens: int) -> None: ... + + +class UnifiedSWATokenToKVPoolAllocator(UnifiedSWAAllocatorBase): + """Two-ended FULL/SWA allocator with asymmetric shared-byte reservations.""" + + def _build_swa_attn_allocator(self, **kwargs) -> MultiEndedAllocator: + """The swa sub-allocator: an END pool in the 2-pool pair.""" + return MultiEndedAllocator( + sub_pool_name="swa", + is_id_owner=False, # non-owner; consumes virtuals minted by full + **kwargs, + ) + + def _wire_peers(self) -> None: + self.full_attn_allocator.bind_peer(self.swa_attn_allocator) + self.swa_attn_allocator.bind_peer(self.full_attn_allocator) + + def _compute_available_size(self) -> int: + """Joint byte budget in TOKENS: each composite alloc(1) consumes one + full-side AND one swa-side page under the same virtual id.""" + fa, sa = self.full_attn_allocator, self.swa_attn_allocator + e_f = fa.entry_bytes_per_page + e_s = sa.entry_bytes_per_page + # Direction-agnostic shared gap: the free byte band between the two pools. + if fa.grow_direction == "up": + gap_bytes = max(0, sa._byte_low_frontier() - fa._byte_high_frontier()) + else: + gap_bytes = max(0, fa._byte_low_frontier() - sa._byte_high_frontier()) + R_f = fa.num_pages - fa.min_page_index - fa._allocated_pages() + R_s = sa.num_pages - sa.min_page_index - sa._allocated_pages() + + if not self.lazy_compaction: + pages_by_bytes = gap_bytes // (e_f + e_s) + return min(pages_by_bytes, R_f, R_s) * self.page_size + + H_f = len(fa._free_phys_pages) + H_s = len(sa._free_phys_pages) + + K1 = min(H_f, H_s) # Phase 1: both drain + + # Phase 2: fewer-holes side extends; more-holes side keeps draining. + if H_f <= H_s: + e_phase2 = e_f + K_phase2_max = H_s + else: + e_phase2 = e_s + K_phase2_max = H_f + K2_room = K_phase2_max - K1 + K2 = min(K2_room, gap_bytes // e_phase2) if e_phase2 > 0 else K2_room + gap_bytes -= K2 * e_phase2 + + K3 = gap_bytes // (e_f + e_s) # Phase 3: both extend + + K_total = K1 + K2 + K3 + K_total = min(K_total, H_f + R_f, H_s + R_s) # index-space caps + return K_total * self.page_size + + def _flush_targets(self): + """Flush ALL members, including ones that are not short themselves: a + one-sided hole is unusable, and compacting it yields SHARED gap.""" + return (self.full_attn_allocator, self.swa_attn_allocator) + + def _ask_float_for_room(self, need_tokens: int) -> None: + """No float in a two-END chain -- nothing can slide.""" + return None + + def reclaim_plan( + self, + full_tokens: int | float, + swa_tokens: int | float, + *, + full_evictable_tokens: int = 0, + swa_evictable_tokens: int = 0, + empty_pool: bool = False, + ) -> Optional[Tuple[int, int]]: + """Return cumulative FULL/SWA eviction targets, or None if impossible. + + The tree evicts FULL before SWA. Minimize required SWA reclaim with all + evictable FULL available, then trim excess FULL reclaim. FULL eviction's + actual SWA cascade is counted by the tree's shared eviction tracker. + """ + if full_tokens < 0 or swa_tokens < 0: + return None + + page_size = self.page_size + full_pages = (math.ceil(full_tokens) + page_size - 1) // page_size + swa_pages = (math.ceil(swa_tokens) + page_size - 1) // page_size + # Restoring host SWA for device-resident FULL can require more new + # SWA pages than new FULL pages; only the shared budget constrains it. + + compacted = empty_pool or not self.lazy_compaction or self._compaction_allowed() + + def fits(full_reclaim_pages: int, swa_reclaim_pages: int) -> bool: + return self._fits_page_demand( + full_pages, + swa_pages, + full_reclaim_pages=full_reclaim_pages, + swa_reclaim_pages=swa_reclaim_pages, + compacted=compacted, + empty_pool=empty_pool, + ) + + if empty_pool: + return (0, 0) if fits(0, 0) else None + if fits(0, 0): + return (0, 0) + + max_full_pages = min( + self.full_attn_allocator.allocated_count() // page_size, + max(0, int(full_evictable_tokens)) // page_size, + ) + max_swa_pages = min( + self.swa_attn_allocator.allocated_count() // page_size, + max(0, int(swa_evictable_tokens)) // page_size, + ) + if not fits(max_full_pages, max_swa_pages): + return None + + def first_fit(high: int, predicate: Callable[[int], bool]) -> int: + low = 0 + while low < high: + mid = (low + high) // 2 + if predicate(mid): + high = mid + else: + low = mid + 1 + return low + + swa_reclaim_pages = first_fit( + max_swa_pages, lambda value: fits(max_full_pages, value) + ) + full_reclaim_pages = first_fit( + max_full_pages, lambda value: fits(value, swa_reclaim_pages) + ) + return ( + full_reclaim_pages * page_size, + swa_reclaim_pages * page_size, + ) + + def can_reserve( + self, + full_tokens: int | float, + swa_tokens: int | float, + *, + full_evictable_tokens: int = 0, + swa_evictable_tokens: int = 0, + empty_pool: bool = False, + require_token_slack: bool = False, + ) -> bool: + """Check pending FULL/SWA demand against the shared byte envelope. + + Scheduler admission keeps the historical one-token strict slack at an + empty-pool boundary. Live admission checks the state reachable after + reclaiming the currently evictable pages from both sides. + """ + if full_tokens < 0 or swa_tokens < 0: + return False + if require_token_slack and ( + full_tokens >= self.size_full or swa_tokens >= self.size_swa + ): + return False + + return ( + self.reclaim_plan( + full_tokens, + swa_tokens, + full_evictable_tokens=full_evictable_tokens, + swa_evictable_tokens=swa_evictable_tokens, + empty_pool=empty_pool, + ) + is not None + ) + + def _compaction_allowed(self) -> bool: + return all( + allocator.disagg_move_gate is None or allocator.disagg_move_gate() + for allocator in (self.full_attn_allocator, self.swa_attn_allocator) + ) + + def _fits_page_demand( + self, + num_full_pages: int, + num_swa_pages: int, + *, + full_reclaim_pages: int = 0, + swa_reclaim_pages: int = 0, + compacted: bool, + empty_pool: bool = False, + ) -> bool: + """Check one FULL/SWA page demand against a single allocator snapshot.""" + if min(num_full_pages, num_swa_pages) < 0: + return False + + fa, sa = self.full_attn_allocator, self.swa_attn_allocator + if empty_pool: + full_live_pages = swa_live_pages = 0 + full_reclaim_pages = swa_reclaim_pages = 0 + else: + full_live_pages = fa.allocated_count() // self.page_size + swa_live_pages = sa.allocated_count() // self.page_size + full_reclaim_pages = min(full_live_pages, max(0, int(full_reclaim_pages))) + swa_reclaim_pages = min(swa_live_pages, max(0, int(swa_reclaim_pages))) + + full_live_pages -= full_reclaim_pages + swa_live_pages -= swa_reclaim_pages + full_total_pages = full_live_pages + num_full_pages + swa_total_pages = swa_live_pages + num_swa_pages + virtual_page_capacity = fa.num_virtual_ids - fa.min_page_index + full_page_capacity = min( + virtual_page_capacity, + fa.num_pages - fa.min_page_index, + ) + swa_page_capacity = min( + virtual_page_capacity, + sa.num_pages - sa.min_page_index, + ) + if full_total_pages > full_page_capacity or swa_total_pages > swa_page_capacity: + return False + + if compacted: + full_virtual_room = virtual_page_capacity - full_live_pages + full_holes = swa_holes = 0 + full_index_room = full_page_capacity - full_live_pages + swa_index_room = swa_page_capacity - swa_live_pages + gap_bytes = ( + self._empty_shared_gap_bytes + - full_live_pages * fa.entry_bytes_per_page + - swa_live_pages * sa.entry_bytes_per_page + ) + else: + full_virtual_room = len(fa.free_virtual_ids) + full_reclaim_pages + full_holes = len(fa._free_phys_pages) + full_reclaim_pages + swa_holes = len(sa._free_phys_pages) + swa_reclaim_pages + full_index_room = fa.num_pages - fa.min_page_index - fa._allocated_pages() + swa_index_room = sa.num_pages - sa.min_page_index - sa._allocated_pages() + gap_bytes = fa._current_gap_bytes() + + if num_full_pages > full_virtual_room: + return False + if num_full_pages > full_holes + full_index_room: + return False + if num_swa_pages > swa_holes + swa_index_room: + return False + full_extensions = max(0, num_full_pages - full_holes) + swa_extensions = max(0, num_swa_pages - swa_holes) + return ( + full_extensions * fa.entry_bytes_per_page + + swa_extensions * sa.entry_bytes_per_page + <= max(0, gap_bytes) + ) + + def ensure_capacity(self, full_tokens: int, swa_tokens: int) -> bool: + """Gate one allocation and compact both sides on shortfall.""" + if full_tokens < 0 or swa_tokens < 0: + return False + if full_tokens == 0 and swa_tokens == 0: + return True + page_size = self.page_size + num_full_pages = (int(full_tokens) + page_size - 1) // page_size + num_swa_pages = (int(swa_tokens) + page_size - 1) // page_size + if self._fits_page_demand( + num_full_pages, + num_swa_pages, + compacted=False, + ): + return True + if not self.lazy_compaction or not self._compaction_allowed(): + return False + if not self._fits_page_demand( + num_full_pages, + num_swa_pages, + compacted=True, + ): + return False + self.full_attn_allocator.flush_for_allocation() + self.swa_attn_allocator.flush_for_allocation() + return self._fits_page_demand( + num_full_pages, + num_swa_pages, + compacted=False, + ) + + def create_prefill_budget(self, tree_cache, *, num_mixed_decode_tokens=0): + from sglang.srt.mem_cache.prefill_budget import SharedSWAPrefillBudget + + return SharedSWAPrefillBudget( + self, tree_cache, num_mixed_decode_tokens=num_mixed_decode_tokens + ) + + def swa_capacity_and_available(self, *, full_capacity, swa_capacity): + return ( + (self.current_full_capacity, self.full_available_size()), + (self.current_swa_capacity, self.swa_available_size()), + ) + + def max_new_tokens_for_memory( + self, + input_tokens: int, + max_new_tokens: int, + *, + token_capacity: int, + sliding_window_size: int | None, + chunk_size: int | None, + ) -> int | None: + from sglang.srt.mem_cache.prefill_budget import estimate_swa_kv_tokens + + def fits(candidate): + return self.can_reserve( + input_tokens + candidate + self.page_size, + estimate_swa_kv_tokens( + input_tokens, + candidate, + sliding_window_size=sliding_window_size, + page_size=self.page_size, + allocation_limit=chunk_size, + ), + empty_pool=True, + require_token_slack=True, + ) + + if not fits(0): + return None + if fits(max_new_tokens): + return max_new_tokens + lo, hi = 0, max_new_tokens + while lo < hi: + mid = (lo + hi + 1) // 2 + if fits(mid): + lo = mid + else: + hi = mid - 1 + return lo + + def evict_to_free_tokens(self, tree_cache, num_tokens: int) -> bool | None: + from sglang.srt.mem_cache.base_prefix_cache import EvictParams + + if tree_cache is None or tree_cache.is_chunk_cache(): + return + reclaim_plan = self.reclaim_plan( + num_tokens, + num_tokens, + full_evictable_tokens=tree_cache.full_evictable_size(), + swa_evictable_tokens=tree_cache.swa_evictable_size(), + ) + if reclaim_plan is None: + return + full_reclaim, swa_reclaim = reclaim_plan + if full_reclaim or swa_reclaim: + tree_cache.evict_for_alloc( + EvictParams(num_tokens=full_reclaim, swa_num_tokens=swa_reclaim) + ) + # A zero-reclaim plan can still depend on compaction before allocation. + return self.ensure_capacity(num_tokens, num_tokens) + + def verify_byte_accounting(self) -> List[str]: + return ( + _chain_byte_accounting_violations( + _end_pair_chain(self.full_attn_allocator, self.swa_attn_allocator) + ) + + self._joint_capacity_memo_violations() + ) + def flush_opportunistic(self) -> int: """Non-urgent flush of BOTH sub-allocators; sync-free.""" with record_function("UnifiedSWAAlloc.flush_opportunistic"): @@ -1108,7 +1130,7 @@ def flush_opportunistic(self) -> int: return fa.flush_opportunistic() + sa.flush_opportunistic() -class UnifiedMambaSWATokenToKVPoolAllocator(UnifiedSWATokenToKVPoolAllocator): +class UnifiedMambaSWATokenToKVPoolAllocator(UnifiedSWAAllocatorBase): """Tri-pool composite for models with full KV + SWA KV + mamba/conv state (both `mambaish_config` and `is_hybrid_swa`). diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index 4810ad31bcab..188ddb3701fc 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/python/sglang/srt/mem_cache/common.py @@ -149,9 +149,11 @@ def maybe_cache_unfinished_req(req: Req, tree_cache: BasePrefixCache, **kwargs): tree_cache.cache_unfinished_req(req, **kwargs) -def evict_from_tree_cache(tree_cache: BasePrefixCache | None, num_tokens: int): +def evict_from_tree_cache( + tree_cache: BasePrefixCache | None, num_tokens: int +) -> bool | None: if tree_cache is not None and not tree_cache.is_chunk_cache(): - tree_cache.token_to_kv_pool_allocator.evict_to_free_tokens( + return tree_cache.token_to_kv_pool_allocator.evict_to_free_tokens( tree_cache, num_tokens ) diff --git a/python/sglang/srt/mem_cache/kv_cache_configurator.py b/python/sglang/srt/mem_cache/kv_cache_configurator.py index 908f43c9a754..8903329d9fc8 100644 --- a/python/sglang/srt/mem_cache/kv_cache_configurator.py +++ b/python/sglang/srt/mem_cache/kv_cache_configurator.py @@ -51,7 +51,7 @@ is_swa_req_ring, ) from sglang.srt.mem_cache.allocator.unified_hybrid_swa import ( - UnifiedSWATokenToKVPoolAllocator, + UnifiedSWAAllocatorBase, ) from sglang.srt.mem_cache.allocator.unified_mamba import ( UnifiedMambaTokenToKVPoolAllocator, @@ -517,7 +517,7 @@ def _init_pools( token_to_kv_pool_allocator, ( UnifiedMambaTokenToKVPoolAllocator, - UnifiedSWATokenToKVPoolAllocator, + UnifiedSWAAllocatorBase, ), ): draft_virtual_id_space = ( @@ -540,7 +540,7 @@ def _init_pools( if ( isinstance( token_to_kv_pool_allocator, - UnifiedSWATokenToKVPoolAllocator, + UnifiedSWAAllocatorBase, ) and self.is_hybrid_swa ): @@ -2143,7 +2143,7 @@ def _build_token_to_kv_pool_allocator( else: swa_allocator = token_to_kv_pool_allocator uses_unified_virtual_ids = isinstance( - swa_allocator, UnifiedSWATokenToKVPoolAllocator + swa_allocator, UnifiedSWAAllocatorBase ) has_draft_swa_layers = ( not self.is_hybrid_swa_mtp_draft or self.draft_swa_full_capacity diff --git a/python/sglang/srt/mem_cache/kv_index_translator.py b/python/sglang/srt/mem_cache/kv_index_translator.py index 1ccb5b35077c..b7c71ae096e4 100644 --- a/python/sglang/srt/mem_cache/kv_index_translator.py +++ b/python/sglang/srt/mem_cache/kv_index_translator.py @@ -65,7 +65,7 @@ build_kv_read_table_packed, ) from sglang.srt.mem_cache.allocator.unified_hybrid_swa import ( - UnifiedSWATokenToKVPoolAllocator, + UnifiedSWAAllocatorBase, ) from sglang.srt.mem_cache.allocator.unified_mamba import ( UnifiedMambaTokenToKVPoolAllocator, @@ -121,7 +121,7 @@ def __init__( self.is_translating = ( isinstance( token_to_kv_pool_allocator, - (UnifiedMambaTokenToKVPoolAllocator, UnifiedSWATokenToKVPoolAllocator), + (UnifiedMambaTokenToKVPoolAllocator, UnifiedSWAAllocatorBase), ) and token_to_kv_pool_allocator.get_kvcache() is token_to_kv_pool ) @@ -139,7 +139,7 @@ def __init__( # DCP read ids stay WIDENED to the consumer: selecting this rank's # share changes the length, so only the production site can do it. self.defer_read_translate = get_parallel().attn_dcp_size > 1 - if isinstance(alloc, UnifiedSWATokenToKVPoolAllocator): + if isinstance(alloc, UnifiedSWAAllocatorBase): self._swa_v2p_table = alloc.swa_v2p_page_table self._swa_page_multiplier = alloc.swa_kernel_page_multiplier self._swa_write_loc_from_full = self._swa_write_loc_unified diff --git a/python/sglang/srt/mem_cache/unified_cache/components/swa.py b/python/sglang/srt/mem_cache/unified_cache/components/swa.py index 31f92dd19e92..f325aa1ef0a9 100644 --- a/python/sglang/srt/mem_cache/unified_cache/components/swa.py +++ b/python/sglang/srt/mem_cache/unified_cache/components/swa.py @@ -226,11 +226,11 @@ def _translate_full_to_swa(self, full_indices: torch.Tensor) -> torch.Tensor: def _unified_allocator(self): """The unified SWA composite, or None when running on the static pool.""" from sglang.srt.mem_cache.allocator.unified_hybrid_swa import ( - UnifiedSWATokenToKVPoolAllocator, + UnifiedSWAAllocatorBase, ) allocator = self.cache.token_to_kv_pool_allocator - if isinstance(allocator, UnifiedSWATokenToKVPoolAllocator): + if isinstance(allocator, UnifiedSWAAllocatorBase): return allocator return None diff --git a/test/registered/unit/mem_cache/test_unified_tri_pool.py b/test/registered/unit/mem_cache/test_unified_tri_pool.py index 755d96848f29..0bf7ecaa644e 100644 --- a/test/registered/unit/mem_cache/test_unified_tri_pool.py +++ b/test/registered/unit/mem_cache/test_unified_tri_pool.py @@ -961,11 +961,19 @@ def _factory_kwargs(self, **over): return kw def test_budget_sizing_and_boot_signature(self): + from sglang.srt.mem_cache.allocator.unified_hybrid_swa import ( + UnifiedSWAAllocatorBase, + UnifiedSWATokenToKVPoolAllocator, + ) + budget = 1 << 20 bundle = init_unified_mamba_swa_pools( **self._factory_kwargs(unified_total_bytes=budget) ) pool = bundle.unified_memory_pool + allocator = bundle.token_to_kv_pool_allocator + self.assertIsInstance(allocator, UnifiedSWAAllocatorBase) + self.assertNotIsInstance(allocator, UnifiedSWATokenToKVPoolAllocator) # Buffer = budget + the state pool's bytes (budget captured AFTER the # state carve-out), never the token-count re-sum. state_bytes = 4 * pool.spec("mamba").entry_bytes() From f0a008571e36e660048228f9f6b7537d5c303ca0 Mon Sep 17 00:00:00 2001 From: Cheng Wan Date: Sun, 13 Sep 2026 19:39:00 -0700 Subject: [PATCH 51/55] Fix shared prefill deferral and capture and draft capacity bounds --- .../srt/mem_cache/kv_cache_configurator.py | 5 +- .../srt/mem_cache/kv_index_translator.py | 5 +- python/sglang/srt/mem_cache/prefill_budget.py | 12 ++-- .../unit/managers/test_prefill_adder.py | 21 +++++- .../mem_cache/test_kv_index_translator.py | 64 +++++++++++++++++++ .../model_executor/test_pool_configurator.py | 24 +++++++ 6 files changed, 123 insertions(+), 8 deletions(-) diff --git a/python/sglang/srt/mem_cache/kv_cache_configurator.py b/python/sglang/srt/mem_cache/kv_cache_configurator.py index 8903329d9fc8..182cc0577c4e 100644 --- a/python/sglang/srt/mem_cache/kv_cache_configurator.py +++ b/python/sglang/srt/mem_cache/kv_cache_configurator.py @@ -441,7 +441,10 @@ def _derive_pool_sizes(self, *, config: MemoryPoolConfig) -> _PoolSizes: max_running_requests=max_running_requests, full_max_total_num_tokens=full_max_total_num_tokens, swa_max_total_num_tokens=swa_max_total_num_tokens, - unified_memory_pool_bytes=config.unified_memory_pool_bytes, + # The target's byte envelope excludes the separate draft allocation. + unified_memory_pool_bytes=( + None if self.is_draft_worker else config.unified_memory_pool_bytes + ), c4_max_total_num_tokens=c4_max_total_num_tokens, c128_max_total_num_tokens=c128_max_total_num_tokens, c4_state_pool_size=c4_state_pool_size, diff --git a/python/sglang/srt/mem_cache/kv_index_translator.py b/python/sglang/srt/mem_cache/kv_index_translator.py index b7c71ae096e4..8d89774a8c16 100644 --- a/python/sglang/srt/mem_cache/kv_index_translator.py +++ b/python/sglang/srt/mem_cache/kv_index_translator.py @@ -127,6 +127,7 @@ def __init__( ) if self.is_translating: alloc = token_to_kv_pool_allocator + self._capture_page_size = alloc.page_size self._full_v2p_table = alloc.full_v2p_page_table self._full_p2v_table = alloc.full_p2v_page_table self._full_page_multiplier = alloc.kernel_page_multiplier @@ -175,10 +176,10 @@ def capture_token_capacity(self, max_token_pool_size: int) -> int: """Host capture rows are indexed by request-token IDs, not kernel IDs. Unified IDs span the whole virtual table even when admission is capped. - The logical page size also covers DCP-widened request-token IDs. + DCP widens allocator pages; the runner's page size stays physical. """ if self.is_translating: - return self._full_v2p_table.numel() * self.page_size + return self._full_v2p_table.numel() * self._capture_page_size return max_token_pool_size + self.page_size # -- per-batch view -------------------------------------------------------- diff --git a/python/sglang/srt/mem_cache/prefill_budget.py b/python/sglang/srt/mem_cache/prefill_budget.py index 965daaa55f9d..caa7cb03e28b 100644 --- a/python/sglang/srt/mem_cache/prefill_budget.py +++ b/python/sglang/srt/mem_cache/prefill_budget.py @@ -20,6 +20,8 @@ from typing import Optional +from sglang.srt.mem_cache.allocator.swa import is_swa_req_ring + def estimate_swa_kv_tokens( extend_input_len: int, @@ -40,8 +42,10 @@ def estimate_swa_kv_tokens( else min(extend_input_len, allocation_limit) ) allocated_tail = max(allocated - sliding_window_size, 0) - # A full window would double-charge short cached-prefix resumes; - # including extend keeps the reservation above the prefill allocation. + # With a roughly two-window SWA pool, a cached prefix can already lock + # one window. Charging another whole window for a short resume can then + # block admission forever on an idle pool. Reserve only its uncached + # tail plus decode headroom, capped at the window. reserved = ( allocated_tail + min(extend_input_len + max_new_tokens, sliding_window_size) @@ -151,7 +155,7 @@ class SWAPrefillBudget(PrefillBudget): def __init__(self, *args, all_swa=False, **kwargs): super().__init__(*args, **kwargs) self.all_swa = all_swa - self.req_ring = getattr(self.allocator, "swa_req_ring", False) is True + self.req_ring = is_swa_req_ring(self.allocator) def _available_and_evictable(self): if self.all_swa: @@ -368,7 +372,7 @@ def has_capacity(self) -> bool: return self._fits(0, 0) def available_chunk_tokens(self, chunk_limit: int) -> int | None: - return chunk_limit + return chunk_limit if chunk_limit > 0 else None def fit_chunk( self, diff --git a/test/registered/unit/managers/test_prefill_adder.py b/test/registered/unit/managers/test_prefill_adder.py index 22aca576ff90..fedf53286a64 100644 --- a/test/registered/unit/managers/test_prefill_adder.py +++ b/test/registered/unit/managers/test_prefill_adder.py @@ -139,7 +139,7 @@ def create_adder(self, running_batch, **kwargs): defaults["token_to_kv_pool_allocator"].page_size = defaults["page_size"] return PrefillAdder(**defaults) - def create_shared_adder(self): + def create_shared_adder(self, *, num_mixed_decode_tokens=0): self.mock_tree_cache.supports_mamba.return_value = False self.mock_tree_cache.sliding_window_size = 8 self.mock_tree_cache.is_tree_cache.return_value = False @@ -166,6 +166,7 @@ def create_shared_adder(self): self.create_running_batch(), page_size=4, rem_chunk_tokens=16, + num_mixed_decode_tokens=num_mixed_decode_tokens, token_to_kv_pool_allocator=allocator, ) @@ -218,6 +219,24 @@ def lock_prefix(_): self.mock_tree_cache.inc_lock_ref.assert_called_once() self.assertEqual(adder.can_run_list, []) + def test_shared_continuation_defers_when_decode_consumes_chunk_budget(self): + """Mixed decode must not commit an empty or negatively sliced prompt.""" + for decode_tokens in (16, 28): + with self.subTest(decode_tokens=decode_tokens): + adder = self.create_shared_adder(num_mixed_decode_tokens=decode_tokens) + req = self.create_shared_req("continuation") + before = ( + adder.memory_budget.total_offset, + adder.memory_budget.swa_offset, + ) + self.assertIs(adder.add_chunked_req(req), req) + self.assertEqual(adder.can_run_list, []) + req.set_extend_range.assert_not_called() + self.assertEqual( + (adder.memory_budget.total_offset, adder.memory_budget.swa_offset), + before, + ) + def test_shared_continuation_uses_memory_chunk_limit(self): adder = self.create_shared_adder() req = self.create_shared_req("continuation", max_new_tokens=80) diff --git a/test/registered/unit/mem_cache/test_kv_index_translator.py b/test/registered/unit/mem_cache/test_kv_index_translator.py index fb497732be86..8f9a92ffe7b0 100644 --- a/test/registered/unit/mem_cache/test_kv_index_translator.py +++ b/test/registered/unit/mem_cache/test_kv_index_translator.py @@ -633,6 +633,70 @@ def test_rebind_translates_full_side_only(self): self.assertTrue(torch.equal(fb.out_cache_loc, want_full)) self.assertTrue(torch.equal(virt, keep)) + def test_capture_capacity_covers_dcp_widened_mamba_ids(self): + """A capturer must store the highest virtual IDs issued under DCP.""" + from sglang.srt.mem_cache.allocator.unified_mamba import ( + UnifiedMambaTokenToKVPoolAllocator, + ) + from sglang.srt.mem_cache.unified_memory_pool import MambaSubPoolSpec + from sglang.srt.runtime_context import get_parallel + + for dcp_size in (1, 2, 4): + with ( + self.subTest(dcp_size=dcp_size), + get_parallel().override(attn_dcp_size=dcp_size), + ): + pool = UnifiedKVPool( + total_bytes=2048, + sub_pool_specs=[ + MHASubPoolSpec( + name="full", + layer_num=1, + head_num=1, + head_dim=4, + store_dtype=torch.float16, + grow_direction="up", + ), + MambaSubPoolSpec( + name="mamba", + layer_num=1, + conv_state_shapes=((2, 2),), + conv_dtype=torch.float16, + temporal_state_shape=(2, 2), + temporal_dtype=torch.float16, + grow_direction="down", + ), + ], + device="cpu", + enable_memory_saver=False, + page_size=4, + ) + allocator = UnifiedMambaTokenToKVPoolAllocator( + unified_buffer=pool, + kvcache=SimpleNamespace(full_kv_pool=None, mamba_pool=None), + device="cpu", + page_size=4, + ) + virt = allocator.alloc(allocator.available_size()) + self.assertIsNotNone(virt) + src = _make_source(allocator, virt[None, :], 4) + cap = object.__new__(BaseTopkCapturer) + cap.topk_size = 1 + expected = torch.arange(len(virt), dtype=torch.int32).reshape(-1, 1, 1) + cap.device_cache = SimpleNamespace(buffer=expected) + cap.host_cache = SimpleNamespace( + buffer=torch.zeros( + src.capture_token_capacity(1), 1, 1, dtype=torch.int32 + ) + ) + fb = _FakeForwardBatch(out_cache_loc=virt) + fb.out_cache_loc_virtual = virt + cap.on_forward_end(fb, False, None, no_copy_to_cpu=False) + req_pool = SimpleNamespace(req_to_token=virt[None, :]) + self.assertTrue( + torch.equal(cap.get_topk(0, len(virt) + 1, req_pool), expected) + ) + def test_topk_capture_round_trips_request_token_ids(self): for ps in (1, 4, 64): for translating in (False, True): diff --git a/test/registered/unit/model_executor/test_pool_configurator.py b/test/registered/unit/model_executor/test_pool_configurator.py index 13963e6a6861..2764bb04ee77 100644 --- a/test/registered/unit/model_executor/test_pool_configurator.py +++ b/test/registered/unit/model_executor/test_pool_configurator.py @@ -339,6 +339,30 @@ def test_memory_utilization(self): self.assertLessEqual(used, available) self.assertGreater(used, available * 0.99) + def test_draft_does_not_inherit_target_shared_byte_budget(self): + """A separate draft pool must not allocate the target's byte envelope again.""" + from sglang.srt.mem_cache.kv_cache_configurator import KVCacheConfigurator + + mr, _, config = self._run(1 << 20, enable_unified_memory=True) + self.assertIsNotNone(config.unified_memory_pool_bytes) + configurator = object.__new__(KVCacheConfigurator) + configurator.model_config = mr.model_config + configurator.is_hybrid_swa = True + configurator.is_draft_worker = False + target = configurator._derive_pool_sizes(config=config) + configurator.is_draft_worker = True + draft = configurator._derive_pool_sizes(config=config) + self.assertEqual( + target.unified_memory_pool_bytes, config.unified_memory_pool_bytes + ) + self.assertIsNone(draft.unified_memory_pool_bytes) + self.assertEqual( + draft.full_max_total_num_tokens, config.full_max_total_num_tokens + ) + self.assertEqual( + draft.swa_max_total_num_tokens, config.swa_max_total_num_tokens + ) + def test_unified_capacity_is_maximal_with_draft_pool(self): page_size = 8 full_layers = 2 From ab2ea29ea34805c0b42f249252057e34d1dea29a Mon Sep 17 00:00:00 2001 From: Cheng Wan Date: Sun, 13 Sep 2026 19:52:48 -0700 Subject: [PATCH 52/55] Validate inherited unified allocator PD and free-group contracts --- .../test_unified_memory_move_gate.py | 110 ++++++++---------- .../test_unified_free_no_host_sync.py | 7 +- .../unit/mem_cache/test_unified_tri_pool.py | 29 ++++- 3 files changed, 84 insertions(+), 62 deletions(-) diff --git a/test/registered/unit/disaggregation/test_unified_memory_move_gate.py b/test/registered/unit/disaggregation/test_unified_memory_move_gate.py index e702a2595c97..2961689662b4 100644 --- a/test/registered/unit/disaggregation/test_unified_memory_move_gate.py +++ b/test/registered/unit/disaggregation/test_unified_memory_move_gate.py @@ -11,13 +11,17 @@ """ import unittest +from types import SimpleNamespace from typing import List, Optional, Set +import torch + from sglang.srt.disaggregation.utils import ( DisaggregationMode, unified_memory_disagg_move_gate, ) from sglang.srt.mem_cache.allocator.unified_sub_pool import MultiEndedAllocator +from sglang.srt.runtime_context import get_parallel from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -187,58 +191,41 @@ def test_null_mode_is_rejected(self): class TestUnifiedAllocatorsPublishTheTransferContract(CustomTestCase): - """Every unified composite allocator must OVERRIDE the two PD hooks. - - `BaseTokenToKVPoolAllocator.translate_kv_indices_for_transfer` is the - IDENTITY, and `set_disagg_move_gate` exists only where a composite defines - it. Inheriting either is silent, not loud: identity puts VIRTUAL ids on the - wire (they address real bytes, so the peer gets plausible garbage), and a - missing gate lets lazy compaction relocate pages under in-flight RDMA. - An AST-level check because instantiating these composites needs a GPU. - """ + """Unified composites must translate virtual IDs before PD transfer. - # Composites that own the full-side virtual ids and so must define the - # transfer translate themselves. - _COMPOSITES = ( - "UnifiedMambaTokenToKVPoolAllocator", - "UnifiedSWATokenToKVPoolAllocator", - ) - # Every composite must define the gate setter, including the tri-pool, - # which inherits the SWA translates (same full side) but has a THIRD - # member the 2-pool setter does not reach. - _GATE_COMPOSITES = _COMPOSITES + ("UnifiedMambaSWATokenToKVPoolAllocator",) + The implementation may be inherited from a shared unified allocator base, + but inheriting the static allocator's identity would put virtual IDs on the + wire and silently corrupt KV. Gate installation must reach every member. + """ @staticmethod - def _own_methods(cls_name: str) -> Set[str]: - """Names this class defines ITSELF, inheritance excluded. - - Resolved off the class object rather than by parsing a named module: - these composites have already been moved once (out of - `multi_ended_allocator` into `allocator/unified_*`), and a hardcoded - module path turns that kind of move into a test failure that says - nothing about the contract. `__dict__` needs no GPU -- it is the class - body, not an instance. - """ - from sglang.srt.mem_cache.allocator import ( - unified_hybrid_swa, - unified_mamba, + def _allocator_class(name): + from sglang.srt.mem_cache.allocator.unified_hybrid_swa import ( + UnifiedMambaSWATokenToKVPoolAllocator, + UnifiedSWATokenToKVPoolAllocator, + ) + from sglang.srt.mem_cache.allocator.unified_mamba import ( + UnifiedMambaTokenToKVPoolAllocator, ) - for mod in (unified_mamba, unified_hybrid_swa): - cls = getattr(mod, cls_name, None) - if cls is not None: - return set(vars(cls)) - raise AssertionError(f"class {cls_name} not found in the unified allocators") + classes = ( + UnifiedMambaTokenToKVPoolAllocator, + UnifiedSWATokenToKVPoolAllocator, + UnifiedMambaSWATokenToKVPoolAllocator, + ) + return {cls.__name__: cls for cls in classes}[name] def test_transfer_translate_is_not_inherited_identity(self): - for name in self._COMPOSITES: - with self.subTest(composite=name): - self.assertIn( - "translate_kv_indices_for_transfer", - self._own_methods(name), - f"{name} inherits the identity transfer translate; PD would " - "ship VIRTUAL ids and corrupt KV without any error", + virtual = torch.tensor([1, 3], dtype=torch.int32) + for name in self._EXPECTED_COVERAGE: + with self.subTest(composite=name), get_parallel().override(attn_dcp_size=1): + alloc = object.__new__(self._allocator_class(name)) + alloc.full_attn_allocator = SimpleNamespace( + translate_kv_loc=lambda ids: ids + 16 ) + physical = alloc.translate_kv_indices_for_transfer(virtual) + self.assertEqual(physical.dtype, torch.int64) + self.assertEqual(physical.tolist(), [17, 19]) # Every sub-allocator attribute a composite can hold. The stub carries all # of them regardless of composite, so the assertion is on what installation @@ -270,11 +257,7 @@ def _members_reached(self, cls_name: str, slot: str) -> Set[str]: `object.__new__` skips `__init__` (which needs a GPU); the setter reads only `lazy_compaction` and the member attributes. """ - from sglang.srt.mem_cache.allocator import unified_hybrid_swa, unified_mamba - - cls = getattr(unified_mamba, cls_name, None) or getattr( - unified_hybrid_swa, cls_name - ) + cls = self._allocator_class(cls_name) alloc = object.__new__(cls) alloc.lazy_compaction = True for attr in self._MEMBER_ATTRS: @@ -313,14 +296,8 @@ def test_gate_setters_do_not_enumerate_members_themselves(self): """ import inspect - from sglang.srt.mem_cache.allocator import unified_hybrid_swa, unified_mamba - for name in self._EXPECTED_COVERAGE: - cls = getattr(unified_mamba, name, None) or getattr( - unified_hybrid_swa, name - ) - if "set_disagg_move_gate" not in vars(cls): - continue # inherited, and the inherited one is checked above + cls = self._allocator_class(name) with self.subTest(composite=name): body = inspect.getsource(cls.set_disagg_move_gate) self.assertIn("install_move_gate", body) @@ -331,10 +308,23 @@ def test_swa_composite_translates_the_swa_side_separately(self): does not name the SWA page holding the same virtual token. The read-path `translate_loc_from_full_to_swa` cannot stand in either: it returns kernel-facing ids, and the transfer addresses raw page envelopes.""" - self.assertIn( - "translate_swa_indices_for_transfer", - self._own_methods("UnifiedSWATokenToKVPoolAllocator"), - ) + virtual = torch.tensor([1, 3], dtype=torch.int32) + for name in ( + "UnifiedSWATokenToKVPoolAllocator", + "UnifiedMambaSWATokenToKVPoolAllocator", + ): + with self.subTest(composite=name), get_parallel().override(attn_dcp_size=1): + alloc = object.__new__(self._allocator_class(name)) + alloc.full_attn_allocator = SimpleNamespace( + translate_kv_loc=lambda ids: ids + 16 + ) + alloc.swa_attn_allocator = SimpleNamespace( + translate_kv_loc=lambda ids: ids + 32, + translate_kv_loc_for_kernel=lambda ids: ids + 64, + ) + physical = alloc.translate_swa_indices_for_transfer(virtual) + self.assertEqual(physical.dtype, torch.int64) + self.assertEqual(physical.tolist(), [33, 35]) class TestEverySwaAllocatorAnswersTheTransferTranslate(CustomTestCase): diff --git a/test/registered/unit/mem_cache/test_unified_free_no_host_sync.py b/test/registered/unit/mem_cache/test_unified_free_no_host_sync.py index cc9453680437..05063d710797 100644 --- a/test/registered/unit/mem_cache/test_unified_free_no_host_sync.py +++ b/test/registered/unit/mem_cache/test_unified_free_no_host_sync.py @@ -327,6 +327,7 @@ def test_all_overridden(self): mea.MultiEndedAllocator, unified_mamba.UnifiedMambaTokenToKVPoolAllocator, unified_hybrid_swa.UnifiedSWATokenToKVPoolAllocator, + unified_hybrid_swa.UnifiedMambaSWATokenToKVPoolAllocator, ): with self.subTest(cls=cls.__name__): self.assertIsNot( @@ -346,9 +347,13 @@ def test_composites_buffer_reps_not_tokens_in_a_group(self): mea.MultiEndedAllocator, unified_mamba.UnifiedMambaTokenToKVPoolAllocator, unified_hybrid_swa.UnifiedSWATokenToKVPoolAllocator, + unified_hybrid_swa.UnifiedMambaSWATokenToKVPoolAllocator, ): with self.subTest(cls=cls.__name__): - self.assertIn("free_page_reps_group", inspect.getsource(cls)) + alloc = object.__new__(cls) + alloc.free_group = None + alloc.free_group_begin() + self.assertEqual(alloc.free_page_reps_group, []) class TestUnifiedSwaFullSideGroup(unittest.TestCase): diff --git a/test/registered/unit/mem_cache/test_unified_tri_pool.py b/test/registered/unit/mem_cache/test_unified_tri_pool.py index 0bf7ecaa644e..688af799012a 100644 --- a/test/registered/unit/mem_cache/test_unified_tri_pool.py +++ b/test/registered/unit/mem_cache/test_unified_tri_pool.py @@ -20,7 +20,7 @@ import inspect import unittest -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch import torch @@ -175,6 +175,33 @@ def test_chain_wiring_and_float_swa(self): self.assertIs(kvcache._full_allocator, fa) self.assertIs(kvcache._swa_allocator, sa) + def test_pd_preallocation_binds_only_swa_tail_pages(self): + _, allocator, _, _ = self._build(page_size=4) + before = allocator.available_size() + prefix = torch.tensor([0], dtype=torch.int64) + seq = torch.tensor([12], dtype=torch.int64) + # With an empty prefix, ordinary allocation supplies the same virtual + # pages without launching the GPU extend kernel. All page binding and + # capacity accounting still run through the real sub-allocators. + full = allocator.full_attn_allocator + with patch.object( + full, "alloc_extend", side_effect=lambda *a, **kw: full.alloc(12) + ): + virtual = allocator.alloc_extend_swa_tail( + prefix, prefix, seq, seq, torch.tensor([-1]), 12, swa_tail_len=5 + ) + self.assertIsNotNone(virtual) + self.assertEqual(len(virtual), 12) + self.assertEqual(allocator.full_attn_allocator.allocated_count(), 12) + self.assertEqual(allocator.swa_attn_allocator.allocated_count(), 8) + swa_pages = allocator.swa_v2p_page_table[virtual[::4] // 4] + self.assertLessEqual(swa_pages[0].item(), 0) + self.assertTrue(torch.all(swa_pages[1:] > 0).item()) + allocator.free(virtual) + self.assertEqual(allocator.full_attn_allocator.allocated_count(), 0) + self.assertEqual(allocator.swa_attn_allocator.allocated_count(), 0) + self.assertEqual(allocator.available_size(), before) + def test_empty_float_is_transparent_to_the_ends(self): _, allocator, _, _ = self._build() fa = allocator.full_attn_allocator From ab7d973d0e7f4e23dcd3f21500ab2e8abacc4c73 Mon Sep 17 00:00:00 2001 From: yhzhuang Date: Mon, 14 Sep 2026 00:43:23 -0700 Subject: [PATCH 53/55] Explain SWA prefill reservation and chunking safety --- python/sglang/srt/mem_cache/prefill_budget.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/python/sglang/srt/mem_cache/prefill_budget.py b/python/sglang/srt/mem_cache/prefill_budget.py index caa7cb03e28b..a1926abe0ebd 100644 --- a/python/sglang/srt/mem_cache/prefill_budget.py +++ b/python/sglang/srt/mem_cache/prefill_budget.py @@ -46,6 +46,8 @@ def estimate_swa_kv_tokens( # one window. Charging another whole window for a short resume can then # block admission forever on an idle pool. Reserve only its uncached # tail plus decode headroom, capped at the window. + # Including the extension also keeps the reservation large enough for + # this pass's prefill allocation. reserved = ( allocated_tail + min(extend_input_len + max_new_tokens, sliding_window_size) @@ -200,6 +202,8 @@ def swa_never_fits(self, extend_input_len, max_new_tokens, **kwargs): ) def _chunk_cap(self, max_new_tokens, swa_host_hit_length=0): + # Only the sliding window stays locked between chunks, so a smaller + # chunk can bound the transient SWA footprint of a longer prompt. headroom = self.swa_tokens( 0, max_new_tokens, swa_host_hit_length=swa_host_hit_length ) From 81d00a3e41254d4333d975e3c5c1e5c54fd1d3f8 Mon Sep 17 00:00:00 2001 From: yhzhuang Date: Mon, 14 Sep 2026 11:53:47 -0700 Subject: [PATCH 54/55] Refresh shared KV budgets during decode admission --- python/sglang/srt/disaggregation/decode.py | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index a77730e7c2a5..3e812f749e2e 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -966,7 +966,10 @@ def resume_retracted_reqs( indices_to_remove.add(i) req.is_retracted = False self._pre_alloc(req) - full_allocatable_tokens -= full_required + full_allocatable_tokens = self._allocatable_token_budgets( + count_retracted=False, + extra_reserved_reqs=len(resumed_reqs), + ) if swa_allocatable_tokens is not None: swa_allocatable_tokens = self._swa_tail_allocatable_token_budget( count_retracted=False, @@ -1843,8 +1846,14 @@ def _swa_tail_allocatable_token_budget( # pool) over-reserves SWA in steady state. Cap by the actual # remaining headroom up to per-req window cap. window_size = self.scheduler.sliding_window_size or 0 - swa_total = self.token_to_kv_pool_allocator.size_swa - swa_available = self.token_to_kv_pool_allocator.swa_available_size() + allocator = self.token_to_kv_pool_allocator + if self._supports_unified_swa_reservation(): + _, (swa_total, swa_available) = allocator.swa_capacity_and_available( + full_capacity=allocator.size_full, swa_capacity=allocator.size_swa + ) + else: + swa_total = allocator.size_swa + swa_available = allocator.swa_available_size() # Per-request SWA ring: cached prefixes still report swa_evictable, but # evicting them frees no ring space. swa_evictable = ( From 405e92b773dca8a201506f84c56f5974cee6b530 Mon Sep 17 00:00:00 2001 From: yhzhuang Date: Mon, 14 Sep 2026 11:54:13 -0700 Subject: [PATCH 55/55] Port profiled SWA byte-budget fix from #36729 Cherry-picked from a6eb82dc77353249fff1b08356533c2dc16ea6ec. --- python/sglang/srt/mem_cache/kv_cache_configurator.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/mem_cache/kv_cache_configurator.py b/python/sglang/srt/mem_cache/kv_cache_configurator.py index a6334254739e..0067e5127691 100644 --- a/python/sglang/srt/mem_cache/kv_cache_configurator.py +++ b/python/sglang/srt/mem_cache/kv_cache_configurator.py @@ -899,7 +899,9 @@ def _init_unified_swa_pools( ] total_bytes = unified_memory_pool_bytes - if total_bytes is None and self.spec_algorithm.is_none(): + # An uncapped, draft-free pool owns the profiled budget, including bytes + # left over after rounding the FULL/SWA boot capacities to pages. + if unified_total_bytes is not None and self.spec_algorithm.is_none(): total_bytes = unified_total_bytes bundle = init_unified_swa_pools(