vllm 优化之 PagedAttention 源码解读

def __init__(
    self,
    scheduler_config: SchedulerConfig,
    cache_config: CacheConfig,
    lora_config: Optional[LoRAConfig],
    pipeline_parallel_size: int = 1,
    output_proc_callback: Optional[Callable] = None,
) -> None:
    self.scheduler_config = scheduler_config
    self.cache_config = cache_config
    # Note for LoRA scheduling: the current policy is extremely
    # simple and NOT fair. It can lead to starvation of some
    # LoRAs. This should be improved in the future.
    self.lora_config = lora_config

    version = "selfattn"
    if (self.scheduler_config.task == "embedding"
            or self.cache_config.is_attention_free):
        version = "placeholder"

    BlockSpaceManagerImpl = BlockSpaceManager.get_block_space_manager_class(
        version)

    # Create the block space manager.
    self.block_manager = BlockSpaceManagerImpl(
        block_size=self.cache_config.block_size,
        num_gpu_blocks=cache_config.num_gpu_blocks,
        num_cpu_blocks=cache_config.num_cpu_blocks,
        sliding_window=self.cache_config.sliding_window,
        enable_caching=self.cache_config.enable_prefix_caching)

class BlockTable:
    """管理特定序列的内存块的类。
    
    BlockTable 将一系列令牌映射到一组块中，每个块代表序列的一部分连续内存分配。这些块由 DeviceAwareBlockAllocator 管理，负责块的分配和释放。
    
    参数：
        block_size (int): 每个内存块可以存储的最大令牌数量。
        block_allocator (DeviceAwareBlockAllocator): 用于管理物理块内存的分配器。
        _blocks (Optional[List[Block]], optional): 可选的现有块列表，用于初始化 BlockTable。如果未提供，则创建一个空的 BlockTable。
        max_block_sliding_window (Optional[int], optional): 每个序列需要保留的最大块数。如果为 None，则保留所有块（例如，当不使用滑动窗口时）。至少应满足模型的滑动窗口大小。
    
    属性：
        _block_size (int): 每个内存块可以存储的最大令牌数量。
        _allocator (DeviceAwareBlockAllocator): 用于管理物理块内存的分配器。
        _blocks (Optional[List[Block]]): 由此 BlockTable 管理的逻辑块列表。
        _num_full_slots (int): 当前存储在块中的令牌数量。
    """

    def __init__(
        self,
        block_size: int,
        block_allocator: DeviceAwareBlockAllocator,
        _blocks: Optional[List[Block]] = None,
        max_block_sliding_window: Optional[int] = None,
    ):
        self._block_size = block_size  # 设置每个块的大小
        self._allocator = block_allocator  # 设置内存分配器
        if _blocks is None:
            _blocks = []
        self._blocks: BlockList = BlockList(_blocks)  # 初始化块列表

        self._max_block_sliding_window = max_block_sliding_window  # 设置滑动窗口的最大块数
        self._num_full_slots = self._get_num_token_ids()  # 获取当前存储的令牌数量


    def allocate(self,
                 token_ids: List[int],
                 device: Device = Device.GPU) -> None:
        """为给定的令牌序列分配内存块。
        
        此方法分配所需数量的块以存储给定的令牌序列。
        
        参数：
            token_ids (List[int]): 要存储的令牌 ID 序列。
            device (Device, optional): 要分配块的设备。默认为 GPU。
        """
        assert not self._is_allocated  # 确保尚未分配块
        assert token_ids  # 确保有令牌需要分配
        blocks = self._allocate_blocks_for_token_ids(prev_block=None,
                                                     token_ids=token_ids,
                                                     device=device)  # 分配块
        self.update(blocks)  # 更新块表
        self._num_full_slots = len(token_ids)  # 更新存储的令牌数量

    def _allocate_blocks_for_token_ids(self, prev_block: Optional[Block],
                                        token_ids: List[int],
                                        device: Device) -> List[Block]:
        """为给定的令牌 ID 分配内存块。
        
        参数：
            prev_block (Optional[Block]): 前一个块。如果没有，则为 None。
            token_ids (List[int]): 要存储的令牌 ID 列表。
            device (Device): 要分配块的设备。
        
        返回：
            List[Block]: 分配的块列表。
        """
        blocks: List[Block] = []

        block_token_ids = []
        tail_token_ids = []
        for cur_token_ids in chunk_list(token_ids, self._block_size):
            if len(cur_token_ids) == self._block_size:
                block_token_ids.append(cur_token_ids)
            else:
                tail_token_ids.append(cur_token_ids)

        if block_token_ids:
            blocks.extend(
                self._allocator.allocate_immutable_blocks(
                    prev_block, block_token_ids=block_token_ids,
                    device=device))  # 分配不可变块
            prev_block = blocks[-1]

        if tail_token_ids:
            assert len(tail_token_ids) == 1  # 仅有一个尾块
            cur_token_ids = tail_token_ids[0]

            block = self._allocator.allocate_mutable_block(
                prev_block=prev_block, device=device)  # 分配可变块
            block.append_token_ids(cur_token_ids)  # 追加令牌 ID

            blocks.append(block)

        return blocks
    
    def update(self, blocks: List[Block]) -> None:
        """重置块表为新提供的块（包括其对应的块 ID）。
        
        参数：
            blocks (List[Block]): 新分配的块列表。
        """
        self._blocks.update(blocks)

