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

[HiCache] Clamp tombstoned SWA locs in UnifiedSWAKVPool translation (#35933)

Browse Source
This commit is contained in:
Zhiqiang Xie
2026-08-21 23:08:56 -07:00
committed by GitHub
parent 4cb5aebfe0
commit c35683fda0
1 changed files with 11 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/mem_cache/unified_memory_pool.py
+8 -2
View File
@@ -1424,12 +1424,18 @@ class UnifiedSWAKVPool(SWAKVPool):
"attach_allocators" "attach_allocators"
) )
ps = self._swa_allocator.page_size ps = self._swa_allocator.page_size
# Tombstone-safety clamp, matching MultiEndedAllocator.translate_kv_loc:
# a tombstoned v2p entry (-1) must not reach the caller as a negative
# loc. Clamp to 0 routes it to the reserved padding sink instead.
if ps == 1: if ps == 1:
return self._swa_allocator.virtual_to_physical[kv_indices].to(torch.int32) swa_locs = self._swa_allocator.virtual_to_physical[kv_indices]
else:
virt_pages = kv_indices // ps virt_pages = kv_indices // ps
offsets = kv_indices % ps offsets = kv_indices % ps
swa_phys_pages = self._swa_allocator.virtual_to_physical[virt_pages] swa_phys_pages = self._swa_allocator.virtual_to_physical[virt_pages]
return (swa_phys_pages * ps + offsets).to(torch.int32) # Tombstoned page: -1 * ps + offset lands in [-ps, -1].
swa_locs = swa_phys_pages * ps + offsets
return swa_locs.clamp(min=0).to(torch.int32)
def get_state_buf_infos(self): def get_state_buf_infos(self):
return self.swa_kv_pool.get_contiguous_buf_infos() return self.swa_kv_pool.get_contiguous_buf_infos()