vLLM block_manager_v1.py 源码精读

class AllocStatus(enum.Enum):
    """Result for BlockSpaceManager.can_allocate

    1. Ok: seq_group can be allocated now.
    2. Later: seq_group cannot be allocated.
      The capacity of allocator is larger than seq_group required.
    3. Never: seq_group can never be allocated.
      The seq_group is too large to allocated in GPU.
    """
    OK = enum.auto()
    LATER = enum.auto()
    NEVER = enum.auto()


class BlockSpaceManager(ABC):

    @abstractmethod
    def can_allocate(self, seq_group: SequenceGroup) -> AllocStatus: ...

    @abstractmethod
    def allocate(self, seq_group: SequenceGroup) -> None: ...

    @abstractmethod
    def can_append_slots(self, seq_group: SequenceGroup,
                         num_lookahead_slots: int) -> bool: ...

    @abstractmethod
    def append_slots(self, seq: Sequence,
                     num_lookahead_slots: int) -> Dict[int, List[int]]: ...

    @abstractmethod
    def fork(self, parent_seq: Sequence, child_seq: Sequence) -> None: ...

    @abstractmethod
    def swap_in(self, seq_group: SequenceGroup,
                num_lookahead_slots: int) -> Dict[int, int]: ...

    @abstractmethod
    def swap_out(self, seq_group: SequenceGroup) -> Dict[int, int]: ...

    @abstractmethod
    def free(self, seq: Sequence) -> None: ...

class LogicalTokenBlock:
    """A block that stores a contiguous chunk of tokens from left to right.

    Logical blocks are used to represent the states of the corresponding
    physical blocks in the KV cache.
    """

    def __init__(
        self,
        block_number: int,
        block_size: int,
    ) -> None:
        self.block_number = block_number
        self.block_size = block_size

        self.token_ids = [_BLANK_TOKEN_ID] * block_size
        self.num_tokens = 0

    def is_empty(self) -> bool:
        return self.num_tokens == 0

    def get_num_empty_slots(self) -> int:
        return self.block_size - self.num_tokens

    def is_full(self) -> bool:
        return self.num_tokens == self.block_size

    def append_tokens(self, token_ids: List[int]) -> None:
        assert len(token_ids) <= self.get_num_empty_slots()
        curr_idx = self.num_tokens
        self.token_ids[curr_idx:curr_idx + len(token_ids)] = token_ids
        self.num_tokens += len(token_ids)

class PhysicalTokenBlock:
    """Represents the state of a block in the KV cache."""

    def __init__(
        self,
        device: Device,
        block_number: int,
        block_size: int,
        block_hash: int,
        num_hashed_tokens: int,
    ) -> None:
        self.device = device           # GPU 或 CPU
        self.block_number = block_number  # 在显存池中的槽位编号
        self.block_size = block_size   # 每块容纳的 token 数
        self.block_hash = block_hash   # 内容哈希（Prefix Caching 用）
        self.num_hashed_tokens = num_hashed_tokens  # 参与哈希的前缀长度

        self.ref_count = 0             # 引用计数（CoW 核心）
        self.last_accessed = DEFAULT_LAST_ACCESSED_TIME  # LRU 时间戳
        self.computed = False          # 是否已计算过 KV（Prefix Cache 命中标记）

# Mapping: logical block number -> physical block.
BlockTable = List[PhysicalTokenBlock]

class UncachedBlockAllocator(BlockAllocatorBase):

    def __init__(
        self,
        device: Device,
        block_size: int,
        num_blocks: int,
    ) -> None:
        self.device = device
        self.block_size = block_size
        self.num_blocks = num_blocks

        # Initialize the free blocks.
        self.free_blocks: BlockTable = []
        for i in range(num_blocks):
            block = PhysicalTokenBlock(device=device,
                                       block_number=i,
                                       block_size=block_size,
                                       block_hash=-1,
                                       num_hashed_tokens=0)
            self.free_blocks.append(block)

    def allocate(self,
                 block_hash: Optional[int] = None,
                 num_hashed_tokens: int = 0) -> PhysicalTokenBlock:
        if not self.free_blocks:
            raise ValueError("Out of memory! No free blocks are available.")
        block = self.free_blocks.pop()
        block.ref_count = 1
        return block

    def free(self, block: PhysicalTokenBlock) -> None:
        if block.ref_count == 0:
            raise ValueError(f"Double free! {block} is already freed.")
        block.ref_count -= 1
        if block.ref_count == 0:
            self.free_blocks.append(block)

