wehub-resource-sync
106 lines
3.6 KiB
Python
106 lines
3.6 KiB
Python
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
|
|
from typing import Iterable, Union
|
|
|
|
import torch
|
|
|
|
|
|
class BlockedAllocator:
|
|
"""
|
|
Allocator class for managing which blocks are free/used in the
|
|
blocked KV-cache. This is a simple allocator that uses a linked list
|
|
to keep track of which blocks are free/used. The cost of allocation/deallocation
|
|
is O(blocks), where blocks is the number of blocks to allocate/deallocate.
|
|
|
|
TODO(cmikeh2): Evaluate performance of this allocator and migrate
|
|
to C++ if necessary.
|
|
"""
|
|
# Number of blocks in the KV-cache(s).
|
|
_num_blocks: int
|
|
|
|
# Array of blocks, where each element is the next block in the linked list.
|
|
_blocks: torch.Tensor
|
|
|
|
# Index of the head of the linked list.
|
|
_head: int
|
|
|
|
# Number of free blocks in the KV-cache.
|
|
_free_blocks: int
|
|
|
|
def __init__(self, num_blocks: int) -> None:
|
|
"""
|
|
Initialize an allocator with `num_blocks` blocks. This requires at least
|
|
`num_blocks` * 4 bytes of host memory.
|
|
|
|
Parameters:
|
|
num_blocks (int): The number of blocks to allocate.
|
|
"""
|
|
|
|
if num_blocks < 1:
|
|
raise ValueError(f'Blocked KV-cache must have at least 1 block, provided {num_blocks}')
|
|
|
|
self._num_blocks = num_blocks
|
|
self._blocks = torch.arange(1, num_blocks + 1, dtype=torch.int32, device='cpu', pin_memory=True)
|
|
self._head = 0
|
|
self._free_blocks = num_blocks
|
|
|
|
def allocate(self, num_blocks: int) -> torch.Tensor:
|
|
"""
|
|
Allocate a list of blocks from the associated KV-caches. This will
|
|
return `num_blocks` blocks from the KV-cache if they are available,
|
|
or raise an exception if there are not enough free blocks.
|
|
|
|
Parameters:
|
|
num_blocks (int): The number of blocks to allocate.
|
|
|
|
Returns:
|
|
List[int]: The list of blocks allocated.
|
|
"""
|
|
if num_blocks > self._free_blocks:
|
|
raise ValueError(f'Not enough free blocks in the KV-cache to allocate {num_blocks} blocks')
|
|
|
|
allocated_blocks = torch.zeros(num_blocks, dtype=torch.int32)
|
|
for i in range(num_blocks):
|
|
allocated_blocks[i] = self._head
|
|
self._head = self._blocks[self._head].item()
|
|
self._blocks[allocated_blocks[i]] = -1 # Mark as used
|
|
self._free_blocks -= 1
|
|
|
|
return allocated_blocks
|
|
|
|
def free(self, blocks: Union[Iterable[int], int]) -> None:
|
|
"""
|
|
Return a list of blocks to the free pool. If a single invalid block is provided (i.e.,
|
|
one that is out of range of the allocator or is already free), then an exception is raised
|
|
and no blocks are freed.
|
|
|
|
Parameters:
|
|
blocks (Union[Iterable[int], int]): The list of blocks to free. If only one block
|
|
is to be freed, this can be alone as an integer.
|
|
"""
|
|
if isinstance(blocks, int):
|
|
blocks = [blocks]
|
|
|
|
for block in blocks:
|
|
# Parse all blocks for validity before mutating the list.
|
|
if block < 0 or block >= self._num_blocks:
|
|
raise ValueError(f'Invalid block {block} provided to free')
|
|
|
|
if self._blocks[block] != -1:
|
|
raise ValueError(f'Block {block} is already free')
|
|
|
|
for block in blocks:
|
|
self._blocks[block] = self._head
|
|
self._head = block
|
|
self._free_blocks += 1
|
|
|
|
@property
|
|
def free_blocks(self) -> int:
|
|
"""
|
|
Return the number of free blocks in the KV-cache.
|
|
"""
|
|
return self._free_blocks
|