diff --git a/python/sglang/srt/mem_cache/kv_cache_configurator.py b/python/sglang/srt/mem_cache/kv_cache_configurator.py index 3cb24e143..3d4d093ab 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 @@ from sglang.srt.mem_cache.memory_pool import ( 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 ( @@ -500,35 +504,43 @@ class KVCacheConfigurator: # 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_hybrid_swa_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: @@ -576,7 +588,8 @@ class KVCacheConfigurator: 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." ) @@ -1979,27 +1992,33 @@ class KVCacheConfigurator: 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). + if isinstance( + token_to_kv_pool_allocator, + DeepSeekV4HiSparseTokenToKVPoolAllocator, + ): + swa_allocator = token_to_kv_pool_allocator.logical_attn_allocator + else: + swa_allocator = token_to_kv_pool_allocator + uses_unified_virtual_ids = isinstance( + swa_allocator, UnifiedSWATokenToKVPoolAllocator + ) + has_draft_swa_layers = ( + 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 + ): + # 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 8371e425f..a5d371395 100644 --- a/python/sglang/srt/mem_cache/multi_ended_allocator.py +++ b/python/sglang/srt/mem_cache/multi_ended_allocator.py @@ -2898,6 +2898,10 @@ class UnifiedMambaTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): # 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 @@ -3401,6 +3405,10 @@ class UnifiedSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator): # (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 305225d9d..dac5dd010 100644 --- a/python/sglang/srt/model_executor/pool_configurator.py +++ b/python/sglang/srt/model_executor/pool_configurator.py @@ -14,6 +14,7 @@ Two entry points, same core computation: from __future__ import annotations import logging +from bisect import bisect_right from dataclasses import dataclass from typing import TYPE_CHECKING, Optional @@ -469,6 +470,7 @@ class HybridSWAPoolConfigurator(MemoryPoolConfigurator): 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. @@ -555,6 +557,9 @@ class HybridSWAPoolConfigurator(MemoryPoolConfigurator): 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 @@ -583,6 +588,50 @@ class HybridSWAPoolConfigurator(MemoryPoolConfigurator): + 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: @@ -634,7 +683,16 @@ class HybridSWAPoolConfigurator(MemoryPoolConfigurator): 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( @@ -720,13 +778,19 @@ class SWAChunkCapPoolConfigurator(HybridSWAPoolConfigurator): * 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, "