chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

SWA Prefix Cache (#7367)

Browse Source 
Co-authored-by: Ying Sheng <sqy1415@gmail.com>
This commit is contained in:
Hanming Lu
2025-07-13 12:31:07 -07:00
committed by GitHub
parent 0c55cbcfc5
commit 9379da77de
16 changed files with 1742 additions and 158 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
python/sglang/srt/mem_cache/allocator.py
View File

@@ -57,11 +57,6 @@ class BaseTokenToKVPoolAllocator(abc.ABC):
def debug_print(self) -> str:
return ""
def log_usage(self, evictable_size: int = 0):
num_used = self.size - (self.available_size() + evictable_size)
msg = f"#token: {num_used}, token usage: {num_used / self.size:.2f}, "
return msg, num_used
def available_size(self):
return len(self.free_pages) * self.page_size
@@ -190,7 +185,7 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
self._kvcache.full_to_swa_index_mapping = self.full_to_swa_index_mapping
def available_size(self):
return min(self.full_available_size(), self.swa_available_size())
raise NotImplementedError()
def full_available_size(self):
return self.full_attn_allocator.available_size()
@@ -214,16 +209,6 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
)
return msg
def log_usage(self, swa_evictable_size: int = 0, full_evictable_size: int = 0):
used_full = self.size_full - (self.full_available_size() + full_evictable_size)
used_swa = self.size_swa - (self.swa_available_size() + swa_evictable_size)
msg = (
f"#token: full={used_full}, swa={used_swa}, "
f"token usage: full={used_full / self.size_full:.2f}, "
f"swa={used_swa / self.size_swa:.2f}, "
)
return msg, used_full
def get_kvcache(self):
return self._kvcache