Clean up ascend allocator (#11152)

2025-10-02 20:34:26 +08:00
parent 083629c235
commit 7d00479950
3 changed files with 29 additions and 20 deletions
@@ -493,7 +493,11 @@ class PagedTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
        if self.debug_mode:
            assert len(torch.unique(out_indices)) == len(out_indices)

-        num_new_pages = get_num_new_pages(prefix_lens_cpu, seq_lens_cpu, self.page_size)
+        num_new_pages = get_num_new_pages(
+            seq_lens=seq_lens_cpu,
+            page_size=self.page_size,
+            prefix_lens=prefix_lens_cpu,
+        )
        if num_new_pages > len(self.free_pages):
            return None

@@ -529,7 +533,9 @@ class PagedTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
            assert len(torch.unique(out_indices)) == len(out_indices)

        num_new_pages = get_num_new_pages(
-            seq_lens_cpu - 1, seq_lens_cpu, self.page_size, decode=True
+            seq_lens=seq_lens_cpu,
+            page_size=self.page_size,
+            decode=True,
        )
        if num_new_pages > len(self.free_pages):
            return None