def test_allocate_blocks_for_token_ids():
    # 初始化分配器和 BlockTable
    allocator = DeviceAwareBlockAllocator()
    block_size = 4
    block_table = BlockTable(block_size=block_size, block_allocator=allocator)

    # 定义测试令牌 ID
    token_ids = [1, 2, 3, 4, 5, 6, 7, 8, 9]

    # 分配块
    allocated_blocks = block_table._allocate_blocks_for_token_ids(
        prev_block=None,
        token_ids=token_ids,
        device=Device.GPU
    )

    # 更新块表
    block_table._blocks.update(allocated_blocks)
    block_table._num_full_slots = len(token_ids)

    # 打印分配结果
    print("分配的块数量:", len(allocated_blocks))
    for i, block in enumerate(allocated_blocks):
        print(f"块 {i}: Token IDs = {block.token_ids}, Mutable = {block.mutable}")

    # 预期：
    # - 前两个块完全填满（4个令牌）
    # - 最后一个块不完全填满（1个令牌，5）
    # 所以应有 3 个块
    assert len(allocated_blocks) == 3, "应分配 3 个块"
    assert allocated_blocks[0].token_ids == [1, 2, 3, 4], "第一个块令牌 ID 错误"
    assert allocated_blocks[1].token_ids == [5, 6, 7, 8], "第二个块令牌 ID 错误"
    assert allocated_blocks[2].token_ids == [9], "第三个块令牌 ID 错误"
    assert not allocated_blocks[0].mutable, "第一个块应为不可变块"
    assert not allocated_blocks[1].mutable, "第二个块应为不可变块"
    assert allocated_blocks[2].mutable, "第三个块应为可变块"

    print("测试通过！")

if __name__ == "__main__":
    test_allocate_blocks_for_token_ids()