class CachedBlockAllocator(BlockAllocatorBase):

    def __init__(self, device, block_size, num_blocks,
                 eviction_policy=EvictionPolicy.LRU) -> None:
        self.device = device
        self.block_size = block_size
        self.num_blocks = num_blocks
        self.current_num_blocks = 0
        self.cached_blocks: Dict[int, PhysicalTokenBlock] = {}
        self.evictor: Evictor = make_evictor(eviction_policy)
        self.default_hash_ctr = count()

    def allocate_block(self, block_hash, num_hashed_tokens):
        if self.current_num_blocks == self.num_blocks:
            # 显存已满：从 evictor 淘汰一块
            block = self.evictor.evict()
            block.block_hash = block_hash
            block.num_hashed_tokens = num_hashed_tokens
            return block
        # 显存有余：分配新的物理块槽位
        block = PhysicalTokenBlock(device=self.device,
                                   block_number=self.current_num_blocks,
                                   block_size=self.block_size,
                                   block_hash=block_hash,
                                   num_hashed_tokens=num_hashed_tokens)
        self.current_num_blocks += 1
        return block

    def allocate(self, block_hash=None, num_hashed_tokens=0):
        if block_hash is None:
            block_hash = next(self.default_hash_ctr)
        # 命中 evictor（已释放但未驱逐的缓存块）
        if block_hash in self.evictor:
            block = self.evictor.remove(block_hash)
            assert block.ref_count == 0
            self.cached_blocks[block_hash] = block
            block.ref_count += 1
            return block
        # 命中 cached_blocks（正在使用的块）→ 共享引用
        if block_hash not in self.cached_blocks:
            self.cached_blocks[block_hash] = self.allocate_block(
                block_hash, num_hashed_tokens)
        block = self.cached_blocks[block_hash]
        block.ref_count += 1
        return block

    def free(self, block: PhysicalTokenBlock) -> None:
        if block.ref_count == 0:
            raise ValueError(f"Double free! {block} is already freed.")
        block.ref_count -= 1
        if block.ref_count == 0:
            # 引用归零：不立即释放，交给 evictor 管理（延迟淘汰）
            self.evictor.add(block)
            del self.cached_blocks[block.block_hash]

def can_allocate(self, seq_group: SequenceGroup) -> AllocStatus:
    # 取 waiting 列表中第一个（代表性）序列
    seq = seq_group.get_seqs(status=SequenceStatus.WAITING)[0]
    # 计算该序列需要多少个物理块
    num_required_blocks = math.ceil(
        seq.get_len() / self.block_size
    )
    # 加上 lookahead 块（投机解码预留）
    num_required_blocks += self.block_sliding_window or 0

    if self.block_sliding_window is not None:
        num_required_blocks = min(num_required_blocks,
                                  self.block_sliding_window)

    num_free_gpu_blocks = self.gpu_allocator.get_num_free_blocks()
    # 若可用块数低于水位线，拒绝调度
    if (self.num_total_gpu_blocks - num_required_blocks
            < self.watermark_blocks):
        return AllocStatus.NEVER
    if num_free_gpu_blocks - num_required_blocks >= self.watermark_blocks:
        return AllocStatus.OK
    else:
        return AllocStatus.LATER

def allocate(self, seq_group: SequenceGroup) -> None:
    # 取代表性序列（所有序列共享相同 prompt，块布局相同）
    seq = seq_group.get_seqs(status=SequenceStatus.WAITING)[0]

    # 为该序列的每个逻辑块分配一个物理块
    block_table: BlockTable = []
    num_prompt_blocks = math.ceil(seq.get_len() / self.block_size)

    for logical_idx in range(num_prompt_blocks):
        # 计算该块的哈希（用于 Prefix Caching）
        block_hash = ... # 基于 token ids 内容
        block = self.gpu_allocator.allocate(
            block_hash, num_hashed_tokens=...)
        block_table.append(block)

    # 将 block_table 注册给序列组中的每个序列
    for seq in seq_group.get_seqs(status=SequenceStatus.WAITING):
        self.block_tables[seq.seq_id] = block_table.copy()

    # 对于并行采样（num_seqs > 1），子序列共享同一批物理块
    # 通过增加 ref_count 实现共享（CoW 的准备）
    for block in block_table:
        block.ref_count += len(seq_group.get_seqs()) - 1

