diff --git a/python/sglang/srt/beam_search/fork.py b/python/sglang/srt/beam_search/fork.py index 1646b02a2..e54270705 100644 --- a/python/sglang/srt/beam_search/fork.py +++ b/python/sglang/srt/beam_search/fork.py @@ -87,15 +87,14 @@ def free_member_rows(group, req_to_token_pool, token_to_kv_pool_allocator) -> No return leader = group.leader start = group.prompt_len - end = leader.kv.kv_allocated_len if leader.kv is not None else start + end = leader.kv.kv_allocated_len if end > start: # The rewind below is required: without it the leader's own per-Req # release frees this decode region a second time. slots = req_to_token_pool.req_to_token[group.all_rows, start:end] token_to_kv_pool_allocator.free(slots.flatten().unique()) - if leader.kv is not None: - leader.kv_committed_len = start - leader.kv.kv_allocated_len = start + leader.kv_committed_len = start + leader.kv.kv_allocated_len = start req_to_token_pool.free_rows(group.member_rows_cpu.tolist()) group.member_rows = None group.member_rows_cpu = None diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index d32cf5e37..ef200f82d 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -67,7 +67,6 @@ from sglang.srt.environ import envs from sglang.srt.managers.schedule_batch import ( FINISH_ABORT, NextBatchPlan, - ReqKvInfo, ScheduleBatch, ) from sglang.srt.managers.schedule_policy import match_prefix_for_req @@ -1906,10 +1905,7 @@ def alloc_for_decode_prealloc_hisparse( uses_swa_tail: bool, swa_tail_len: int, ) -> torch.Tensor: - if req.kv is None: - req.kv = ReqKvInfo(kv_allocated_len=fill_len, swa_evicted_seqlen=0) - else: - req.kv.kv_allocated_len = fill_len + req.kv.kv_allocated_len = fill_len device = allocator.device prefix_lens = torch.tensor([0], dtype=torch.int64, device=device) prefix_lens_cpu = torch.tensor([0], dtype=torch.int64) @@ -1954,10 +1950,7 @@ def alloc_for_decode_prealloc( swa_tail_len: int, req_to_token_pool: Optional[ReqToTokenPool] = None, ) -> torch.Tensor: - if req.kv is None: - req.kv = ReqKvInfo(kv_allocated_len=fill_len, swa_evicted_seqlen=0) - else: - req.kv.kv_allocated_len = fill_len + req.kv.kv_allocated_len = fill_len if allocator.page_size == 1: kv_loc = allocator.alloc(delta_len) else: diff --git a/python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py b/python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py index 8eecf7e9c..412bb9b80 100644 --- a/python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py +++ b/python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py @@ -281,7 +281,7 @@ class DecodeKVCacheOffloadManager: self.token_to_kv_pool_allocator.free(overalloc_indices) self.req_to_token_pool.free(req) - req.kv = None + req.kv.mark_released() self.tree_cache.protected_size_ -= len(req.prefix_indices) if req.rid in self.offloaded_state: del self.offloaded_state[req.rid] diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index c363a35b5..b2310f133 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -1036,11 +1036,7 @@ class SchedulerDisaggregationPrefillMixin: else: logger.warning(error_message) req.time_stats.trace_ctx.abort(abort_info={"reason": error_message}) - if ( - req.req_pool_idx is not None - or req.kv is not None - or req.mamba_pool_idx is not None - ): + if req.is_holding_kv or req.mamba_pool_idx is not None: release_kv_cache(req, self.tree_cache) maybe_release_metadata_buffer(req, self.req_to_metadata_buffer_idx_allocator) req.pending_bootstrap = False diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 81ba9403a..f3b9e0fd1 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -815,13 +815,23 @@ class ReqLogprob: @dataclasses.dataclass(slots=True, kw_only=True) class ReqKvInfo: - kv_allocated_len: int + # Device KV a request holds outside the prefix cache. Always present on the Req; + # whether any KV is held is `req.req_pool_idx is not None` (Req.is_holding_kv). + kv_allocated_len: int = 0 # 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 + swa_evicted_seqlen: int = 0 + + @property + def is_released(self) -> bool: + return self.kv_allocated_len == 0 and self.swa_evicted_seqlen == 0 + + def mark_released(self) -> None: + self.kv_allocated_len = 0 + self.swa_evicted_seqlen = 0 class Req(ReqDllmMixin): @@ -904,7 +914,7 @@ class Req(ReqDllmMixin): # For req-level memory management self.kv_committed_len = 0 - self.kv: Optional[ReqKvInfo] = None + self.kv = ReqKvInfo() self.retraction_backup: Optional[RetractionBackup] = None # for cross-encoder model @@ -1265,6 +1275,10 @@ class Req(ReqDllmMixin): or self.mamba_host_hit_length > 0 ) + @property + def is_holding_kv(self) -> bool: + return self.req_pool_idx is not None + def effective_kv_committed_len(self) -> int: # Report only the prompt prefix so thinking + answer fall into the # overallocated range and are reclaimed by release_kv_cache. #22373. @@ -1731,7 +1745,7 @@ class Req(ReqDllmMixin): self.mamba_cow_src_index = None self.mamba_needs_clear = False self.already_computed = 0 - assert self.kv is None, "expect it is already released" + assert not self.is_holding_kv, "expect it is already released" self.kv_committed_len = 0 self.extend_batch_idx = 0 self.decode_batch_idx = 0 @@ -3481,7 +3495,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): # seqlen progress is monotonic per KV handle. if ( req.decode_batch_idx >= 1 - and req.kv is not None + and req.is_holding_kv and req.seqlen - 1 - sliding_window_size >= req.kv.swa_evicted_seqlen + eviction_interval ): diff --git a/python/sglang/srt/managers/scheduler_components/invariant_checker.py b/python/sglang/srt/managers/scheduler_components/invariant_checker.py index 551b77d93..2cb0b9b5e 100644 --- a/python/sglang/srt/managers/scheduler_components/invariant_checker.py +++ b/python/sglang/srt/managers/scheduler_components/invariant_checker.py @@ -252,7 +252,7 @@ class SchedulerInvariantChecker: swa_uncached = 0 for batch in batches: for req in batch.reqs: - if req.kv is None: + if not req.is_holding_kv: continue allocated_len = req.kv.kv_allocated_len @@ -324,7 +324,7 @@ class SchedulerInvariantChecker: batch = self.get_last_batch() if batch is not None: for req in batch.reqs: - if req.kv is None: + if not req.is_holding_kv: continue _add_owner( req, diff --git a/python/sglang/srt/managers/scheduler_pp_mixin.py b/python/sglang/srt/managers/scheduler_pp_mixin.py index e7eb92fd3..c58b1f6d3 100644 --- a/python/sglang/srt/managers/scheduler_pp_mixin.py +++ b/python/sglang/srt/managers/scheduler_pp_mixin.py @@ -723,7 +723,7 @@ class SchedulerPPMixin: latencies.append(latency_ms) # Release KV and Mamba cache - if req.req_pool_idx is not None: + if req.is_holding_kv: kv_indices = self.req_to_token_pool.req_to_token[ req.req_pool_idx, : req.extend_range.end ] @@ -731,7 +731,7 @@ class SchedulerPPMixin: if req.mamba_pool_idx is not None: self.req_to_token_pool.free_mamba_cache(req) self.req_to_token_pool.free(req) - req.kv = None + req.kv.mark_released() logger.info( f"[PP Dynamic Chunk] [PP0] Profiled {len(seq_lens)} samples: " diff --git a/python/sglang/srt/mem_cache/allocation.py b/python/sglang/srt/mem_cache/allocation.py index 45ff96aa1..211183b05 100644 --- a/python/sglang/srt/mem_cache/allocation.py +++ b/python/sglang/srt/mem_cache/allocation.py @@ -385,13 +385,8 @@ def alloc_for_extend( batch.seq_lens_cpu, ) - from sglang.srt.managers.schedule_batch import ReqKvInfo - for req, seq_len in zip(batch.reqs, batch.seq_lens_cpu.tolist()): - if req.kv is None: - req.kv = ReqKvInfo(kv_allocated_len=seq_len, swa_evicted_seqlen=0) - else: - req.kv.kv_allocated_len = seq_len + req.kv.kv_allocated_len = seq_len return out_cache_loc, req_pool_indices_device, req_pool_indices_cpu @@ -416,7 +411,7 @@ def _alloc_extend_loc_with_kv_reuse( retained_len = len(req.dllm_incomplete_ids) if extend_len != retained_len: raise RuntimeError("dLLM FDFO retained KV must be reused as a full block.") - if req.kv is None or prefix_len + extend_len > req.kv.kv_allocated_len: + if prefix_len + extend_len > req.kv.kv_allocated_len: raise RuntimeError("dLLM FDFO retained KV is missing.") alloc_extend_lens = [ diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index 524950b3f..46c7a41cc 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/python/sglang/srt/mem_cache/common.py @@ -62,7 +62,7 @@ def free_swa_out_of_window_slots( is_chunk_cache: bool = False, retain_floor: int | None = None, ) -> None: - if req.kv is None: + if not req.is_holding_kv: 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 @@ -200,10 +200,9 @@ def retraction_discard(req: Req, tree_cache: BasePrefixCache, backend: str) -> N def release_kv_cache(req: Req, tree_cache: BasePrefixCache, is_insert: bool = True): - # the two resources currently have the same lifecycle, thus simplify logic below - assert (req.req_pool_idx is None) == (req.kv is None) + assert (not req.is_holding_kv) == req.kv.is_released # MambaRadixCache may alloc mamba state before alloc KV cache - if req.req_pool_idx is None: + if not req.is_holding_kv: assert ( tree_cache.supports_mamba() ), "Only MambaRadixCache allow freeing before alloc" @@ -224,8 +223,8 @@ def release_kv_cache(req: Req, tree_cache: BasePrefixCache, is_insert: bool = Tr # StreamingSession.cache_finished_req handles speculative tail trim # internally, then sets req_pool_idx = None. - assert (req.req_pool_idx is None) == (req.kv is None) - if req.req_pool_idx is None and req.kv is None: + assert (not req.is_holding_kv) == req.kv.is_released + if not req.is_holding_kv: return start_p, end_p = effective_kv_committed_len, req.kv.kv_allocated_len @@ -242,7 +241,7 @@ def release_kv_cache(req: Req, tree_cache: BasePrefixCache, is_insert: bool = Tr # The DSV4-NPU ReqToTokenPool subclass's free() additionally releases the # c4/c128 state pages; other ReqToTokenPool subclasses are a no-op here. tree_cache.req_to_token_pool.free(req) - req.kv = None + req.kv.mark_released() def _release_overallocated_kv_indices( diff --git a/python/sglang/srt/session/streaming_session.py b/python/sglang/srt/session/streaming_session.py index 9489b1e26..c80d49a97 100644 --- a/python/sglang/srt/session/streaming_session.py +++ b/python/sglang/srt/session/streaming_session.py @@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Any, Dict, Optional import torch +from sglang.srt.managers.schedule_batch import ReqKvInfo from sglang.srt.mem_cache.base_prefix_cache import ( BasePrefixCache, DecLockRefParams, @@ -21,7 +22,7 @@ from sglang.srt.mem_cache.base_prefix_cache import ( from sglang.srt.utils.common import ceil_align, is_npu if TYPE_CHECKING: - from sglang.srt.managers.schedule_batch import Req, ReqKvInfo + from sglang.srt.managers.schedule_batch import Req logger = logging.getLogger(__name__) @@ -43,10 +44,10 @@ class SessionSlot: virtual_node: _VirtualNode = field(default_factory=_VirtualNode) - # KV pool state (None means no KV is currently held by this slot) + # KV pool state req_pool_idx: Optional[int] = None kv_committed_len: int = 0 - kv: Optional[ReqKvInfo] = None + kv: ReqKvInfo = field(default_factory=ReqKvInfo) # First req's radix tree node (for dec_lock_ref on session close) last_node: Any = None @@ -67,7 +68,7 @@ class SessionSlot: @property def is_holding_kv(self) -> bool: """Whether this slot currently holds KV pool resources.""" - return self.kv is not None + return self.req_pool_idx is not None def save_from_req(self, req: Req, is_first: bool): """Save KV state from a finishing request into this slot.""" @@ -97,7 +98,7 @@ class SessionSlot: # the slot's tensor to be reused by a new req and leaked when # the slot is later freed. req.req_pool_idx = None - req.kv = None + req.kv = ReqKvInfo() req.mamba_pool_idx = None req.mamba_ping_pong_track_buffer = None req.mamba_next_track_idx = None @@ -221,7 +222,7 @@ class StreamingSession(BasePrefixCache): if not _is_streaming(req): return None slot = self.slots.get(req.session.session_id) - if slot is None or slot.kv is None: + if slot is None or not slot.is_holding_kv: return None if req.to_finish is not None: req.session.abort_req() @@ -350,7 +351,7 @@ class StreamingSession(BasePrefixCache): ) self.release_session(session_id) req.req_pool_idx = None - req.kv = None + req.kv = ReqKvInfo() req.session.abort_req() return True diff --git a/python/sglang/test/scripted_runtime/req_handle.py b/python/sglang/test/scripted_runtime/req_handle.py index 5638ceb51..d3459d8e9 100644 --- a/python/sglang/test/scripted_runtime/req_handle.py +++ b/python/sglang/test/scripted_runtime/req_handle.py @@ -42,7 +42,7 @@ class ScriptedReqHandle: @property def kv_pages(self) -> int: req = self.req - if req is None or req.kv is None: + if req is None or not req.is_holding_kv: return 0 page_size = self.context.scheduler.page_size return (req.kv.kv_allocated_len + page_size - 1) // page_size diff --git a/test/registered/unit/disaggregation/test_specv2_kvcache_offloading.py b/test/registered/unit/disaggregation/test_specv2_kvcache_offloading.py index 806f85681..67522f853 100644 --- a/test/registered/unit/disaggregation/test_specv2_kvcache_offloading.py +++ b/test/registered/unit/disaggregation/test_specv2_kvcache_offloading.py @@ -8,7 +8,6 @@ Requires: torch, sglang (run in an environment with sglang installed) """ import unittest -from types import SimpleNamespace from unittest.mock import MagicMock import torch @@ -18,6 +17,7 @@ from sglang.srt.disaggregation.decode_kvcache_offload_manager import ( ) from sglang.srt.disaggregation.kv_events import OffloadedState from sglang.srt.managers.cache_controller import HiCacheAck +from sglang.srt.managers.schedule_batch import ReqKvInfo from sglang.test.ci.ci_register import register_cpu_ci register_cpu_ci(est_time=8, suite="base-a-test-cpu") @@ -35,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 = SimpleNamespace(kv_allocated_len=kv_allocated_len) + req.kv = ReqKvInfo(kv_allocated_len=kv_allocated_len) req.prefix_indices = list(range(prefix_indices_len)) req.effective_kv_committed_len = lambda: req.kv_committed_len return req diff --git a/test/registered/unit/managers/test_kv_page_invariants.py b/test/registered/unit/managers/test_kv_page_invariants.py index 62ba150ca..73ece119f 100644 --- a/test/registered/unit/managers/test_kv_page_invariants.py +++ b/test/registered/unit/managers/test_kv_page_invariants.py @@ -49,6 +49,7 @@ class _FakeReq: self.req_pool_idx = rpi self.kv_committed_len = committed self.kv = SimpleNamespace(kv_allocated_len=allocated, swa_evicted_seqlen=0) + self.is_holding_kv = True class _FakeSlot: diff --git a/test/registered/unit/mem_cache/test_dllm_fdfo_kv_reuse.py b/test/registered/unit/mem_cache/test_dllm_fdfo_kv_reuse.py index 79222d6e9..abe5f34d4 100644 --- a/test/registered/unit/mem_cache/test_dllm_fdfo_kv_reuse.py +++ b/test/registered/unit/mem_cache/test_dllm_fdfo_kv_reuse.py @@ -65,10 +65,8 @@ def _make_req(rid, prefix, block_size, *, req_pool_idx=None, reuse=False): dllm_incomplete_ids=array("q", range(block_size)) if reuse else array("q"), inflight_middle_chunks=1 if req_pool_idx is not None else 0, kv_committed_len=len(prefix) if req_pool_idx is not None else 0, - kv=( - SimpleNamespace(kv_allocated_len=len(prefix) + block_size) - if req_pool_idx is not None - else None + kv=SimpleNamespace( + kv_allocated_len=len(prefix) + block_size if req_pool_idx is not None else 0 ), ) diff --git a/test/registered/unit/mem_cache/test_hisparse_allocator.py b/test/registered/unit/mem_cache/test_hisparse_allocator.py index 4dac3419f..d1994f883 100644 --- a/test/registered/unit/mem_cache/test_hisparse_allocator.py +++ b/test/registered/unit/mem_cache/test_hisparse_allocator.py @@ -6,6 +6,7 @@ from unittest.mock import MagicMock import numpy as np import torch +from sglang.srt.managers.schedule_batch import ReqKvInfo from sglang.srt.mem_cache.allocator.hisparse import ( DeepSeekV4HiSparseTokenToKVPoolAllocator, ) @@ -91,7 +92,7 @@ class TestDeepSeekV4HiSparseAllocator(CustomTestCase): rid="req-0", origin_input_ids=list(range(fill_len)), output_ids=[], - kv=None, + kv=ReqKvInfo(), ) def set_extend_range(start, end): 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 7ab78ffe0..bcc4df1c8 100644 --- a/test/registered/unit/mem_cache/test_swa_eviction_boundary.py +++ b/test/registered/unit/mem_cache/test_swa_eviction_boundary.py @@ -104,6 +104,7 @@ def _make_req(req_pool_idx, token_ids, cache_protected_len, tree): """Mock Req with fields needed by _evict_swa and cache_finished_req.""" req = SimpleNamespace( req_pool_idx=req_pool_idx, + is_holding_kv=True, origin_input_ids=token_ids, output_ids=[], cache_protected_len=cache_protected_len, diff --git a/test/registered/unit/mem_cache/test_unified_radix_cache_bench.py b/test/registered/unit/mem_cache/test_unified_radix_cache_bench.py index 40584284d..78f5f6092 100644 --- a/test/registered/unit/mem_cache/test_unified_radix_cache_bench.py +++ b/test/registered/unit/mem_cache/test_unified_radix_cache_bench.py @@ -240,7 +240,7 @@ def create_bench_cache( _rid = [0] def make_req(): - from sglang.srt.managers.schedule_batch import Req, ReqKvInfo + from sglang.srt.managers.schedule_batch import Req from sglang.srt.sampling.sampling_params import SamplingParams req = Req( @@ -251,8 +251,6 @@ def create_bench_cache( ) _rid[0] += 1 req_to_token_pool.alloc([req]) - # fabricated reqs bypass alloc_for_extend, the normal creator of req.kv - req.kv = ReqKvInfo(kv_allocated_len=0, swa_evicted_seqlen=0) return req return tree, allocator, req_to_token_pool, make_req 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 fec8124de..9d3b4fd6c 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 @@ -24,7 +24,7 @@ from sglang.srt.disaggregation.kv_events import ( StorageMedium, ) from sglang.srt.environ import envs -from sglang.srt.managers.schedule_batch import Req, ReqKvInfo +from sglang.srt.managers.schedule_batch import Req from sglang.srt.mem_cache.allocator import TokenToKVPoolAllocator from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator from sglang.srt.mem_cache.base_prefix_cache import ( @@ -1030,7 +1030,6 @@ class UnifiedRadixCacheSuite: ) self._rid += 1 req_to_token_pool.alloc([req]) - req.kv = ReqKvInfo(kv_allocated_len=0, swa_evicted_seqlen=0) return req def _apply_match_to_req(self, req, match): @@ -1265,7 +1264,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 = ReqKvInfo(kv_allocated_len=kv_len, swa_evicted_seqlen=0) + req.kv.kv_allocated_len = kv_len req.last_node = cache.root_node_handle() req.cache_protected_len = 0 req.swa_uuid_for_lock = None @@ -2200,7 +2199,6 @@ class UnifiedRadixCacheSuite: req.cache_protected_len = 0 req.swa_uuid_for_lock = None req.extra_key = None - req.kv = ReqKvInfo(kv_allocated_len=0, swa_evicted_seqlen=0) swa_avail_before = allocator.swa_attn_allocator.available_size() @@ -2290,7 +2288,6 @@ class UnifiedRadixCacheSuite: req.cache_protected_len = 0 req.swa_uuid_for_lock = None req.extra_key = None - req.kv = ReqKvInfo(kv_allocated_len=0, swa_evicted_seqlen=0) with envs.SGLANG_OPT_UNIFIED_CACHE_FREE_OUT_OF_WINDOW_SLOTS.override(True): cache.cache_unfinished_req(req) @@ -6603,7 +6600,7 @@ class TestUnifiedRadixCacheInt8MambaCheckpoint(CustomTestCase): req_to_token_pool.alloc([req]) req.output_ids = array("q") req.kv_committed_len = len(tokens) - req.kv = ReqKvInfo(kv_allocated_len=len(tokens), swa_evicted_seqlen=0) + req.kv.kv_allocated_len = len(tokens) req.cache_protected_len = 0 req.swa_uuid_for_lock = None req.extra_key = None @@ -7934,7 +7931,6 @@ class TestSWAWindowUnderBigramKey(CustomTestCase): req.cache_protected_len = 0 req.swa_uuid_for_lock = None req.extra_key = None - req.kv = ReqKvInfo(kv_allocated_len=0, swa_evicted_seqlen=0) with envs.SGLANG_OPT_UNIFIED_CACHE_FREE_OUT_OF_WINDOW_SLOTS.override(True): cache.cache_unfinished_req(req)