class CpuGpuBlockAllocator(DeviceAwareBlockAllocator):
    """一个能够在 CPU 和 GPU 内存中分配块的块分配器。
    
    该类实现了 `DeviceAwareBlockAllocator` 接口，提供了在 CPU 和 GPU 设备上分配和管理内存块的功能。
    
    `CpuGpuBlockAllocator` 维护了独立的 CPU 和 GPU 内存块池，并允许在这些内存池之间进行分配、释放、分叉和交换操作。
    """
    
    @staticmethod
    def create(
        allocator_type: str,
        num_gpu_blocks: int,
        num_cpu_blocks: int,
        block_size: int,
    ) -> DeviceAwareBlockAllocator:
        """创建一个具有指定配置的 CpuGpuBlockAllocator 实例。
        
        这个静态方法根据提供的参数创建并返回一个 CpuGpuBlockAllocator 实例。它初始化了 CPU 和 GPU 块分配器，指定块的数量、块大小和分配器类型。
        
        参数：
            allocator_type (str): 用于 CPU 和 GPU 块的块分配器类型。目前支持的值为 "naive" 和 "prefix_caching"。
            num_gpu_blocks (int): 要为 GPU 内存分配的块数量。
            num_cpu_blocks (int): 要为 CPU 内存分配的块数量。
            block_size (int): 每个块的大小，以令牌数量表示。
        
        返回：
            DeviceAwareBlockAllocator: 一个具有指定配置的 CpuGpuBlockAllocator 实例。
        
        注意：
            - 块 ID 是连续分配的，GPU 块 ID 在前，CPU 块 ID 在后。
        """
        # 对于 HPU，块 ID 0 仅用于填充
        reserved_blocks = 1 if current_platform.is_hpu() else 0
        block_ids = list(
            range(reserved_blocks, num_gpu_blocks + num_cpu_blocks))
        num_gpu_blocks -= reserved_blocks
        gpu_block_ids = block_ids[:num_gpu_blocks]
        cpu_block_ids = block_ids[num_gpu_blocks:]
    
        # 根据 allocator_type 创建不同类型的块分配器
        if allocator_type == "naive":
            gpu_allocator: BlockAllocator = NaiveBlockAllocator(
                create_block=NaiveBlock,  # 创建不可变块的函数
                num_blocks=num_gpu_blocks,
                block_size=block_size,
                block_ids=gpu_block_ids,
            )
    
            cpu_allocator: BlockAllocator = NaiveBlockAllocator(
                create_block=NaiveBlock,  # 创建不可变块的函数
                num_blocks=num_cpu_blocks,
                block_size=block_size,
                block_ids=cpu_block_ids,
            )
        elif allocator_type == "prefix_caching":
            gpu_allocator = PrefixCachingBlockAllocator(
                num_blocks=num_gpu_blocks,
                block_size=block_size,
                block_ids=gpu_block_ids,
            )
    
            cpu_allocator = PrefixCachingBlockAllocator(
                num_blocks=num_cpu_blocks,
                block_size=block_size,
                block_ids=cpu_block_ids,
            )
        else:
            raise ValueError(f"未知的分配器类型 {allocator_type=}")
    
        return CpuGpuBlockAllocator(
            cpu_block_allocator=cpu_allocator,
            gpu_block_allocator=gpu_allocator,
        )
    
    def __init__(self, cpu_block_allocator: BlockAllocator,
                 gpu_block_allocator: BlockAllocator):
        """初始化 CpuGpuBlockAllocator 实例。
        
        参数：
            cpu_block_allocator (BlockAllocator): 用于管理 CPU 内存块的分配器。
            gpu_block_allocator (BlockAllocator): 用于管理 GPU 内存块的分配器。
        """
        # 确保 CPU 和 GPU 分配器的块 ID 没有交集
        assert not (
            cpu_block_allocator.all_block_ids
            & gpu_block_allocator.all_block_ids
        ), "CPU 和 GPU 块分配器的块 ID 不能有交集"
    
        # 将 CPU 和 GPU 分配器存储在字典中
        self._allocators = {
            Device.CPU: cpu_block_allocator,
            Device.GPU: gpu_block_allocator,
        }
    
        self._swap_mapping: Dict[int, int] = {}  # 记录交换操作的块 ID 映射
        self._null_block: Optional[Block] = None  # 用于存储空块
    
        # 记录每个块 ID 对应的分配器
        self._block_ids_to_allocator: Dict[int, BlockAllocator] = {}
        for _, allocator in self._allocators.items():
            for block_id in allocator.all_block_ids:
                self._block_ids_to_allocator[block_id] = allocator
    
    def allocate_mutable_block(self, prev_block: Optional[Block],
                               device: Device) -> Block:
        """在指定设备上分配一个新的可变块。
        
        参数：
            prev_block (Optional[Block]): 序列中的前一个块。用于前缀哈希。
            device (Device): 要分配新块的设备。
        
        返回：
            Block: 新分配的可变块。
        """
        return self._allocators[device].allocate_mutable_block(prev_block)
    
    def allocate_immutable_blocks(self, prev_block: Optional[Block],
                                  block_token_ids: List[List[int]],
                                  device: Device) -> List[Block]:
        """在指定设备上分配一组包含提供的块令牌 ID 的不可变块。
        
        参数：
            prev_block (Optional[Block]): 序列中的前一个块。用于前缀哈希。
            block_token_ids (List[int]): 要存储在新块中的块令牌 ID 列表。
            device (Device): 要分配新块的设备。
        
        返回：
            List[Block]: 新分配的包含提供的块令牌 ID 的不可变块列表。
        """
        return self._allocators[device].allocate_immutable_blocks(
            prev_block, block_token_ids)

