diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index e06ed21a2..a5e613f08 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -545,6 +545,9 @@ class Envs: # Periodically log lazy-compaction stats per sub-pool (observability only). SGLANG_LOG_LAZY_COMPACTION_STATS = EnvBool(False) SGLANG_LOG_LAZY_COMPACTION_STATS_INTERVAL_SEC = EnvInt(30) + # Per-call move cap on a non-urgent lazy-compaction flush, so a large + # backlog cannot stall the scheduler loop; urgent flushes are uncapped. + SGLANG_LAZY_COMPACTION_MAX_MOVES_PER_CALL = EnvInt(4096) # HND KV layout folds (page, head) into one paged index for per-kv-head sparse # page tables (DP attn); paged backends like trtllm_mha consume it directly. SGLANG_USE_HND_KVCACHE = EnvBool(False) diff --git a/python/sglang/srt/mem_cache/allocator/base.py b/python/sglang/srt/mem_cache/allocator/base.py index 478bcc030..d5efd3b37 100644 --- a/python/sglang/srt/mem_cache/allocator/base.py +++ b/python/sglang/srt/mem_cache/allocator/base.py @@ -122,11 +122,9 @@ class BaseTokenToKVPoolAllocator(abc.ABC): return kv_indices def get_cpu_copy(self, indices, mamba_indices=None): - # FIXME: reuse the get_cpu_copy after paged allocator is implemented raise NotImplementedError() def load_cpu_copy(self, kv_cache_cpu, indices, mamba_indices=None): - # FIXME: reuse the load_cpu_copy after paged allocator is implemented raise NotImplementedError() def alloc_extend(self, *args, **kwargs): diff --git a/python/sglang/srt/mem_cache/allocator/hisparse.py b/python/sglang/srt/mem_cache/allocator/hisparse.py index f43ba2d30..5647154f7 100644 --- a/python/sglang/srt/mem_cache/allocator/hisparse.py +++ b/python/sglang/srt/mem_cache/allocator/hisparse.py @@ -472,7 +472,8 @@ class DeepSeekV4HiSparseTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): self.hisparse_attn_allocator.free(buffer_indices[buffer_indices > 0]) def get_last_loc_compressed(self, last_locs: torch.Tensor): - return (last_locs - 3) // self.compress_ratio + # Last complete C4 block of a prefix of last_loc + 1 tokens; -1 stays -1. + return (last_locs - (self.compress_ratio - 1)) // self.compress_ratio def get_last_loc_hisparse_device(self, last_locs: torch.Tensor): return self.hisparse_kvcache._translate_loc_to_hisparse_device( diff --git a/python/sglang/srt/mem_cache/allocator/paged.py b/python/sglang/srt/mem_cache/allocator/paged.py index 5aaf63aef..0d5ec1d9c 100755 --- a/python/sglang/srt/mem_cache/allocator/paged.py +++ b/python/sglang/srt/mem_cache/allocator/paged.py @@ -15,11 +15,6 @@ limitations under the License. from __future__ import annotations -""" -Page-aligned memory pool. -""" - - from typing import TYPE_CHECKING import torch 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 0b15392ee..fb964cfca 100644 --- a/python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py +++ b/python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py @@ -878,6 +878,7 @@ class UnifiedMambaSWATokenToKVPoolAllocator(UnifiedSWATokenToKVPoolAllocator): 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: diff --git a/python/sglang/srt/mem_cache/allocator/unified_mamba.py b/python/sglang/srt/mem_cache/allocator/unified_mamba.py index 87b1cf52e..5d78eabc5 100644 --- a/python/sglang/srt/mem_cache/allocator/unified_mamba.py +++ b/python/sglang/srt/mem_cache/allocator/unified_mamba.py @@ -306,11 +306,6 @@ class UnifiedMambaTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): self.full_attn_allocator.clear_inverse_history() self.mamba_allocator.clear_inverse_history() - def clear(self) -> None: - self.full_attn_allocator.clear() - self.mamba_allocator.clear() - self.free_group = None - def free_segment(self, free_index: torch.Tensor, *, start_pos: int) -> None: """Fixed-shape counterpart of `free()`; see `MultiEndedAllocator._page_reps`. The mamba sub-pool is slot-granular and untouched by a token free.""" 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 4f87ca90e..9ee3d1685 100644 --- a/python/sglang/srt/mem_cache/allocator/unified_sub_pool.py +++ b/python/sglang/srt/mem_cache/allocator/unified_sub_pool.py @@ -361,8 +361,8 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator): # Per-call move cap on NON-urgent `_flush`: bounds work per `on_idle()` so # a large backlog doesn't block ZMQ IPC. Urgent retries are uncapped. - self._lazy_max_moves_per_call = int( - os.environ.get("SGLANG_LAZY_COMPACTION_MAX_MOVES_PER_CALL", "4096") + self._lazy_max_moves_per_call = ( + envs.SGLANG_LAZY_COMPACTION_MAX_MOVES_PER_CALL.get() ) # Epoch-keyed memos for the capacity views: pure between mutations, but 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 14017da3c..91a1cfcd1 100644 --- a/test/registered/unit/disaggregation/test_unified_memory_move_gate.py +++ b/test/registered/unit/disaggregation/test_unified_memory_move_gate.py @@ -136,7 +136,7 @@ class TestPrefillMoveGate(CustomTestCase): class TestGatedPeerHolesAreNotSchedulable(CustomTestCase): """`schedulable_available_size` credits holes a peer urgent-flush would release. While the move gate is closed that flush relocates nothing, so - crediting them lets the scheduler admit work `_flush_peer_for_alloc` cannot + crediting them lets the scheduler admit work `_relieve_for_alloc` cannot satisfy; the alloc then returns None and the decode prealloc path treats that as a memory-estimation bug and aborts the scheduler. """