diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index d0995088a..5e48302d3 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -1420,7 +1420,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin): ), "req_pool_indices is full! There is a bug in memory estimation." fill_len = self._pre_alloc_fill_len(req) - req.kv_allocated_len = fill_len + req.kv.kv_allocated_len = fill_len req.kv_committed_len = fill_len if prefix_len > 0: @@ -1525,7 +1525,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin): extend_num_tokens=fill_len, swa_tail_len=self._swa_tail_len(fill_len), ) - req.swa_evicted_seqlen = fill_len - self._swa_tail_len(fill_len) + req.kv.swa_evicted_seqlen = fill_len - self._swa_tail_len(fill_len) else: kv_loc = self.token_to_kv_pool_allocator.alloc_extend( prefix_lens=torch.tensor( diff --git a/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_allocator.py b/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_allocator.py index 20d39e7b2..4e702bd4f 100644 --- a/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_allocator.py +++ b/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_allocator.py @@ -663,7 +663,7 @@ class DSV4NPUTokenToKVPoolAllocator(SWATokenToKVPoolAllocator): if req is None or req_to_token_pool is None: return - kv_len = max(req.kv_committed_len, req.kv_allocated_len) + kv_len = max(req.kv_committed_len, req.kv.kv_allocated_len) req_pool_idx = req.req_pool_idx if kv_len <= 0 or req_pool_idx is None: return diff --git a/python/sglang/srt/managers/hisparse_coordinator.py b/python/sglang/srt/managers/hisparse_coordinator.py index 0b9ac7190..68c9d918d 100644 --- a/python/sglang/srt/managers/hisparse_coordinator.py +++ b/python/sglang/srt/managers/hisparse_coordinator.py @@ -268,7 +268,7 @@ class HiSparseCoordinator: """ self.alloc_device_buffer(req) - host_len = self.host_token_len(req.kv_allocated_len) + host_len = self.host_token_len(req.kv.kv_allocated_len) if host_len <= self.device_buffer_size: # Short sequences (seq_len <= device_buffer_size): the kernel fast path # returns device_buffer_locs directly without any host loading, so we @@ -293,7 +293,7 @@ class HiSparseCoordinator: def _preload_to_device_buffer(self, req: Req) -> None: """Preload all tokens from host pool into the device buffer.""" - n = self.host_token_len(req.kv_allocated_len) + n = self.host_token_len(req.kv.kv_allocated_len) host_indices = self.req_to_host_pool[req.req_pool_idx, :n] device_locs = self.req_to_device_buffer[req.req_pool_idx, :n] @@ -311,7 +311,7 @@ class HiSparseCoordinator: allocated_len = req.extend_range.end alloc_size = self.padded_buffer_size else: - allocated_len = req.kv_allocated_len + allocated_len = req.kv.kv_allocated_len page_size = self.mem_pool_device.page_size # Allocate only enough for current tokens (page-aligned). # When prefill already fills device_buffer_size, include the reserved page. @@ -766,7 +766,7 @@ class HiSparseCoordinator: # we just freed via free_hisparse_indices(all_hi). If left set, the # subsequent release_kv_cache -> allocator.free -> free_hisparse path # re-frees them (double-free into the page allocator's free list). - allocated_len = req.kv_allocated_len + allocated_len = req.kv.kv_allocated_len # release memory -- only free actually-allocated buffer indices current_cap = int(self.req_device_buffer_size[req.req_pool_idx]) diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 83548ae0b..46fa64016 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -661,6 +661,17 @@ class ReqLogprob: output_token_ids_logprobs_idx: Optional[list] = None +@dataclasses.dataclass(slots=True, kw_only=True) +class ReqKvInfo: + kv_allocated_len: int + # The length of KV that have been removed in swa cache. + # SWA KV cache eviction behavior differs by cache type: + # - Radix cache: KV in range [cache_protected_len, swa_evicted_seqlen) is freed manually in + # `ScheduleBatch.maybe_evict_swa`; KV in range [0, cache_protected_len) is freed during radix cache eviction. + # - Chunk cache: KV in range [0, swa_evicted_seqlen) is freed manually in `ScheduleBatch.maybe_evict_swa`. + swa_evicted_seqlen: int + + class Req(ReqDllmMixin): """The input and output status of a request.""" @@ -737,19 +748,13 @@ class Req(ReqDllmMixin): # For req-level memory management self.kv_committed_len = 0 - self.kv_allocated_len = 0 + self.kv: ReqKvInfo = ReqKvInfo(kv_allocated_len=0, swa_evicted_seqlen=0) self.kv_committed_freed = False self.kv_overallocated_freed = False # for cross-encoder model self.token_type_ids = token_type_ids - # The length of KV that have been removed in swa cache. - # SWA KV cache eviction behavior differs by cache type: - # - Radix cache: KV in range [cache_protected_len, swa_evicted_seqlen) is freed manually in - # `ScheduleBatch.maybe_evict_swa`; KV in range [0, cache_protected_len) is freed during radix cache eviction. - # - Chunk cache: KV in range [0, swa_evicted_seqlen) is freed manually in `ScheduleBatch.maybe_evict_swa`. - self.swa_evicted_seqlen = 0 # Tokens in [0, swa_evict_floor) are protected from SWA window eviction. # This is used by prefill-aware SWA models such as Unlimited-OCR to keep prompt/image KV visible during decode. self.swa_evict_floor: int = 0 @@ -1094,9 +1099,9 @@ class Req(ReqDllmMixin): # e.g., speculative decoding may allocate more KV cache than actually used. assert ( not self.kv_overallocated_freed - ), f"Overallocated KV cache already freed, {self.kv_committed_len=}, {self.kv_allocated_len=}" + ), f"Overallocated KV cache already freed, {self.kv_committed_len=}, {self.kv.kv_allocated_len=}" self.kv_overallocated_freed = True - return self._cache_commit_len(), self.kv_allocated_len + return self._cache_commit_len(), self.kv.kv_allocated_len def update_spec_correct_drafts_histogram(self, num_correct_drafts: int): """Update the speculative decoding acceptance histogram. @@ -1518,11 +1523,11 @@ class Req(ReqDllmMixin): self.mamba_cow_src_index = None self.mamba_needs_clear = False self.already_computed = 0 - self.kv_allocated_len = 0 + self.kv.kv_allocated_len = 0 self.kv_committed_len = 0 self.kv_committed_freed = False self.kv_overallocated_freed = False - self.swa_evicted_seqlen = 0 + self.kv.swa_evicted_seqlen = 0 self.extend_batch_idx = 0 self.decode_batch_idx = 0 @@ -2190,7 +2195,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): # update req-level memory management fields req.kv_committed_len = seq_len - req.kv_allocated_len = seq_len + req.kv.kv_allocated_len = seq_len # If input_embeds are available, store them if req.input_embeds is not None: @@ -2558,8 +2563,8 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): reserve = get_alloc_reserve_per_decode() total = 0 for r in requests: - x = max(0, r.kv_committed_len + reserve - r.kv_allocated_len) - cur = r.kv_allocated_len + x = max(0, r.kv_committed_len + reserve - r.kv.kv_allocated_len) + cur = r.kv.kv_allocated_len nxt = cur + x total += ceil_align(nxt, page_size) - ceil_align(cur, page_size) return total @@ -2790,7 +2795,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): for req in self.reqs: req.decode_batch_idx += 1 req.kv_committed_len += 1 - req.kv_allocated_len += 1 + req.kv.kv_allocated_len += 1 # New-tensor avoids racing model_worker_batch refs queued for # overlap forward. diff --git a/python/sglang/srt/managers/scheduler_components/invariant_checker.py b/python/sglang/srt/managers/scheduler_components/invariant_checker.py index bdd1b7019..b7e0c5123 100644 --- a/python/sglang/srt/managers/scheduler_components/invariant_checker.py +++ b/python/sglang/srt/managers/scheduler_components/invariant_checker.py @@ -250,7 +250,7 @@ class SchedulerInvariantChecker: if req.kv_committed_freed or req.req_pool_idx is None: continue - allocated_len = req.kv_allocated_len + allocated_len = req.kv.kv_allocated_len if self.page_size > 1: allocated_len = ceil_align(allocated_len, self.page_size) assert req.cache_protected_len % self.page_size == 0 @@ -258,7 +258,7 @@ class SchedulerInvariantChecker: full_uncached += allocated_len - req.cache_protected_len if self.is_hybrid_swa: swa_uncached += allocated_len - max( - req.cache_protected_len, req.swa_evicted_seqlen + req.cache_protected_len, req.kv.swa_evicted_seqlen ) return full_uncached, swa_uncached @@ -321,7 +321,7 @@ class SchedulerInvariantChecker: f"req {req.rid}", req.req_pool_idx, req.kv_committed_len, - req.kv_allocated_len, + req.kv.kv_allocated_len, ) sess = getattr(self.tree_cache, "slots", None) if sess: @@ -332,7 +332,7 @@ class SchedulerInvariantChecker: f"slot {sid[:8]}", slot.req_pool_idx, slot.kv_committed_len, - slot.kv_allocated_len, + slot.kv.kv_allocated_len, ) active = [ diff --git a/python/sglang/srt/mem_cache/chunk_cache.py b/python/sglang/srt/mem_cache/chunk_cache.py index 7bc00000a..e155c5553 100644 --- a/python/sglang/srt/mem_cache/chunk_cache.py +++ b/python/sglang/srt/mem_cache/chunk_cache.py @@ -156,7 +156,7 @@ class PureSWAChunkCache(SWAChunkCache): req.req_pool_idx, :kv_committed_len ] evict_floor = req.swa_evict_floor - evicted_seqlen = req.swa_evicted_seqlen + evicted_seqlen = req.kv.swa_evicted_seqlen if evicted_seqlen > evict_floor: parts = [] if evict_floor > 0: diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index 902519614..a0a7f5708 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/python/sglang/srt/mem_cache/common.py @@ -77,7 +77,7 @@ def free_swa_out_of_window_slots( evict_floor = max(req.cache_protected_len, getattr(req, "swa_evict_floor", 0)) if page_size > 1 and evict_floor > req.cache_protected_len: evict_floor = -(-evict_floor // page_size) * page_size - req.swa_evicted_seqlen = max(req.swa_evicted_seqlen, evict_floor) + req.kv.swa_evicted_seqlen = max(req.kv.swa_evicted_seqlen, evict_floor) if is_chunk_cache: # Chunk cache builds no radix tree, so no tombstone-leaf concern; evict @@ -90,22 +90,22 @@ def free_swa_out_of_window_slots( # No extra page margin is needed. evict_threshold = pre_len - max(sliding_window_size, page_size) new_swa_evicted_seqlen = max( - req.swa_evicted_seqlen, + req.kv.swa_evicted_seqlen, evict_threshold, ) if page_size > 1: new_swa_evicted_seqlen = (new_swa_evicted_seqlen // page_size) * page_size - if new_swa_evicted_seqlen > req.swa_evicted_seqlen: + if new_swa_evicted_seqlen > req.kv.swa_evicted_seqlen: free_slots = req_to_token_pool.req_to_token[ - req.req_pool_idx, req.swa_evicted_seqlen : new_swa_evicted_seqlen + req.req_pool_idx, req.kv.swa_evicted_seqlen : new_swa_evicted_seqlen ] token_to_kv_pool_allocator.free_swa(free_slots) maybe_evict_dsv4_state_on_swa( token_to_kv_pool_allocator, req_to_token_pool, req, new_swa_evicted_seqlen ) - req.swa_evicted_seqlen = new_swa_evicted_seqlen + req.kv.swa_evicted_seqlen = new_swa_evicted_seqlen def maybe_cache_unfinished_req(req: Req, tree_cache: BasePrefixCache, **kwargs): @@ -661,7 +661,7 @@ def release_kv_cache(req: Req, tree_cache: BasePrefixCache, is_insert: bool = Tr if spec_algo is None and not global_server_args.strip_thinking_cache: assert ( start_p == end_p - ), f"Unexpected overallocated KV cache, {req.kv_committed_len=}, {req.kv_allocated_len=}" + ), f"Unexpected overallocated KV cache, {req.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/pure_swa_radix_cache.py b/python/sglang/srt/mem_cache/pure_swa_radix_cache.py index 309060fe4..639749c3e 100644 --- a/python/sglang/srt/mem_cache/pure_swa_radix_cache.py +++ b/python/sglang/srt/mem_cache/pure_swa_radix_cache.py @@ -93,7 +93,7 @@ class PureSWARadixCache(RadixCache): old_prefix_len = req.cache_protected_len swa_evict_floor = req.swa_evict_floor - swa_evicted_seqlen = req.swa_evicted_seqlen + swa_evicted_seqlen = req.kv.swa_evicted_seqlen if self.page_size > 1 and swa_evict_floor > 0: swa_evict_floor = -(-swa_evict_floor // self.page_size) * self.page_size diff --git a/python/sglang/srt/mem_cache/swa_radix_cache.py b/python/sglang/srt/mem_cache/swa_radix_cache.py index 516516734..ce036adc3 100644 --- a/python/sglang/srt/mem_cache/swa_radix_cache.py +++ b/python/sglang/srt/mem_cache/swa_radix_cache.py @@ -486,7 +486,7 @@ class SWARadixCache(KVCacheEventMixin, BasePrefixCache): key=radix_key, value=values, prev_prefix_len=old_prefix_len, - swa_evicted_seqlen=req.swa_evicted_seqlen, + swa_evicted_seqlen=req.kv.swa_evicted_seqlen, ) ) else: 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 fe005d1af..827b0a8e0 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 @@ -589,7 +589,7 @@ class SWAComponent(TreeComponent): ) -> Optional[int]: # Unfinished requests can already have an SWA-evicted prefix; preserve # that boundary so insertion creates a tombstone instead of live SWA KV. - insert_params.swa_evicted_seqlen = req.swa_evicted_seqlen + insert_params.swa_evicted_seqlen = req.kv.swa_evicted_seqlen return None def free_out_of_window_slots( @@ -604,7 +604,7 @@ class SWAComponent(TreeComponent): req_to_token_pool=self.cache.req_to_token_pool, token_to_kv_pool_allocator=self.cache.token_to_kv_pool_allocator, ) - insert_params.swa_evicted_seqlen = req.swa_evicted_seqlen + insert_params.swa_evicted_seqlen = req.kv.swa_evicted_seqlen # ---- HiCache Hooks ---- diff --git a/python/sglang/srt/session/streaming_session.py b/python/sglang/srt/session/streaming_session.py index 0d0c29b3f..a818aeb7d 100644 --- a/python/sglang/srt/session/streaming_session.py +++ b/python/sglang/srt/session/streaming_session.py @@ -1,5 +1,6 @@ from __future__ import annotations +import copy import logging from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, Dict, Optional @@ -20,7 +21,7 @@ from sglang.srt.mem_cache.base_prefix_cache import ( from sglang.srt.utils.common import ceil_align if TYPE_CHECKING: - from sglang.srt.managers.schedule_batch import Req + from sglang.srt.managers.schedule_batch import Req, ReqKvInfo logger = logging.getLogger(__name__) @@ -36,6 +37,12 @@ class _VirtualNode: pass +def _new_kv() -> ReqKvInfo: + from sglang.srt.managers.schedule_batch import ReqKvInfo + + return ReqKvInfo(kv_allocated_len=0, swa_evicted_seqlen=0) + + @dataclass class SessionSlot: """Holds KV state between streaming session turns.""" @@ -45,16 +52,13 @@ class SessionSlot: # KV pool state (None means no KV is currently held by this slot) req_pool_idx: Optional[int] = None kv_committed_len: int = 0 - kv_allocated_len: int = 0 + kv: ReqKvInfo = field(default_factory=_new_kv) # First req's radix tree node (for dec_lock_ref on session close) last_node: Any = None cache_protected_len: int = 0 swa_uuid_for_lock: Optional[str] = None - # SWA state - swa_evicted_seqlen: int = 0 - # Mamba states mamba_pool_idx: Any = None mamba_ping_pong_track_buffer: Any = None @@ -71,8 +75,7 @@ class SessionSlot: """Save KV state from a finishing request into this slot.""" self.req_pool_idx = req.req_pool_idx self.kv_committed_len = req.kv_committed_len - self.kv_allocated_len = req.kv_allocated_len - self.swa_evicted_seqlen = req.swa_evicted_seqlen + self.kv = copy.copy(req.kv) if is_first: self.last_node = req.last_node @@ -104,8 +107,7 @@ class SessionSlot: """Restore KV state from this slot into an incoming request.""" req.req_pool_idx = self.req_pool_idx req.kv_committed_len = self.kv_committed_len - req.kv_allocated_len = self.kv_allocated_len - req.swa_evicted_seqlen = self.swa_evicted_seqlen + req.kv = copy.copy(self.kv) req.swa_uuid_for_lock = self.swa_uuid_for_lock req.mamba_pool_idx = self.mamba_pool_idx @@ -305,7 +307,7 @@ class StreamingSession(BasePrefixCache): # the mamba pool; otherwise the abort orphans them. slot = SessionSlot( req_pool_idx=req.req_pool_idx, - kv_allocated_len=req.kv_allocated_len, + kv=copy.copy(req.kv), last_node=req.last_node, cache_protected_len=req.cache_protected_len, swa_uuid_for_lock=req.swa_uuid_for_lock, @@ -317,7 +319,9 @@ class StreamingSession(BasePrefixCache): # the abort fall-through doesn't double-free. req.mamba_pool_idx = None req.mamba_ping_pong_track_buffer = None - slot.kv_allocated_len = max(slot.kv_allocated_len, req.kv_allocated_len) + slot.kv.kv_allocated_len = max( + slot.kv.kv_allocated_len, req.kv.kv_allocated_len + ) self.release_session(session_id) req.req_pool_idx = None req.session.abort_req() @@ -339,7 +343,7 @@ class StreamingSession(BasePrefixCache): # req clock (under overlap + honest committed the clock lags the in-flight # verify by ~1, which would short-change inheritance). Clamp to allocated # to keep committed <= allocated for prepare_for_decode. - slot.kv_committed_len = min(target, slot.kv_allocated_len) + slot.kv_committed_len = min(target, slot.kv.kv_allocated_len) # Update req_nodes to this successfully finished request. req.session.finish_req(req) @@ -411,7 +415,9 @@ class StreamingSession(BasePrefixCache): protected_len = slot.cache_protected_len lock_node = slot.last_node tokens_freed = ( - max(0, slot.kv_allocated_len - protected_len) if slot.is_holding_kv else 0 + max(0, slot.kv.kv_allocated_len - protected_len) + if slot.is_holding_kv + else 0 ) logger.info( "Session KV released: %s (%d tokens freed)", session_id, tokens_freed @@ -428,7 +434,7 @@ class StreamingSession(BasePrefixCache): if slot.is_holding_kv: start = protected_len - end = slot.kv_allocated_len + end = slot.kv.kv_allocated_len if start < end: kv_indices = self.req_to_token_pool.req_to_token[ slot.req_pool_idx, start:end @@ -454,7 +460,7 @@ class StreamingSession(BasePrefixCache): active_pool_idxs is not None and slot.req_pool_idx in active_pool_idxs ) if slot.is_holding_kv and not in_batch: - allocated = ceil_align(slot.kv_allocated_len, self.page_size) + allocated = ceil_align(slot.kv.kv_allocated_len, self.page_size) total += allocated - slot.cache_protected_len return total @@ -470,9 +476,9 @@ class StreamingSession(BasePrefixCache): active_pool_idxs is not None and slot.req_pool_idx in active_pool_idxs ) if slot.is_holding_kv and not in_batch: - allocated = ceil_align(slot.kv_allocated_len, self.page_size) + allocated = ceil_align(slot.kv.kv_allocated_len, self.page_size) total += allocated - max( - slot.cache_protected_len, slot.swa_evicted_seqlen + slot.cache_protected_len, slot.kv.swa_evicted_seqlen ) return total @@ -527,13 +533,13 @@ class StreamingSession(BasePrefixCache): decoding pushes allocated above committed, or when retract retry's logit-reserve pulls prefix_len below committed. """ - self._free_kv_aligned(slot.req_pool_idx, prefix_len, slot.kv_allocated_len) - slot.kv_allocated_len = prefix_len + self._free_kv_aligned(slot.req_pool_idx, prefix_len, slot.kv.kv_allocated_len) + slot.kv.kv_allocated_len = prefix_len slot.kv_committed_len = min(slot.kv_committed_len, prefix_len) - slot.swa_evicted_seqlen = min(slot.swa_evicted_seqlen, prefix_len) - req.kv_allocated_len = prefix_len + slot.kv.swa_evicted_seqlen = min(slot.kv.swa_evicted_seqlen, prefix_len) + req.kv.kv_allocated_len = prefix_len req.kv_committed_len = min(req.kv_committed_len, prefix_len) - req.swa_evicted_seqlen = min(req.swa_evicted_seqlen, prefix_len) + req.kv.swa_evicted_seqlen = min(req.kv.swa_evicted_seqlen, prefix_len) def _trim_overshoot(self, req: Req, finished_len: int) -> None: """Trim slot KV to finished_len boundary. Spec v2 may overshoot @@ -542,10 +548,10 @@ class StreamingSession(BasePrefixCache): be released to avoid token/KV mismatch. """ target = len(req.origin_input_ids) + finished_len - self._free_kv_aligned(req.req_pool_idx, target, req.kv_allocated_len) - req.kv_allocated_len = min(req.kv_allocated_len, target) + self._free_kv_aligned(req.req_pool_idx, target, req.kv.kv_allocated_len) + req.kv.kv_allocated_len = min(req.kv.kv_allocated_len, target) req.kv_committed_len = min(req.kv_committed_len, target) - req.swa_evicted_seqlen = min(req.swa_evicted_seqlen, target) + req.kv.swa_evicted_seqlen = min(req.kv.swa_evicted_seqlen, target) req.output_ids = req.output_ids[:finished_len] def _free_kv_aligned(self, pool_idx: int, target: int, end: int) -> None: diff --git a/python/sglang/srt/speculative/dflash_info_v2.py b/python/sglang/srt/speculative/dflash_info_v2.py index 360c80ce1..4c966a0db 100644 --- a/python/sglang/srt/speculative/dflash_info_v2.py +++ b/python/sglang/srt/speculative/dflash_info_v2.py @@ -155,7 +155,7 @@ class DFlashDraftInputV2(SpecInput): for i, req in enumerate(batch.reqs): committed_len = int(req.kv_committed_len) # Read the allocation watermark from the req object like EAGLE. - cur_alloc_len = int(req.kv_allocated_len) + cur_alloc_len = int(req.kv.kv_allocated_len) reserved_len = max(cur_alloc_len, committed_len + 2 * block_size) top_k = int(req.sampling_params.top_k) @@ -233,7 +233,9 @@ class DFlashDraftInputV2(SpecInput): # This request-side high-water mark is what release_kv_cache() uses to # reclaim any DFLASH over-allocation if the request finishes later. for i, req in enumerate(batch.reqs): - req.kv_allocated_len = max(req.kv_allocated_len, int(nxt_kv_lens_cpu_t[i])) + req.kv.kv_allocated_len = max( + req.kv.kv_allocated_len, int(nxt_kv_lens_cpu_t[i]) + ) # Seed committed; overlap's resolve overwrites it with the published value. batch.seq_lens_cpu = batch_seq_lens_cpu_t diff --git a/python/sglang/srt/speculative/eagle_utils.py b/python/sglang/srt/speculative/eagle_utils.py index a25b9cbf7..dad18a053 100644 --- a/python/sglang/srt/speculative/eagle_utils.py +++ b/python/sglang/srt/speculative/eagle_utils.py @@ -833,7 +833,7 @@ def eagle_prepare_for_decode(batch: ScheduleBatch): nxt_kv_lens = [0] * bs num_needed_tokens = 0 for i, r in enumerate(batch.reqs): - cur = r.kv_allocated_len + cur = r.kv.kv_allocated_len # max(cur, ...) clamps so adaptive downswitch cannot make nxt < cur. # kv_committed_len is honest (bonus committed in resolve, not here), # so it lags batch.seq_lens by ~1 verify in overlap; 2*alloc absorbs. @@ -841,7 +841,7 @@ def eagle_prepare_for_decode(batch: ScheduleBatch): cur_kv_lens[i] = cur nxt_kv_lens[i] = nxt num_needed_tokens += nxt - cur - r.kv_allocated_len = nxt + r.kv.kv_allocated_len = nxt r.decode_batch_idx += 1 cur_kv_lens_cpu = torch.tensor(cur_kv_lens, dtype=torch.int32, device="cpu") diff --git a/python/sglang/test/scripted_runtime/req_handle.py b/python/sglang/test/scripted_runtime/req_handle.py index ee68a3a16..a08c1aa8a 100644 --- a/python/sglang/test/scripted_runtime/req_handle.py +++ b/python/sglang/test/scripted_runtime/req_handle.py @@ -45,7 +45,7 @@ class ScriptedReqHandle: if req is None or req.req_pool_idx is None: return 0 page_size = self.context.scheduler.page_size - return (req.kv_allocated_len + page_size - 1) // page_size + return (req.kv.kv_allocated_len + page_size - 1) // page_size @property def lock_refs(self) -> int: diff --git a/test/registered/disaggregation/test_specv2_kvcache_offloading.py b/test/registered/disaggregation/test_specv2_kvcache_offloading.py index 296138eaa..bcb0e096e 100644 --- a/test/registered/disaggregation/test_specv2_kvcache_offloading.py +++ b/test/registered/disaggregation/test_specv2_kvcache_offloading.py @@ -8,6 +8,7 @@ Requires: torch, sglang (run in an environment with sglang installed) """ import unittest +from types import SimpleNamespace from unittest.mock import MagicMock import torch @@ -34,7 +35,7 @@ def _make_mock_req( req.rid = rid req.req_pool_idx = req_pool_idx req.kv_committed_len = kv_committed_len - req.kv_allocated_len = kv_allocated_len + req.kv = SimpleNamespace(kv_allocated_len=kv_allocated_len) req.kv_committed_freed = False req.kv_overallocated_freed = False req.prefix_indices = list(range(prefix_indices_len)) @@ -47,7 +48,7 @@ def _make_mock_req( def pop_overallocated(): assert not req.kv_overallocated_freed req.kv_overallocated_freed = True - return req.kv_committed_len, req.kv_allocated_len + return req.kv_committed_len, req.kv.kv_allocated_len req.pop_committed_kv_cache = pop_committed req.pop_overallocated_kv_cache = pop_overallocated diff --git a/test/registered/unit/managers/test_hisparse_unit.py b/test/registered/unit/managers/test_hisparse_unit.py index 779f8b365..7fbafbae6 100644 --- a/test/registered/unit/managers/test_hisparse_unit.py +++ b/test/registered/unit/managers/test_hisparse_unit.py @@ -50,7 +50,7 @@ def _make_req(rid="test-req-0", origin_input_ids=None, output_ids=None): fill_ids=origin_input_ids + output_ids, seqlen=len(origin_input_ids) + len(output_ids), req_pool_idx=None, - kv_allocated_len=0, + kv=SimpleNamespace(kv_allocated_len=0), kv_committed_len=0, finished_reason=None, hisparse_staging=False, @@ -218,7 +218,7 @@ class TestHiSparseUnit(unittest.TestCase): ) self.assertIsNotNone(kv_loc, "KV alloc failed") self.req_to_token_pool.write((req.req_pool_idx, slice(0, len(kv_loc))), kv_loc) - req.kv_allocated_len = fill_len + req.kv.kv_allocated_len = fill_len req.kv_committed_len = fill_len req.full_untruncated_fill_ids = array("q", range(fill_len)) req.extend_range = Range(0, fill_len) @@ -578,7 +578,7 @@ class TestHiSparseUnit(unittest.TestCase): seq_len = fill_len + 1 self.req_to_token_pool.write((req.req_pool_idx, fill_len), out_loc) - req.kv_allocated_len = seq_len + req.kv.kv_allocated_len = seq_len req.kv_committed_len = seq_len self.coordinator.map_last_loc_to_buffer( @@ -769,7 +769,7 @@ class TestHiSparseUnit(unittest.TestCase): self.coordinator.req_to_host_pool[req.req_pool_idx, :fill_len], ) ) - self.assertEqual(req.kv_allocated_len, fill_len) + self.assertEqual(req.kv.kv_allocated_len, fill_len) self.assertEqual(req.kv_committed_len, fill_len) self.assertEqual(req.extend_range.length, fill_len) @@ -786,7 +786,7 @@ class TestHiSparseUnit(unittest.TestCase): self.assertEqual(allocated_host_indices.numel(), rounded_len) kv_loc = self.req_to_token_pool.req_to_token[ - req.req_pool_idx, : req.kv_allocated_len + req.req_pool_idx, : req.kv.kv_allocated_len ].clone() self._cleanup_req(req, kv_loc, logical_only=True) self._assert_sizes_restored(initial, "pd_decode_prealloc_hisparse") diff --git a/test/registered/unit/managers/test_kv_page_invariants.py b/test/registered/unit/managers/test_kv_page_invariants.py index 939acf1e3..62ba150ca 100644 --- a/test/registered/unit/managers/test_kv_page_invariants.py +++ b/test/registered/unit/managers/test_kv_page_invariants.py @@ -48,14 +48,14 @@ class _FakeReq: self.rid = rid self.req_pool_idx = rpi self.kv_committed_len = committed - self.kv_allocated_len = allocated + self.kv = SimpleNamespace(kv_allocated_len=allocated, swa_evicted_seqlen=0) class _FakeSlot: def __init__(self, rpi, committed, allocated): self.req_pool_idx = rpi self.kv_committed_len = committed - self.kv_allocated_len = allocated + self.kv = SimpleNamespace(kv_allocated_len=allocated, swa_evicted_seqlen=0) self.is_holding_kv = True 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 52d170e71..8336b7a74 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 @@ -27,6 +27,7 @@ register_amd_ci(est_time=10, suite="stage-b-test-1-gpu-small-amd") import unittest from array import array +from types import SimpleNamespace from unittest.mock import MagicMock import torch @@ -81,7 +82,7 @@ class MockReq: self.prefix_indices = torch.empty(0, dtype=torch.int64) self.priority = 0 self.kv_committed_len = len(fill_ids) - self.kv_allocated_len = len(fill_ids) + self.kv = SimpleNamespace(kv_allocated_len=len(fill_ids)) self.kv_committed_freed = False def get_fill_ids(self): @@ -92,7 +93,7 @@ class MockReq: return self.kv_committed_len def pop_overallocated_kv_cache(self): - return (self.kv_committed_len, self.kv_allocated_len) + return (self.kv_committed_len, self.kv.kv_allocated_len) def _make_req(fill_ids, req_pool_idx=0, cache_protected_len=0, last_node=None): diff --git a/test/registered/unit/mem_cache/test_pure_swa_chunk_cache.py b/test/registered/unit/mem_cache/test_pure_swa_chunk_cache.py index 7c8d90f3e..a4cd9c129 100644 --- a/test/registered/unit/mem_cache/test_pure_swa_chunk_cache.py +++ b/test/registered/unit/mem_cache/test_pure_swa_chunk_cache.py @@ -23,7 +23,7 @@ class _FakeAllocator: class _FakeReq: req_pool_idx = 0 swa_evict_floor = 3 - swa_evicted_seqlen = 6 + kv = SimpleNamespace(swa_evicted_seqlen=6) def pop_committed_kv_cache(self): return 8 diff --git a/test/registered/unit/mem_cache/test_streaming_session_unit.py b/test/registered/unit/mem_cache/test_streaming_session_unit.py index b0fe8cd9e..cab86c332 100644 --- a/test/registered/unit/mem_cache/test_streaming_session_unit.py +++ b/test/registered/unit/mem_cache/test_streaming_session_unit.py @@ -57,13 +57,15 @@ class _FakeReq: ) self.req_pool_idx = req_pool_idx self.kv_committed_len = committed - self.kv_allocated_len = allocated + self.kv = SimpleNamespace( + kv_allocated_len=allocated, + swa_evicted_seqlen=0, + ) self.kv_committed_freed = False self.kv_overallocated_freed = False self.origin_input_ids = list(range(committed)) self.output_ids = [] self.extra_key = None - self.swa_evicted_seqlen = 0 self.last_node = None self.cache_protected_len = 0 self.swa_uuid_for_lock = None @@ -86,7 +88,7 @@ class _FakeReq: assert not self.kv_overallocated_freed self.pop_overallocated_calls += 1 self.kv_overallocated_freed = True - return self.kv_committed_len, self.kv_allocated_len + return self.kv_committed_len, self.kv.kv_allocated_len def test_preabort_detaches_session_and_preserves_slot(): @@ -112,7 +114,7 @@ def test_preabort_detaches_session_and_preserves_slot(): tree_cache.slots["session-a"] = SessionSlot( req_pool_idx=0, kv_committed_len=48, - kv_allocated_len=48, + kv=SimpleNamespace(kv_allocated_len=48, swa_evicted_seqlen=0), cache_protected_len=16, ) @@ -132,7 +134,7 @@ def test_preabort_detaches_session_and_preserves_slot(): slot = tree_cache.slots["session-a"] assert slot.req_pool_idx == 0 assert slot.kv_committed_len == 48 - assert slot.kv_allocated_len == 48 + assert slot.kv.kv_allocated_len == 48 assert len(result.device_indices) == 0 @@ -179,7 +181,7 @@ def test_nth_mid_abort_nukes_session_slot(): tree_cache.slots["session-a"] = SessionSlot( req_pool_idx=0, kv_committed_len=50, - kv_allocated_len=50, + kv=SimpleNamespace(kv_allocated_len=50, swa_evicted_seqlen=0), last_node=None, cache_protected_len=0, ) @@ -230,14 +232,14 @@ def test_trim_overshoot_postcondition(): req = _FakeReq("session-a", req_pool_idx=0, committed=40, allocated=44) req.origin_input_ids = list(range(26)) req.output_ids = list(range(14)) - req.swa_evicted_seqlen = 42 + req.kv.swa_evicted_seqlen = 42 tree_cache._trim_overshoot(req, finished_len=12) target = 38 assert req.kv_committed_len == target - assert req.kv_allocated_len == target - assert req.swa_evicted_seqlen == target + assert req.kv.kv_allocated_len == target + assert req.kv.swa_evicted_seqlen == target assert len(req.output_ids) == 12 # Tail [38, 44) freed by _free_kv_aligned. assert len(allocator.freed) == 1 diff --git a/test/registered/unit/mem_cache/test_swa_eviction_boundary.py b/test/registered/unit/mem_cache/test_swa_eviction_boundary.py index 763284ceb..b1f554d5f 100644 --- a/test/registered/unit/mem_cache/test_swa_eviction_boundary.py +++ b/test/registered/unit/mem_cache/test_swa_eviction_boundary.py @@ -106,7 +106,9 @@ def _make_req(req_pool_idx, token_ids, cache_protected_len, tree): origin_input_ids=token_ids, output_ids=[], cache_protected_len=cache_protected_len, - swa_evicted_seqlen=0, + kv=SimpleNamespace( + swa_evicted_seqlen=0, + ), extra_key=None, last_node=tree.root_node, swa_uuid_for_lock=None, @@ -162,7 +164,7 @@ class TestSWAEvictionBoundary(unittest.TestCase): insert_len = seq_len // page_size * page_size self.assertLess( - req.swa_evicted_seqlen, + req.kv.swa_evicted_seqlen, insert_len, f"page={page_size}, win={window}, seq={seq_len}", ) @@ -187,7 +189,7 @@ class TestSWAEvictionBoundary(unittest.TestCase): ScheduleBatch._evict_swa(batch, req, seq_len - 1) insert_len = seq_len // page_size * page_size - self.assertLess(req.swa_evicted_seqlen, insert_len) + self.assertLess(req.kv.swa_evicted_seqlen, insert_len) tree.cache_finished_req(req, is_insert=True) tree.sanity_check() @@ -209,9 +211,9 @@ class TestSWAEvictionBoundary(unittest.TestCase): batch = _make_batch(tree, allocator, pool) ScheduleBatch._evict_swa(batch, req, seq_len - 1) - self.assertLess(req.swa_evicted_seqlen, seq_len) + self.assertLess(req.kv.swa_evicted_seqlen, seq_len) self.assertEqual( - req.swa_evicted_seqlen, max(0, seq_len - 1 - max(window, page_size)) + req.kv.swa_evicted_seqlen, max(0, seq_len - 1 - max(window, page_size)) ) tree.cache_finished_req(req, is_insert=True) @@ -235,7 +237,7 @@ class TestSWAEvictionBoundary(unittest.TestCase): batch = _make_batch(tree, allocator, pool) ScheduleBatch._evict_swa(batch, req, seq_len - 1) - self.assertEqual(req.swa_evicted_seqlen, 0) + self.assertEqual(req.kv.swa_evicted_seqlen, 0) # -- Insert case 1: swa_evicted <= total_prefix_length -- @@ -264,7 +266,7 @@ class TestSWAEvictionBoundary(unittest.TestCase): # pre_len=15: 15-2-8=5, floor to 8 -> 0. Eviction stays within matched. ScheduleBatch._evict_swa(batch, req2, first_len - 1) - self.assertLessEqual(req2.swa_evicted_seqlen, first_len) + self.assertLessEqual(req2.kv.swa_evicted_seqlen, first_len) swa_evictable_before = tree.swa_evictable_size_ tree.cache_finished_req(req2, is_insert=True) @@ -295,12 +297,12 @@ class TestSWAEvictionBoundary(unittest.TestCase): ScheduleBatch._evict_swa(batch, req, seq_len - 1) insert_len = seq_len // page_size * page_size - self.assertGreater(req.swa_evicted_seqlen, 0, "Should have some eviction") - self.assertLess(req.swa_evicted_seqlen, insert_len, "Should be partial") + self.assertGreater(req.kv.swa_evicted_seqlen, 0, "Should have some eviction") + self.assertLess(req.kv.swa_evicted_seqlen, insert_len, "Should be partial") tree.cache_finished_req(req, is_insert=True) - non_tombstone = insert_len - req.swa_evicted_seqlen + non_tombstone = insert_len - req.kv.swa_evicted_seqlen self.assertEqual(tree.swa_evictable_size_, swa_evictable_before + non_tombstone) self.assertGreater(tree.full_evictable_size_, 0) tree.sanity_check() @@ -333,7 +335,7 @@ class TestSWAEvictionBoundary(unittest.TestCase): allocator.free_swa(pool.req_to_token[0, :old_evicted]) req = _make_req(0, list(range(seq_len)), 0, tree) - req.swa_evicted_seqlen = old_evicted + req.kv.swa_evicted_seqlen = old_evicted swa_evictable_before = tree.swa_evictable_size_ tree.cache_finished_req(req, is_insert=True) @@ -363,7 +365,7 @@ class TestSWAEvictionBoundary(unittest.TestCase): ScheduleBatch._evict_swa(batch, req, seq_len - 1) insert_len = seq_len // page_size * page_size - self.assertLess(req.swa_evicted_seqlen, insert_len, f"turn {turn}") + self.assertLess(req.kv.swa_evicted_seqlen, insert_len, f"turn {turn}") tree.cache_finished_req(req, is_insert=True) tree.sanity_check() @@ -386,7 +388,7 @@ class TestSWAEvictionBoundary(unittest.TestCase): ScheduleBatch._evict_swa(batch, req, seq_len - 1) self.assertEqual( - req.swa_evicted_seqlen, max(0, seq_len - 1 - max(window, page_size)) + req.kv.swa_evicted_seqlen, max(0, seq_len - 1 - max(window, page_size)) ) tree.cache_finished_req(req, is_insert=True) diff --git a/test/registered/unit/mem_cache/test_swa_unittest.py b/test/registered/unit/mem_cache/test_swa_unittest.py index 1abd3b135..446277153 100644 --- a/test/registered/unit/mem_cache/test_swa_unittest.py +++ b/test/registered/unit/mem_cache/test_swa_unittest.py @@ -1,5 +1,6 @@ import unittest from array import array +from types import SimpleNamespace import torch @@ -35,6 +36,7 @@ class _DummyReq: def __init__(self): self._kv_committed_len = 0 self.swa_prefix_lock_released = False + self.kv = SimpleNamespace(swa_evicted_seqlen=0) def pop_committed_kv_cache(self): return self._kv_committed_len @@ -575,7 +577,7 @@ class TestSWA(unittest.TestCase): req.extra_key = None req.last_node = tree.root_node req.swa_uuid_for_lock = None - req.swa_evicted_seqlen = 0 + req.kv.swa_evicted_seqlen = 0 req.cache_protected_len = 1 # Intentionally mismatch to ensure code does not use len(prefix_indices). req.prefix_indices = torch.tensor([7, 8, 9, 10, 11], device=tree.device) @@ -610,7 +612,7 @@ class TestSWA(unittest.TestCase): req2.extra_key = None req2.last_node = tree.root_node req2.swa_uuid_for_lock = None - req2.swa_evicted_seqlen = 0 + req2.kv.swa_evicted_seqlen = 0 req2.cache_protected_len = 1 req2.prefix_indices = torch.tensor([21, 22, 23, 24, 25], device=tree.device) 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 7f0c972a8..b0651ab1d 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 @@ -881,7 +881,7 @@ class UnifiedRadixCacheSuite: kv_indices = self._alloc(allocator, kv_len) req_to_token_pool.write((req.req_pool_idx, slice(0, kv_len)), kv_indices) req.kv_committed_len = kv_len - req.kv_allocated_len = kv_len + req.kv.kv_allocated_len = kv_len req.last_node = cache.root_node req.cache_protected_len = 0 req.swa_uuid_for_lock = None @@ -998,7 +998,7 @@ class UnifiedRadixCacheSuite: req.cache_protected_len = 0 req.swa_uuid_for_lock = None req.extra_key = None - req.swa_evicted_seqlen = evicted_len + req.kv.swa_evicted_seqlen = evicted_len cache.cache_unfinished_req(req) @@ -1220,7 +1220,7 @@ class UnifiedRadixCacheSuite: req.cache_protected_len = 0 req.swa_uuid_for_lock = None req.extra_key = None - req.swa_evicted_seqlen = 0 + req.kv.swa_evicted_seqlen = 0 full_available_before_insert = allocator.full_attn_allocator.available_size() @@ -1761,7 +1761,7 @@ class UnifiedRadixCacheSuite: req.cache_protected_len = 0 req.swa_uuid_for_lock = None req.extra_key = None - req.swa_evicted_seqlen = 0 + req.kv.swa_evicted_seqlen = 0 swa_avail_before = allocator.swa_attn_allocator.available_size() @@ -1771,10 +1771,10 @@ class UnifiedRadixCacheSuite: cushion = max(self.cfg.sliding_window_size, self.cfg.page_size) expected_evicted = (pre_len - 1) - cushion self.assertEqual( - req.swa_evicted_seqlen, + req.kv.swa_evicted_seqlen, expected_evicted, f"swa_evicted_seqlen should advance to (pre_len-1) - cushion = " - f"{expected_evicted}, got {req.swa_evicted_seqlen}", + f"{expected_evicted}, got {req.kv.swa_evicted_seqlen}", ) swa_avail_after = allocator.swa_attn_allocator.available_size() @@ -1813,13 +1813,13 @@ class UnifiedRadixCacheSuite: req.cache_protected_len = 0 req.swa_uuid_for_lock = None req.extra_key = None - req.swa_evicted_seqlen = 0 + req.kv.swa_evicted_seqlen = 0 with envs.SGLANG_OPT_UNIFIED_CACHE_FREE_OUT_OF_WINDOW_SLOTS.override(True): cache.cache_unfinished_req(req) self.assertEqual( - req.swa_evicted_seqlen, + req.kv.swa_evicted_seqlen, 0, "Nothing should be evicted when prefill fits inside the cushion", ) diff --git a/test/registered/unit/spec/test_decode_bookkeeping_ownership.py b/test/registered/unit/spec/test_decode_bookkeeping_ownership.py index ba505a990..191d441ea 100644 --- a/test/registered/unit/spec/test_decode_bookkeeping_ownership.py +++ b/test/registered/unit/spec/test_decode_bookkeeping_ownership.py @@ -80,9 +80,7 @@ _OWNER_SITES = { ): 1, # streaming session slot save/restore and tail trimming (_SS, "SessionSlot.save_from_req", "kv_committed_len"): 1, - (_SS, "SessionSlot.save_from_req", "kv_allocated_len"): 1, (_SS, "SessionSlot.restore_to_req", "kv_committed_len"): 1, - (_SS, "SessionSlot.restore_to_req", "kv_allocated_len"): 1, (_SS, "StreamingSession._free_tail", "kv_committed_len"): 2, (_SS, "StreamingSession._free_tail", "kv_allocated_len"): 2, (_SS, "StreamingSession._trim_overshoot", "kv_committed_len"): 1,