class NaiveBlockAllocator(BlockAllocator):
    """一个简单的块分配器，不支持前缀缓存。
    
    该类实现了 `BlockAllocator` 接口，提供了基本的内存块分配和释放功能。
    
    参数：
        create_block (Block.Factory): 用于创建新块的工厂函数。当 NaiveBlockAllocator 被前缀缓存分配器组合使用时，必须能够创建前缀缓存块（但不应了解其余细节）。
        num_blocks (int): 要管理的块的总数量。
        block_size (int): 每个块的大小，以令牌数量表示。
        block_ids (Optional[Iterable[int]], optional): 可选的块 ID 可迭代对象。如果未提供，块 ID 将从 0 到 `num_blocks - 1` 顺序分配。
    """
    
    def __init__(
        self,
        create_block,
        num_blocks: int,
        block_size: int,
        block_ids: Optional[Iterable[int]] = None,
        block_pool: Optional[BlockPool] = None,
    ):
        if block_ids is None:
            block_ids = range(num_blocks)  # 如果未提供块 ID，则顺序分配
        
        self._free_block_indices: deque = deque(block_ids)  # 初始化自由块队列
        self._all_block_indices = frozenset(block_ids)  # 所有块 ID 的集合
        assert len(self._all_block_indices) == num_blocks, "块 ID 数量应与 num_blocks 相等"
    
        self._refcounter = RefCounter(
            all_block_indices=self._free_block_indices)
        self._block_size = block_size
    
        self._cow_tracker = CopyOnWriteTracker(
            refcounter=self._refcounter)
    
        if block_pool is None:
            extra_factor = 4
            # 预分配 "num_blocks * extra_factor" 个块对象。
            # "* extra_factor" 是为了允许分配比物理块更多的块对象
            self._block_pool = BlockPool(self._block_size, create_block, self,
                                         num_blocks * extra_factor)
        else:
            # 在这种情况下，块池由调用者提供，意味着可能需要在分配器之间共享块池
            self._block_pool = block_pool
    
    def allocate_immutable_blocks(
            self,
            prev_block: Optional[Block],
            block_token_ids: List[List[int]],
            device: Optional[str] = None) -> List[Block]:
        """分配一组新的不可变块，包含提供的块令牌 ID，并链接到前一个块。
        
        参数：
            prev_block (Optional[Block]): 序列中的前一个块。如果为 None，则分配的块为序列中的第一个块。
            block_token_ids (List[List[int]]): 要存储在新块中的块令牌 ID 列表。
            device (Optional[str], optional): 分配块的设备。对于 NaiveBlockAllocator，通常为 None。
        
        返回：
            List[Block]: 新分配的不可变块列表。
        """
        assert device is None, "NaiveBlockAllocator 不支持设备参数"
        num_blocks = len(block_token_ids)  # 需要分配的块数量
    
        block_ids = []
        for i in range(num_blocks):
            block_ids.append(self._allocate_block_id())  # 分配块 ID
    
        blocks = []
        for i in range(num_blocks):
            # 初始化块，设置前一个块、令牌 ID、块大小和物理块 ID
            prev_block = self._block_pool.init_block(
                prev_block=prev_block,
                token_ids=block_token_ids[i],
                block_size=self._block_size,
                physical_block_id=block_ids[i])
            blocks.append(prev_block)
    
        return blocks
    
    def allocate_mutable_block(self,
                               prev_block: Optional[Block],
                               device: Optional[str] = None) -> Block:
        """分配一个新的可变块，并链接到前一个块。
        
        参数：
            prev_block (Optional[Block]): 序列中的前一个块。如果为 None，则分配的块为序列中的第一个块。
            device (Optional[str], optional): 分配块的设备。对于 NaiveBlockAllocator，通常为 None。
        
        返回：
            Block: 新分配的可变块。
        """
        assert device is None, "NaiveBlockAllocator 不支持设备参数"
        block_id = self._allocate_block_id()  # 分配一个块 ID
        block = self._block_pool.init_block(prev_block=prev_block,
                                            token_ids=[],  # 初始化为空令牌 ID
                                            block_size=self._block_size,
                                            physical_block_id=block_id)
        block.mutable = True  # 设置为可变块
        return block