def can_append_slots(self,
                     seq_group: SequenceGroup,
                     num_lookahead_slots: int = 0) -> bool:
    """Simple heuristic: a sequence can append if there are enough
    free blocks for all sequences in the group that need a new block.
    """
    num_free_gpu_blocks = self.gpu_allocator.get_num_free_blocks()

    # 统计当前 running 中需要新块的序列数
    num_seqs = seq_group.num_seqs(status=SequenceStatus.RUNNING)
    # 每个序列最多需要 1 个新块（当前最后一块已满时）
    # 加上 lookahead 预留（投机解码）
    num_required_blocks = num_seqs * (1 + num_lookahead_slots)

    return num_free_gpu_blocks >= num_required_blocks

def append_slots(
    self,
    seq: Sequence,
    num_lookahead_slots: int = 0,
) -> Dict[int, List[int]]:
    """Allocate a physical slot for the new token(s).

    Returns a dict of {src_block: [dst_block]} for Copy-on-Write.
    """
    block_table = self.block_tables[seq.seq_id]

    # 可能需要升级（promote）最后一块：
    # 若最后一块已满且是只读的前缀缓存块，需要 CoW
    block_table, maybe_promote_block = \
        self._maybe_promote_last_block(seq, block_table)

    # 检查最后一块是否已满
    if self._is_last_block_full(seq):
        # 分配新的物理块
        new_block = self._allocate_last_physical_block(seq)
        block_table.append(new_block)

    # lookahead 预留：额外分配 num_lookahead_slots 个块
    for _ in range(num_lookahead_slots):
        new_block = self.gpu_allocator.allocate()
        block_table.append(new_block)

    # 返回 CoW 映射表（供 CUDA kernel 执行内存复制）
    return maybe_promote_block

def _is_last_block_full(self, seq: Sequence) -> bool:
    token_ids_len = seq.data.get_len()
    return token_ids_len > 0 and token_ids_len % self.block_size == 0

def _allocate_last_physical_block(self, seq: Sequence) -> PhysicalTokenBlock:
    # 计算新块对应的 token 哈希（Prefix Caching）
    block_table = self.block_tables[seq.seq_id]
    num_hashed_tokens = self.block_size * len(block_table)

    # 计算内容哈希（含前缀块哈希，形成链式哈希）
    new_block = self.gpu_allocator.allocate(
        block_hash=hash of token slice,
        num_hashed_tokens=num_hashed_tokens
    )
    return new_block

def _promote_last_block(
    self,
    seq: Sequence,
    last_block: PhysicalTokenBlock,
) -> Tuple[PhysicalTokenBlock, Dict[int, List[int]]]:
    """If the last block is shared, copy it to a new block.

    Returns:
      - new_block: 专属的新物理块
      - cow_map: {旧块号: [新块号]}，传给 CUDA kernel 执行内存复制
    """
    if last_block.ref_count == 1:
        # 独占块，无需 CoW
        return last_block, {}

    # 触发 CoW：分配新块
    new_block = self.gpu_allocator.allocate()
    # 释放对旧块的引用（ref_count-1）
    self.gpu_allocator.free(last_block)

    # 返回 CoW 映射：调度器将此映射传给 Worker 执行 GPU 内存复制
    return new_block, {last_block.block_number: [new_block.block_number]}

def _maybe_promote_last_block(
    self,
    seq: Sequence,
    block_table: BlockTable,
) -> Tuple[BlockTable, Dict[int, List[int]]]:
    if self._is_last_block_full(seq):
        # 最后一块已满，即将分配新块，不需要 promote
        return block_table, {}
    last_block = block_table[-1]
    new_last_block, cow_map = self._promote_last_block(seq, last_block)
    if cow_map:
        block_table[-1] = new_last_block
    return block_table, cow_map

def fork(self, parent_seq: Sequence, child_seq: Sequence) -> None:
    """Fork a sequence by sharing all its physical blocks.

    Used for Beam Search: child sequence shares parent's blocks.
    """
    # 子序列的 block_table 是父序列的浅拷贝（共享物理块对象）
    src_block_table = self.block_tables[parent_seq.seq_id]
    self.block_tables[child_seq.seq_id] = src_block_table.copy()
    # 每个物理块的引用计数加 1
    for block in src_block_table:
        block.ref_count += 1

