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

[NPU][Bugfix] move free_page logics to cpu (#16608)

Browse Source
This commit is contained in:
hw-csong
2026-01-08 12:00:18 +08:00
committed by GitHub
parent 154740bd4d
commit 261860e17b
1 changed files with 18 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
python/sglang/srt/hardware_backend/npu/allocator_npu.py
View File

@@ -182,3 +182,21 @@ class NPUPagedTokenToKVPoolAllocator(PagedTokenToKVPoolAllocator):
self.free_pages = self.free_pages[num_new_pages:]
return out_indices.int()
def free(self, free_index: torch.Tensor):
if free_index.numel() == 0:
return
if self.is_not_in_free_group:
device = free_index.device
free_page_indices = torch.unique(free_index.cpu() // self.page_size)
free_page_indices = free_page_indices.to(device)
if self.need_sort:
self.release_pages = torch.cat((free_page_indices, self.release_pages))
else:
self.free_pages = torch.cat((free_page_indices, self.free_pages))
else:
self.free_group.append(free_index)
if self.debug_mode:
assert len(torch.unique(self.free_pages)) == len(self.free_pages)