[NPU][Bugfix] Fix OOB gather in decode KV allocation when free pool is tight (#35727)

2026-08-27 19:44:45 +08:00
parent 97ba99067d
commit b647ae82f5
1 changed files with 1 additions and 0 deletions
@@ -126,6 +126,7 @@ class NPUPagedTokenToKVPoolAllocator(PagedTokenToKVPoolAllocator):
        if num_new_pages == 0:
            out_indices = last_loc + 1
        else:
+            start_new_pages = start_new_pages.clamp(max=num_new_pages - 1)
            out_indices = (last_loc + 1) * (1 - need_new_pages) + self.free_pages[
                start_new_pages
            ] * self.page_size * need_new_pages