Mini-Sglang:KVCache源码管理2

# engine/engine.py：注意 +1 是给 CUDA graph 的 dummy 请求留的行self.page_table = torch.zeros((config.max_running_req + 1, aligned_max_seq_len),dtype=torch.int32, device=self.device)

def allocate_paged(self, reqs: List[Req]) -> None:    needed_pages = 0 ## 计算总共需要的页面数目    allocation_info: List[Tuple[int, int, int]] = []    for req in reqs:        first_page = div_ceil(req.cached_len, self.page_size) ## 已缓存的前缀,按页管理        last_page = div_ceil(req.device_len, self.page_size) ## 请求总长折算成的页面数量        ##  div_ceil()向上取整，不足一页的算作一页        if last_page > first_page: ## 需要prefiil的页面数             needed_pages += last_page - first_page             allocation_info.append((req.table_idx, first_page, last_page))    if needed_pages > 0:        allocated = self._page_to_token(self._allocate(needed_pages))        _write_page_table(self.page_table, allocated, allocation_info, self.page_size)

 def _allocate(self, needed_pages: int) -> torch.Tensor:     if needed_pages > (free_pages := len(self.free_slots)): ## 判断是否需要驱逐页面        evicted = self.prefix_cache.evict((needed_pages - free_pages) * self.page_size) ## 驱逐页面        self.free_slots = torch.cat([self.free_slots, evicted[:: self.page_size]]) ## ::按pages_size采样进行驱逐        assert len(self.free_slots) >= needed_pages, "Eviction did not free enough space."     allocated = self.free_slots[:needed_pages] ## 分配需要的页面     self.free_slots = self.free_slots[needed_pages:] ## 总是从头部进行切片     return allocated

self.free_slots = torch.arange(num_pages, dtype=torch.int32, device=device) * page_size

def _page_to_token(self, pages: torch.Tensor) -> torch.Tensor:        if self.page_size == 1:            return pages        # [X * page_size] -> [X * page_size, ..., X * page_size + page_size - 1]        offsets = torch.arange(self.page_size, device=self.device, dtype=torch.int32) ## 生成页内偏移量        return (pages.unsqueeze(1) + offsets).flatten() ## 生成页偏移量的token索引        ## unsqueeze(1)将pages的形状从[X]变为[X, 1]，然后加上offsets得到[X, page_size]，最后flatten()将其展平为一维数组，得到每个页面对应的所有token索引。

ex: pages = [4,8],offsets = [0,1,2,3]    pages.unsqueeze(1) = [[4],[8]]    pages.unsqueeze(1) + offsets = [[4,5,6,7],[8,9,10,11]]    allocated = flatten() = [4,5,6,7,8,9,10,11]

def _write_page_table(    page_table: torch.Tensor,    allocated: torch.Tensor,    allocation_info: List[Tuple[int, int, int]],    page_size: int,) -> None:    needed_tokens = len(allocated)    table_idx_host = torch.empty(needed_tokens, dtype=torch.int64, pin_memory=True)     ## host 侧索引张量，专门存放每个待写入物理槽位在 page_table 中的行坐标。    positions_host = torch.empty(needed_tokens, dtype=torch.int64, pin_memory=True)     ##  host 侧索引张量，专门存放每个待写入物理槽位在 page_table 中的列坐标（token 逻辑位置）。    offset = 0    for table_idx, first_page, last_page in allocation_info: ## table_Idx:哪个请求         first_pos, last_pos = first_page * page_size, last_page * page_size        length = last_pos - first_pos ## 需要新填充的长度        table_idx_host[offset : offset + length].fill_(table_idx)        torch.arange(first_pos, last_pos, out=positions_host[offset : offset + length])        # first_pos=12, last_pos=20        # → positions_host[offset:offset+8] = [12,13,14,15,16,17,18,19]        offset += length    assert offset == needed_tokens, "Mismatch in allocated tokens and filled tokens."    table_idxs = table_idx_host.to(page_table.device, non_blocking=True)    offsets = positions_host.to(page_table.device, non_blocking=True)    page_table[table_idxs, offsets] = allocated     ## 逻辑 token 位置到物理 KV Cache 槽位的映射表    # 行：请求（或请求槽位）。每一行对应一个正在运行的请求。    # 列：该请求内的 token 逻辑位置（第几个 token）    # 值：该逻辑位置对应的物理 KV 槽位索引，也就是 K/V 实际存储在 KV Pool 中的地址    ## position_host:是token在token序列中的位置    ## allocated：是KVCache的物理页号

def __init__(self, num_pages: int, page_size: int, page_table: torch.Tensor, type: str):  # The `_free_slots` follows a page-aligned manner. For example, if page_size = 2,  # the `_free_slots` may look like [0, 2, 4, 6, ...], and each slot represents a page.  device = page_table.device  self.free_slots = torch.arange(num_pages, dtype=torch.int32, device=device) * page_size ##存储空闲页面的页偏移  self.prefix_cache = create_prefix_cache(device=device, type=type)  self.device = device  self.num_pages = num_pages  self.page_table = page_table  self.page_size = page_size

def match_req(self, req: PendingReq) -> MatchResult:  input_len = req.input_len  assert input_len > 0, "Input length must be greater than 0."  return self.prefix_cache.match_prefix(req.input_ids[: input_len - 1]) ##需要去掉最后一个token

 def cache_req(self, req: Req, *, finished: bool) -> None:        # ==================================== valid cache region ====================================        # [0, req.cached_len)                       This part is valid for attention kernel read/write.        # [0, old_handle.cached_len)                This part is in the prefix cache before prefill.        # [old_handle.cached_len, req.cached_len)   This part is allocated by cache manager for this request.        # ================================== allocated cache region ==================================        # [old_handle.cached_len, cached_len)       This part was not in the prefix cache when prefill,        #                                           but later cached by other requests.        #                                           We must free them to avoid memory leak.        # [cached_len, new_handle.cached_len)       This part is newly inserted into the prefix cache.        # [new_handle.cached_len, req.cached_len)   This part is tailing part that can not inserted into the prefix cache.        #                                           We should free it if the request has finished.        # [0, old_handle.cached_len)                prefill 前就在前缀缓存里（已命中）        # [old_handle.cached_len, cached_len)       prefill 时新分配，但之后被别的请求缓存了        # [cached_len, new_handle.cached_len)       本次新插入前缀缓存的部分        # [new_handle.cached_len, req.cached_len)   尾部，无法插入前缀缓存（未页对齐）        insert_ids = req.input_ids[: req.cached_len]        page_indices = self.page_table[req.table_idx, : req.cached_len]        old_handle = req.cache_handle        cached_len, new_handle = self.prefix_cache.insert_prefix(insert_ids, page_indices) ## 把经过prefill和前向计算后已经计算出KV的token缓存进树        # unlock until all operations on handle is done        self.unlock(old_handle)        # this part is already in the prefix cache, free it        self._free(page_indices[old_handle.cached_len : cached_len])        if finished:  # this tail part should be freed            self._free(page_indices[new_handle.cached_len :])        else:  # keep the tail part, update the handle            req.cache_handle = new_handle            self.lock(new_handle)

if finished:  # this tail part should be freed  self._free(page_indices[new_handle.cached_len :])

def allocate(self) -> int:   return self._free_slots.pop()

def free(self, slot: int) -> None:  self._free_slots.append(slot)