def swap_out(self, seq_group: SequenceGroup) -> Dict[int, int]:
    """Swap out all sequences in the group from GPU to CPU.

    Returns:
      mapping: {gpu_block_number: cpu_block_number}
    """
    mapping: Dict[PhysicalTokenBlock, PhysicalTokenBlock] = {}
    for seq in seq_group.get_seqs(status=SequenceStatus.RUNNING):
        new_block_table: BlockTable = []
        block_table = self.block_tables[seq.seq_id]

        for gpu_block in block_table:
            if gpu_block in mapping:
                # 该 GPU 块已被映射（共享块，只换出一次）
                cpu_block = mapping[gpu_block]
                cpu_block.ref_count += 1
            else:
                # 在 CPU 上分配对应的块
                cpu_block = self.cpu_allocator.allocate(
                    gpu_block.block_hash, gpu_block.num_hashed_tokens)
                mapping[gpu_block] = cpu_block
            new_block_table.append(cpu_block)

            # 释放 GPU 块
            self.gpu_allocator.free(gpu_block)

        self.block_tables[seq.seq_id] = new_block_table

    # 返回 {gpu_block_number: cpu_block_number} 供 Worker 执行 cudaMemcpy
    return {gpu_block.block_number: cpu_block.block_number
            for gpu_block, cpu_block in mapping.items()}

def swap_in(self, seq_group: SequenceGroup,
            num_lookahead_slots: int = 0) -> Dict[int, int]:
    """Swap in all sequences in the group from CPU to GPU.

    Returns:
      mapping: {cpu_block_number: gpu_block_number}
    """
    mapping: Dict[PhysicalTokenBlock, PhysicalTokenBlock] = {}
    for seq in seq_group.get_seqs(status=SequenceStatus.SWAPPED):
        new_block_table: BlockTable = []
        block_table = self.block_tables[seq.seq_id]

        for cpu_block in block_table:
            if cpu_block in mapping:
                gpu_block = mapping[cpu_block]
                gpu_block.ref_count += 1
            else:
                gpu_block = self.gpu_allocator.allocate(
                    cpu_block.block_hash, cpu_block.num_hashed_tokens)
                mapping[cpu_block] = gpu_block
            new_block_table.append(gpu_block)
            self.cpu_allocator.free(cpu_block)

        self.block_tables[seq.seq_id] = new_block_table

    return {cpu_block.block_number: gpu_block.block_number
            for cpu_block, gpu_block in mapping.items()}

def can_swap_in(self, seq_group: SequenceGroup,
                num_lookahead_slots: int = 0) -> AllocStatus:
    # 计算换入所需的 GPU 块数
    blocks = self._get_physical_blocks(seq_group)
    num_swapped_seqs = seq_group.num_seqs(status=SequenceStatus.SWAPPED)
    num_free_blocks = self.gpu_allocator.get_num_free_blocks()
    # 加上 lookahead 预留
    num_required = len(blocks) + num_swapped_seqs * num_lookahead_slots
    if self.num_total_gpu_blocks - num_required < self.watermark_blocks:
        return AllocStatus.NEVER
    if num_free_blocks - num_required >= self.watermark_blocks:
        return AllocStatus.OK
    return AllocStatus.LATER

def can_swap_out(self, seq_group: SequenceGroup) -> bool:
    # 检查 CPU 是否有足够空闲块
    blocks = self._get_physical_blocks(seq_group)
    return len(blocks) <= self.cpu_allocator.get_num_free_blocks()

class LRUEvictor(Evictor):
    """Evicts in a least-recently-used order using the last_accessed timestamp
    that's recorded in the PhysicalTokenBlock. If there are multiple blocks with
    the same last_accessed time, then the one with the largest num_hashed_tokens
    will be evicted.
    """

    def __init__(self):
        # OrderedDict 保持插入顺序，块按 free 时间排列
        self.free_table: OrderedDict[int, PhysicalTokenBlock] = OrderedDict()

    def evict(self) -> PhysicalTokenBlock:
        if len(self.free_table) == 0:
            raise ValueError("No usable cache memory left")

        evicted_block = next(iter(self.free_table.values()))
        # 遍历找到最旧（last_accessed 最小）且前缀最长的块
        for _, block in self.free_table.items():
            if evicted_block.last_accessed < block.last_accessed:
                break  # 时间戳更大，跳过
            if evicted_block.num_hashed_tokens < block.num_hashed_tokens:
                evicted_block = block  # 同时间戳，优先淘汰更长前缀块

        self.free_table.pop(evicted_block.block_hash)
        evicted_block.computed = False
        return evicted_block

    def add(self, block: PhysicalTokenBlock):
        # 块被 free() 时调用，加入待淘汰队列
        self.free_table[block.block_hash] = block

    def remove(self, block_hash: int) -> PhysicalTokenBlock:
        # 块被重新激活时调用（从 free 状态复活）
        if block_hash not in self.free_table:
            raise ValueError("Attempting to remove block not in evictor")
        block = self.free_table[block_hash]
        self.free_table.pop(block_hash)
        return block