diff --git a/python/sglang/srt/pic/picache.py b/python/sglang/srt/pic/picache.py index 2be8689558..c5803a731e 100644 --- a/python/sglang/srt/pic/picache.py +++ b/python/sglang/srt/pic/picache.py @@ -2,10 +2,10 @@ See qianyou/2026-05-28-pic-sglang-design.md §5.4. """ + from __future__ import annotations import logging -import os import time from dataclasses import dataclass from typing import Any, Dict, List, Optional @@ -38,10 +38,11 @@ class SegmentEntry: that the existing `EvictionStrategy` classes (LRUStrategy/FIFOStrategy/...) duck-type on us without modification. """ + seg_hash: bytes - full_kv_slots: torch.Tensor # int64, len == segment length - mamba_state_slot: int # MambaPool slot id - token_ids: torch.Tensor # for hash-collision fallback compare + full_kv_slots: torch.Tensor # int64, len == segment length + mamba_state_slot: int # MambaPool slot id + token_ids: torch.Tensor # for hash-collision fallback compare lock_ref: int = 0 last_access_time: float = 0.0 creation_time: float = 0.0 @@ -93,6 +94,17 @@ def __init__( def supports_mamba(self) -> bool: return True + def _free_mamba_slots(self, slots: List[int]) -> None: + if not slots: + return + self.mamba_allocator.free( + torch.tensor( + slots, + dtype=torch.int64, + device=self.mamba_allocator.device, + ) + ) + def is_chunk_cache(self) -> bool: return False @@ -131,7 +143,10 @@ def mamba_evictable_size(self) -> int: return sum(1 for e in self._entries.values() if e.lock_ref == 0) def mamba_protected_size(self) -> int: - return sum(1 for e in self._entries.values() if e.lock_ref > 0) + self._inflight_mamba_slots + return ( + sum(1 for e in self._entries.values() if e.lock_ref > 0) + + self._inflight_mamba_slots + ) def full_protected_size(self) -> int: return self.protected_size() + self._inflight_full_tokens @@ -142,7 +157,9 @@ def add_inflight(self, num_tokens: int, num_mamba_slots: int) -> None: def remove_inflight(self, num_tokens: int, num_mamba_slots: int) -> None: self._inflight_full_tokens = max(0, self._inflight_full_tokens - num_tokens) - self._inflight_mamba_slots = max(0, self._inflight_mamba_slots - num_mamba_slots) + self._inflight_mamba_slots = max( + 0, self._inflight_mamba_slots - num_mamba_slots + ) def supports_swa(self) -> bool: return False @@ -190,7 +207,9 @@ def pretty_print(self) -> None: self.protected_size(), ) - def _match_segment(self, seg_hash: bytes, token_ids: torch.Tensor) -> Optional[SegmentEntry]: + def _match_segment( + self, seg_hash: bytes, token_ids: torch.Tensor + ) -> Optional[SegmentEntry]: entry = self._entries.get(seg_hash) if entry is None: return None @@ -234,12 +253,16 @@ def _insert_segment( self._entries[seg_hash] = entry return entry - def inject_received_segment(self, seg_hash, token_ids, full_kv_slots, mamba_state_slot): + def inject_received_segment( + self, seg_hash, token_ids, full_kv_slots, mamba_state_slot + ): existing = self._entries.get(seg_hash) if existing is not None: existing.lock_ref += 1 return existing - entry = self._insert_segment(seg_hash, token_ids, full_kv_slots, mamba_state_slot) + entry = self._insert_segment( + seg_hash, token_ids, full_kv_slots, mamba_state_slot + ) entry.lock_ref += 1 return entry @@ -248,8 +271,11 @@ def match_prefix(self, params: MatchPrefixParams) -> MatchResult: Lock acquired later in schedule_policy._req_inc_lock_ref. """ req = params.req - device = (self.token_to_kv_pool_allocator.device - if hasattr(self.token_to_kv_pool_allocator, "device") else "cpu") + device = ( + self.token_to_kv_pool_allocator.device + if hasattr(self.token_to_kv_pool_allocator, "device") + else "cpu" + ) empty_indices = torch.empty((0,), dtype=torch.int64, device=device) empty = MatchResult( device_indices=empty_indices, @@ -258,7 +284,12 @@ def match_prefix(self, params: MatchPrefixParams) -> MatchResult: best_match_node=None, pic_segment_entries=None, ) - if self.disable or req is None or getattr(req, "pic_segments", None) is None or len(req.pic_segments) == 0: + if ( + self.disable + or req is None + or getattr(req, "pic_segments", None) is None + or len(req.pic_segments) == 0 + ): return empty segments = req.pic_segments @@ -268,6 +299,7 @@ def match_prefix(self, params: MatchPrefixParams) -> MatchResult: is_recompute = self.policy.recompute if is_recompute: from sglang.srt.pic import SEAM_SINK_DEFAULT, resolve_seam_sink_tokens + seam_sink = SEAM_SINK_DEFAULT req.pic_hit_seam_positions = {} @@ -277,7 +309,9 @@ def match_prefix(self, params: MatchPrefixParams) -> MatchResult: min_tokens = getattr(get_global_server_args(), "pic_segment_min_tokens", -1) if min_tokens > 0 and end - start < min_tokens: continue - seg_token_ids = torch.tensor(req.origin_input_ids[start:end], dtype=torch.int64) + seg_token_ids = torch.tensor( + req.origin_input_ids[start:end], dtype=torch.int64 + ) h_bytes = segment_hash(seg_token_ids) entry = self._match_segment(h_bytes, seg_token_ids) if entry is not None: @@ -310,7 +344,11 @@ def match_prefix(self, params: MatchPrefixParams) -> MatchResult: hit_tokens = sum(e.full_kv_slots.numel() for e in per_seg if e is not None) logger.info( "PICache.match_prefix: %d segments, %d hit (%d tokens), %d miss, cache_size=%d entries", - len(segments), num_hit, hit_tokens, num_miss, len(self._entries), + len(segments), + num_hit, + hit_tokens, + num_miss, + len(self._entries), ) return MatchResult( @@ -343,8 +381,9 @@ def cache_unfinished_req(self, req, **kwargs) -> None: inflight_mamba = 0 req.pic_cache_owned_miss_segments = set() req.pic_freed_miss_segments = set() + duplicate_mamba_slots = [] min_tokens = getattr(get_global_server_args(), "pic_segment_min_tokens", -1) - for (start, end) in miss_segments: + for start, end in miss_segments: # ponytail: the last segment (Q) is normally not cached (recomputed by # decode). A scatter single-seg sub-request has exactly one segment that # IS the "last" one but MUST be cached + pushed to combine — don't skip. @@ -367,19 +406,19 @@ def cache_unfinished_req(self, req, **kwargs) -> None: slot_tuple = miss_slots[(start, end)] if is_rope: _private_slots, public_slots, mamba_slot = slot_tuple - assert public_slots is not None, ( - "local miss segment must have a public slot in rope PIC modes" - ) + assert ( + public_slots is not None + ), "local miss segment must have a public slot in rope PIC modes" kv_slots = public_slots else: kv_slots, mamba_slot = slot_tuple - assert mamba_slot is not None, ( - "cacheable miss segment must have a mamba slot" - ) + assert ( + mamba_slot is not None + ), "cacheable miss segment must have a mamba slot" existing = self._match_segment(seg_hash, seg_ids) if existing is not None: self.token_to_kv_pool_allocator.free(kv_slots) - self.mamba_allocator.free(torch.tensor([mamba_slot], dtype=torch.int64, device=kv_slots.device)) + duplicate_mamba_slots.append(mamba_slot) req.pic_segment_entries[seg_hash] = existing req.pic_freed_miss_segments.add((start, end)) else: @@ -395,6 +434,8 @@ def cache_unfinished_req(self, req, **kwargs) -> None: inflight_tokens += kv_slots.numel() inflight_mamba += 1 + self._free_mamba_slots(duplicate_mamba_slots) + # Transition mode: inflight -> lock_ref handoff (for the slots now owned # by entries — i.e. public-slot inflight for rope, kv-slot inflight for # plain transition). Private-slot inflight (rope only) is removed by @@ -405,12 +446,14 @@ def cache_unfinished_req(self, req, **kwargs) -> None: if inserted > 0: logger.info( "PICache.cache_unfinished_req: inserted %d new segments, total cache=%d entries", - inserted, len(self._entries), + inserted, + len(self._entries), ) def cache_finished_req(self, req, is_insert: bool = True, **kwargs) -> None: miss_slots = getattr(req, "pic_miss_segment_slots", None) if miss_slots: + mamba_slots_to_free = [] is_transition = self.policy.compose is PICCompose.TRANSITION is_rope = self.policy.rope if is_insert: @@ -423,7 +466,8 @@ def cache_finished_req(self, req, is_insert: bool = True, **kwargs) -> None: for priv, pub, _mamba in miss_slots.values() ) inflight_mamba = sum( - 1 for _priv, _pub, mamba in miss_slots.values() + 1 + for _priv, _pub, mamba in miss_slots.values() if mamba is not None ) else: @@ -431,8 +475,7 @@ def cache_finished_req(self, req, is_insert: bool = True, **kwargs) -> None: slots.numel() for slots, _mamba in miss_slots.values() ) inflight_mamba = sum( - 1 for _slots, mamba in miss_slots.values() - if mamba is not None + 1 for _slots, mamba in miss_slots.values() if mamba is not None ) self.remove_inflight(inflight_tokens, inflight_mamba) # Free last segment's slots (never cached by design). @@ -445,16 +488,11 @@ def cache_finished_req(self, req, is_insert: bool = True, **kwargs) -> None: last_priv, last_pub, last_mamba = miss_slots[last_seg] assert last_pub is None self.token_to_kv_pool_allocator.free(last_priv) - last_dev = last_priv.device else: last_kv, last_mamba = miss_slots[last_seg] self.token_to_kv_pool_allocator.free(last_kv) - last_dev = last_kv.device if last_mamba is not None: - self.mamba_allocator.free( - torch.tensor([last_mamba], dtype=torch.int64, - device=last_dev) - ) + mamba_slots_to_free.append(last_mamba) if is_transition and not is_rope: self.remove_inflight( last_kv.numel(), @@ -485,18 +523,9 @@ def cache_finished_req(self, req, is_insert: bool = True, **kwargs) -> None: if skipped_slots: combined = torch.cat(skipped_slots) self.token_to_kv_pool_allocator.free(combined) - if skipped_mamba_slots: - self.mamba_allocator.free( - torch.tensor( - skipped_mamba_slots, - dtype=torch.int64, - device=combined.device, - ) - ) + mamba_slots_to_free.extend(skipped_mamba_slots) if is_transition: - self.remove_inflight( - skipped_tokens, len(skipped_mamba_slots) - ) + self.remove_inflight(skipped_tokens, len(skipped_mamba_slots)) # transition_rope: free ALL miss private slots (local miss private was # never registered to an entry; global private was just freed above). @@ -522,12 +551,14 @@ def cache_finished_req(self, req, is_insert: bool = True, **kwargs) -> None: gp = miss_slots[global_seg][0] self.remove_inflight(gp.numel(), 1) # +1 for global mamba + self._free_mamba_slots(mamba_slots_to_free) + # transition_rope[_recompute]: free private hit slots (rerotated copies, not cached). if self.policy.rope: hit_private = getattr(req, "pic_rope_hit_private_slots", None) if hit_private: all_private = [] - for (private_slots, _entry_slots) in hit_private.values(): + for private_slots, _entry_slots in hit_private.values(): all_private.append(private_slots) if all_private: combined = torch.cat(all_private) @@ -559,9 +590,8 @@ def cache_finished_req(self, req, is_insert: bool = True, **kwargs) -> None: # HybridReqToTokenPool.alloc for the decode-time recurrent state. # release_kv_cache skips this branch for caches with # supports_mamba()=True (PIC), so we must free it ourselves. - if ( - getattr(req, "mamba_pool_idx", None) is not None - and hasattr(self.req_to_token_pool, "free_mamba_cache") + if getattr(req, "mamba_pool_idx", None) is not None and hasattr( + self.req_to_token_pool, "free_mamba_cache" ): self.req_to_token_pool.free_mamba_cache(req) @@ -577,14 +607,17 @@ def evict(self, params: EvictParams) -> EvictResult: candidates.sort(key=lambda e: self.eviction_strategy.get_priority(e)) evicted_tokens = 0 evicted_mamba = 0 + evicted_entries = [] for e in candidates: if evicted_tokens >= need_tokens and evicted_mamba >= need_mamba: break self.token_to_kv_pool_allocator.free(e.full_kv_slots) - self.mamba_allocator.free(torch.tensor([e.mamba_state_slot], dtype=torch.int64, - device=e.full_kv_slots.device)) evicted_tokens += int(e.full_kv_slots.numel()) evicted_mamba += 1 + evicted_entries.append(e) + + self._free_mamba_slots([entry.mamba_state_slot for entry in evicted_entries]) + for e in evicted_entries: del self._entries[e.seg_hash] self.update_eviction_metrics(evicted_tokens, start_time) return EvictResult( diff --git a/test/registered/unit/pic/test_picache_mamba_free.py b/test/registered/unit/pic/test_picache_mamba_free.py new file mode 100644 index 0000000000..bf5503b8ad --- /dev/null +++ b/test/registered/unit/pic/test_picache_mamba_free.py @@ -0,0 +1,185 @@ +from types import SimpleNamespace + +import pytest +import torch + +from sglang.srt.mem_cache.allocator.mamba import MambaSlotAllocator +from sglang.srt.mem_cache.base_prefix_cache import EvictParams +from sglang.srt.pic.picache import PICache +from sglang.srt.pic.segmenter import segment_hash +from sglang.srt.server_args import ( + get_global_server_args, + set_global_server_args_for_scheduler, +) +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + + +@pytest.fixture +def set_pic_min_tokens(): + try: + previous_server_args = get_global_server_args() + except ValueError: + previous_server_args = None + + def set_value(value): + set_global_server_args_for_scheduler( + SimpleNamespace(pic_segment_min_tokens=value) + ) + + yield set_value + set_global_server_args_for_scheduler(previous_server_args) + + +class RecordingAllocator: + page_size = 1 + device = "cpu" + + def __init__(self): + self.freed = [] + + def available_size(self): + return 4096 + + def free(self, slots): + self.freed.append(slots.clone()) + + +class RecordingReqPool: + def __init__(self): + self.mamba_allocator = RecordingAllocator() + + +def make_picache(): + req_pool = RecordingReqPool() + kv_allocator = RecordingAllocator() + cache = PICache( + req_to_token_pool=req_pool, + token_to_kv_pool_allocator=kv_allocator, + mamba_pool=SimpleNamespace(), + page_size=1, + disable=False, + ) + return cache, kv_allocator, req_pool.mamba_allocator + + +def insert_entry(cache, token_ids, kv_slots, mamba_slot, last_access_time): + token_ids = torch.tensor(token_ids, dtype=torch.int64) + entry = cache._insert_segment( + segment_hash(token_ids), + token_ids, + torch.tensor(kv_slots, dtype=torch.int64), + mamba_slot, + ) + entry.last_access_time = last_access_time + return entry + + +def test_duplicate_cleanup_returns_all_mamba_slots_together(set_pic_min_tokens): + set_pic_min_tokens(-1) + cache, kv_allocator, mamba_allocator = make_picache() + first = insert_entry(cache, [1, 2], [1, 2], 1, 1) + second = insert_entry(cache, [3, 4], [3, 4], 2, 2) + req = SimpleNamespace( + origin_input_ids=[1, 2, 3, 4, 5], + pic_segments=[(0, 2), (2, 4), (4, 5)], + pic_miss_segments=[(0, 2), (2, 4)], + pic_miss_segment_slots={ + (0, 2): (torch.tensor([101, 102]), 11), + (2, 4): (torch.tensor([103, 104]), 12), + }, + pic_segment_entries={}, + ) + + cache.cache_unfinished_req(req) + + assert [slots.tolist() for slots in kv_allocator.freed] == [ + [101, 102], + [103, 104], + ] + assert [slots.tolist() for slots in mamba_allocator.freed] == [[11, 12]] + assert req.pic_segment_entries == { + first.seg_hash: first, + second.seg_hash: second, + } + assert req.pic_freed_miss_segments == {(0, 2), (2, 4)} + assert len(cache._entries) == 2 + + +def test_finished_request_returns_last_and_skipped_mamba_slots_together( + set_pic_min_tokens, +): + set_pic_min_tokens(3) + cache, kv_allocator, mamba_allocator = make_picache() + req = SimpleNamespace( + pic_segments=[(0, 2), (2, 4), (4, 5)], + pic_miss_segment_slots={ + (0, 2): (torch.tensor([101, 102]), 11), + (2, 4): (torch.tensor([103, 104]), 12), + (4, 5): (torch.tensor([105]), 13), + }, + pic_segment_entries={}, + req_pool_idx=None, + mamba_pool_idx=None, + ) + + cache.cache_finished_req(req, is_insert=False) + + assert [slots.tolist() for slots in kv_allocator.freed] == [ + [105], + [101, 102, 103, 104], + ] + assert [slots.tolist() for slots in mamba_allocator.freed] == [[13, 11, 12]] + assert cache._entries == {} + + +def test_evict_returns_selected_mamba_slots_in_lru_order_together(): + cache, kv_allocator, mamba_allocator = make_picache() + first = insert_entry(cache, [1], [101, 102], 11, 1) + locked = insert_entry(cache, [2], [201], 21, 0) + third = insert_entry(cache, [3], [301, 302, 303], 31, 2) + retained = insert_entry(cache, [4], [401], 41, 3) + locked.lock_ref = 1 + + result = cache.evict(EvictParams(num_tokens=5, mamba_num=2)) + + assert result.num_tokens_evicted == 5 + assert result.mamba_num_evicted == 2 + assert [slots.tolist() for slots in kv_allocator.freed] == [ + [101, 102], + [301, 302, 303], + ] + assert [slots.tolist() for slots in mamba_allocator.freed] == [[11, 31]] + assert list(cache._entries) == [locked.seg_hash, retained.seg_hash] + assert first.seg_hash not in cache._entries + assert third.seg_hash not in cache._entries + + +def test_evict_restores_real_mamba_allocator_capacity_and_order(): + req_pool = RecordingReqPool() + req_pool.mamba_allocator = MambaSlotAllocator(4, "cpu") + allocated = req_pool.mamba_allocator.alloc(4) + cache = PICache( + req_to_token_pool=req_pool, + token_to_kv_pool_allocator=RecordingAllocator(), + mamba_pool=SimpleNamespace(), + page_size=1, + disable=False, + ) + for offset, slot in enumerate(allocated.tolist()): + insert_entry( + cache, + [offset], + [100 + offset], + slot, + last_access_time=offset, + ) + + result = cache.evict(EvictParams(num_tokens=3, mamba_num=3)) + + assert result.num_tokens_evicted == 3 + assert result.mamba_num_evicted == 3 + assert req_pool.mamba_allocator.available_size() == 3 + assert req_pool.mamba_allocator.free_slots.tolist() == [1, 2, 3] + assert [entry.mamba_state_slot for entry in cache._entries.values()] == [4]