From 261860e17b2727f2de5f5f44166008ffb4d56fc8 Mon Sep 17 00:00:00 2001 From: hw-csong <226476376+hw-csong@users.noreply.github.com> Date: Thu, 8 Jan 2026 12:00:18 +0800 Subject: [PATCH] [NPU][Bugfix] move free_page logics to cpu (#16608) --- .../srt/hardware_backend/npu/allocator_npu.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/python/sglang/srt/hardware_backend/npu/allocator_npu.py b/python/sglang/srt/hardware_backend/npu/allocator_npu.py index dd7221287..e863de62b 100644 --- a/python/sglang/srt/hardware_backend/npu/allocator_npu.py +++ b/python/sglang/srt/hardware_backend/npu/allocator_npu.py @@ -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)