chore: import upstream snapshot with attribution

python/tokenspeed/runtime/cache/allocator.py

@@ -0,0 +1,224 @@
+# Copyright (c) 2026 LightSeek Foundation
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+"""Allocators for KV-cache page metadata and slot management."""
+
+import torch
+
+from tokenspeed.runtime.utils import get_colorful_logger
+
+logger = get_colorful_logger(__name__)
+
+
+class KVAllocator:
+    """Operate on token slots and block-table metadata only.
+
+    Physical KV storage is managed separately so the same metadata operations
+    can work with different memory backends.
+    """
+
+    def __init__(
+        self,
+        size: int,
+        device: str,
+        max_batch_size: int,
+        max_context_len: int,
+        page_size: int,
+    ):
+        self.free_slots = None
+        self.token_slot_refs = None
+        self.size = size
+        self.is_not_in_free_group = True
+        self.free_group = []
+        self.page_size = page_size
+        self.device = device
+        self.last_slot = torch.ones(max_batch_size, dtype=torch.int32) * (page_size - 1)
+        self.num_pages = torch.zeros(max_batch_size, dtype=torch.int32)
+        self.max_context_len = max_context_len
+        self.max_page_num = (max_context_len + page_size - 1) // page_size
+        self.max_batch_size = max_batch_size
+        self.req_to_page = None
+        self.req_to_page_cpu = None
+        self.clear()
+
+    def available_size(self) -> int:
+        return len(self.free_slots)
+
+    def alloc(self, req_pool_index: int, need_size: int, alloced_len: int):
+        page_offset = alloced_len % self.page_size
+        page_num = (alloced_len + self.page_size - 1) // self.page_size
+        last_page_remain = page_num * self.page_size - alloced_len
+        last_page_id = self.req_to_page_cpu[req_pool_index, page_num - 1].item()
+        # if last_page_remain is zero, kv_loc is Tensor([])
+        kv_loc = (
+            last_page_id * self.page_size
+            + page_offset
+            + torch.arange(0, min(last_page_remain, need_size), dtype=torch.int32)
+        )
+        if last_page_remain >= need_size:
+            return kv_loc.to(self.device, non_blocking=True)
+
+        remain_size = need_size - last_page_remain
+        need_new_page_num = (remain_size + self.page_size - 1) // self.page_size
+        if need_new_page_num > len(self.free_slots):
+            # do not change self.seq_lens
+            return None
+
+        # Check if we have enough space in req_to_page tensor
+        if page_num + need_new_page_num > self.max_page_num:
+            logger.warning(
+                "Requested page range [%s:%s] exceeds max_page_num %s. alloced_len=%s, need_size=%s, page_num=%s",
+                page_num,
+                page_num + need_new_page_num,
+                self.max_page_num,
+                alloced_len,
+                need_size,
+                page_num,
+            )
+            # Do not change self.seq_lens
+            return None
+
+        new_pages = self.free_slots[:need_new_page_num]
+        self.free_slots = self.free_slots[need_new_page_num:]
+        # update req_to_page
+        self.req_to_page[req_pool_index, page_num : page_num + need_new_page_num] = (
+            new_pages.to(self.device)
+        )
+        self.req_to_page_cpu[
+            req_pool_index, page_num : page_num + need_new_page_num
+        ] = new_pages
+        # construct kv_loc
+        kv_loc1 = new_pages.unsqueeze(1) * self.page_size
+        offsets = torch.arange(0, self.page_size, dtype=torch.int32)
+        kv_loc1 = kv_loc1 + offsets
+        kv_loc1 = kv_loc1.flatten()[:remain_size]
+        final_kv_loc = torch.concat([kv_loc, kv_loc1]).to(
+            self.device, non_blocking=True
+        )
+        return final_kv_loc
+
+    def free_extra_pages_not_cached(
+        self, req_pool_index: int, real_seq_len: int, alloced_len: int
+    ):
+        full_page_num = real_seq_len // self.page_size
+        alloced_page_num = (alloced_len + self.page_size - 1) // self.page_size
+        page_num_to_free = alloced_page_num - full_page_num
+        if page_num_to_free == 0:
+            return
+        page_ids_to_free = self.req_to_page[
+            req_pool_index, full_page_num : full_page_num + page_num_to_free
+        ]
+        self.need_to_free.append(page_ids_to_free)
+
+    def free_req_cache(self, req_pool_index: int, alloced_len: int):
+        """Release all pages of the request when prefix cache is not used."""
+        alloced_page_num = (alloced_len + self.page_size - 1) // self.page_size
+        if alloced_page_num == 0:
+            return
+        page_ids_to_free = self.req_to_page[req_pool_index, :alloced_page_num]
+        self.need_to_free.append(page_ids_to_free)
+
+    def free_with_diff(self, new_prefix_page_ids, old_page_ids):
+        # New KV pages come from the prefix tree and are already cached, so only
+        # release the pages that differ from the old allocation.
+        if len(new_prefix_page_ids) != len(old_page_ids):
+            raise ValueError(
+                "[free with diff] new_prefix_page_ids and old_page_ids "
+                "should have the same length"
+            )
+        diff = new_prefix_page_ids != old_page_ids
+        if torch.any(diff):
+            logger.debug(
+                "[DebugTrace] free_with_diff free page=%s", old_page_ids[diff].tolist()
+            )
+            self.need_to_free.append(old_page_ids[diff])
+        else:
+            logger.debug(
+                "[DebugTrace] free_with_diff: no pages to free, all pages are cached"
+            )
+        return diff
+
+    def append_to_later_free(self, page_ids: torch.Tensor) -> None:
+        self.need_to_free.append(page_ids)
+
+    def free(self, req_pool_index: int, indices=None) -> None:
+        if self.is_not_in_free_group:
+            num_pages = self.num_pages[req_pool_index]
+            pages = self.req_to_page[req_pool_index, :num_pages].cpu()
+            free_slots = [self.free_slots]
+            for i in range(num_pages):
+                page_index = pages[i]
+                free_slots.append(
+                    torch.arange(
+                        page_index * self.page_size,
+                        (page_index + 1) * self.page_size,
+                        dtype=torch.int32,
+                    )
+                )
+            self.free_slots = torch.concat(free_slots)
+            self.num_pages[req_pool_index] = 0
+            self.last_slot[req_pool_index] = self.page_size - 1
+        else:
+            self.free_group.append(req_pool_index)
+
+    def free_group_end(self) -> None:
+        self.is_not_in_free_group = True
+        if self.need_to_free:
+            pages_need_to_free = torch.concat(self.need_to_free)
+            logger.debug(
+                "[DebugTrace] free_group_end pages_need_to_free=%s",
+                pages_need_to_free.tolist(),
+            )
+            token_level_offsets = torch.arange(self.page_size, device=self.device)
+            slots_to_free = (
+                pages_need_to_free[:, None] * self.page_size + token_level_offsets
+            ).flatten()
+            writted_positions = slots_to_free[self.token_slot_refs[slots_to_free] >= 1]
+            self.token_slot_refs[writted_positions] += -1
+            self.free_slots = torch.concat([self.free_slots, pages_need_to_free.cpu()])
+            self.need_to_free = []
+
+    def clear(self) -> None:
+        # Page 0 is used for padding
+        self.free_slots = torch.arange(
+            1, self.size // self.page_size, dtype=torch.int32
+        )
+        if self.token_slot_refs is None:
+            self.token_slot_refs = torch.zeros(
+                self.size, dtype=torch.int32, device=self.device
+            )
+        else:
+            self.token_slot_refs.zero_()
+        if self.req_to_page is None:
+            self.req_to_page = torch.zeros(
+                (self.max_batch_size, self.max_page_num),
+                dtype=torch.int32,
+                device=self.device,
+            )
+            self.req_to_page_cpu = torch.zeros(
+                (self.max_batch_size, self.max_page_num),
+                dtype=torch.int32,
+                pin_memory=True,
+            )
+        else:
+            self.req_to_page.zero_()
+            self.req_to_page_cpu.zero_()
+        self.free_group = []
+        self.need_to_free = []