lightseekorg--tokenspeed/python/tokenspeed/runtime/cache/allocator.py

# 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 = []