class BlockList:
    """This class is an optimization to allow fast-access to physical 
    block ids. It maintains a block id list that is updated with the 
    block list and this avoids the need to reconstruct the block id 
    list on every iteration of the block manager
    """

    def __init__(self, blocks: List[Block]):
        self._blocks: List[Block] = []
        self._block_ids: List[int] = []

        self.update(blocks)

    def _add_block_id(self, block_id: Optional[BlockId]) -> None:
        assert block_id is not None
        self._block_ids.append(block_id)

    def _update_block_id(self, block_index: int,
                         new_block_id: Optional[BlockId]) -> None:
        assert new_block_id is not None
        self._block_ids[block_index] = new_block_id
    #####省略代码######

block_tables: Optional[torch.Tensor]

@dataclass
class FlashAttentionMetadata:
    # NOTE(sang): Definition of context_len, query_len, and seq_len.
    # |---------- N-1 iteration --------|
    # |---------------- N iteration ---------------------|
    # |- tokenA -|......................|-- newTokens ---|
    # |---------- context_len ----------|
    # |-------------------- seq_len ---------------------|
    #                                   |-- query_len ---|

    num_actual_tokens: int  # Number of tokens excluding padding.
    max_query_len: int
    query_start_loc: torch.Tensor
    max_seq_len: int
    seq_start_loc: torch.Tensor
    block_table: torch.Tensor
    slot_mapping: torch.Tensor

# vllm/attention/ops/paged_attn.py
class PagedAttention:
    def get_kv_cache_shape(
        num_blocks: int,
        block_size: int,
        num_kv_heads: int,
        head_size: int,
    ) -> Tuple[int, ...]:
        return (2, num_blocks, block_size * num_kv_heads * head_size)

cache_block_size = self.get_cache_block_size_bytes()
# self.cache_config.cpu_kvcache_space_bytes 可通过环境变量 `VLLM_CPU_KVCACHE_SPACE` 定义
num_cpu_blocks = int(self.cache_config.cpu_kvcache_space_bytes // cache_block_size)

@staticmethod
def get_cache_block_size(
    block_size: int,
    cache_dtype: str,
    model_config: ModelConfig,
    parallel_config: ParallelConfig,
) -> int:
    head_size = model_config.get_head_size()
    num_heads = model_config.get_num_kv_heads(parallel_config)
    num_layers = model_config.get_num_layers(parallel_config)

    key_cache_block = block_size * num_heads * head_size
    value_cache_block = key_cache_block
    
    # 每层 layer 的 block size
    total = num_layers * (key_cache_block + value_cache_block)
    if cache_dtype == "auto":
        dtype = model_config.dtype
    else:
        dtype = STR_DTYPE_TO_TORCH_DTYPE[cache_dtype]
    dtype_size = torch.tensor([], dtype=dtype).element_size()
    return dtype_size * total

@torch.inference_mode()
def determine_num_available_blocks(model_config, gpu_memory_utilization = 0.9) -> Tuple[int, int]:
    """
    评估模型的峰值内存使用情况，以确定在不发生内存溢出的情况下可以分配的 KV（键值）缓存块的数量。

    该方法首先清理 CUDA 缓存，然后使用虚拟输入执行一次前向传播，以评估模型的内存使用情况。
    接着，计算在剩余可用内存下，最多可以分配的 GPU 和 CPU 缓存块数量。

    提示：
        可以通过调整 `gpu_memory_utilization` 参数来限制 GPU 内存的使用。
    """
    # 清理 CUDA 缓存，以确保获取准确的内存使用信息
    torch.cuda.empty_cache()

    # 使用虚拟输入执行一次前向传播，以评估模型的内存使用情况
    self.model_runner.profile_run()

    # 同步 CUDA 操作，确保内存信息准确
    torch.cuda.synchronize()
    
    # 获取当前 GPU 的空闲内存和总内存（单位：字节）
    free_memory_pre_profile, total_gpu_memory = torch.cuda.mem_get_info()
    # 计算模型加载后的峰值内存使用量
    # Get the peak memory allocation recorded by torch
    peak_memory = torch.cuda.memory_stats()["allocated_bytes.all.peak"]
    
    # 清理未使用的缓存，计算非Torch分配的内存
    torch.cuda.empty_cache()
    torch_allocated_bytes = torch.cuda.memory_stats()["allocated_bytes.all.current"]

    total_allocated_bytes = torch.cuda.mem_get_info()[1] - torch.cuda.mem_get_info()[0]
    non_torch_allocations = total_allocated_bytes - torch_allocated_bytes
    
    if non_torch_allocations > 0:
        peak_memory += non_torch_allocations

    available_kv_cache_memory = (
        total_gpu_memory * gpu_memory_utilization -
        peak_memory)
    
    # 计算每个缓存块的大小
    cache_block_size = _get_cache_block_size(model_config)
    # 计算在剩余可用内存下，最多可以分配的 GPU 缓存块数量
    num_gpu_blocks = int(
        (total_gpu_memory * gpu_memory_utilization -
         peak_memory) // cache_block_size
    )
    # 确保缓存块数量不为负数
    num_gpu_blocks = max(num_gpu_blocks, 0)

    logger.info(
            "Memory profiling results: total_gpu_memory=%.2fGiB \n"
            " initial_memory_usage=%.2fGiB peak_torch_memory=%.2fGiB \n"
            " memory_usage_post_profile=%.2fGib \n"
            " non_torch_memory=%.2fGiB kv_cache_size=%.2fGiB \n"
            " gpu_memory_utilization=%.2f", total_gpu_memory / (1024**3),
            (total_gpu_memory - free_memory_pre_profile) / (1024**3),
            (peak_memory - non_torch_allocations) / (1024**3),
            total_allocated_bytes / (1024**3),
            non_torch_allocations / (1024**3),
            available_kv_cache_memory / (1024**3),
            gpu_memory_utilization)

    # 进行垃圾回收，释放未使用的内存
    gc.collect()
    # 再次清理 CUDA 缓存
    torch.cuda.empty_cache()
    # 返回可分配的 GPU 和 CPU 缓存块数量（此处 CPU 块数量为 0）

    return num_gpu_blocks, 0