vLLM Block Manager V2 与 Prefix Caching 源码精读

class Block(ABC):

    @abstractmethod
    def append_token_ids(self, token_ids: List[int]) -> None:
        pass

    @property
    @abstractmethod
    def block_id(self) -> Optional[int]:
        pass

    @property
    @abstractmethod
    def token_ids(self) -> List[int]:
        pass

    @property
    @abstractmethod
    def num_empty_slots(self) -> int:
        pass

    @property
    @abstractmethod
    def is_full(self) -> bool:
        pass

    @property
    @abstractmethod
    def prev_block(self) -> Optional["Block"]:
        pass

    @property
    @abstractmethod
    def computed(self) -> bool:
        raise NotImplementedError

    @property
    @abstractmethod
    def last_accessed(self) -> float:
        raise NotImplementedError

    @property
    @abstractmethod
    def content_hash(self) -> Optional[int]:
        """Return the content-based hash of the current block, or None if it is
        not yet defined or not supported. For the content-based hash to be
        defined, the current block must be full."""
        return None

class BlockTable:
    """A class to manage blocks for a specific sequence.

    The BlockTable maps a sequence of tokens to a list of blocks, where each
    block represents a contiguous memory allocation for a portion of the
    sequence.
    """

    def __init__(self, block_size, block_allocator, _blocks=None):
        self._block_size = block_size
        self._allocator = block_allocator
        self._blocks: List[Block] = _blocks or []
        self._num_full_slots = len(self._get_all_token_ids())

    @staticmethod
    def get_num_required_blocks(token_ids, block_size) -> int:
        return cdiv(len(token_ids), block_size)

    def allocate(self, token_ids, device=Device.GPU):
        """首次分配：为整个 token 序列分配所有需要的 block"""
        self._blocks = self._allocate_blocks_for_token_ids(
            prev_block=None, token_ids=token_ids, device=device)
        self._num_full_slots = len(token_ids)

    def append_token_ids(self, token_ids, num_lookahead_slots=0):
        """追加 token：用于 decode 阶段逐步生成"""
        self.ensure_num_empty_slots(len(token_ids) + num_lookahead_slots)
        blocks = self._blocks[self._num_full_slots // self._block_size:]
        token_blocks = self._chunk_token_blocks_for_append(token_ids)
        for block, token_block in zip(blocks, token_blocks):
            block.append_token_ids(token_block)
        self._num_full_slots += len(token_ids)

def ensure_num_empty_slots(self, num_empty_slots: int) -> None:
    if self._num_empty_slots >= num_empty_slots:
        return

    slots_to_allocate = num_empty_slots - self._num_empty_slots
    blocks_to_allocate = cdiv(slots_to_allocate, self._block_size)

    for _ in range(blocks_to_allocate):
        self._blocks.append(
            self._allocator.allocate_mutable(
                prev_block=self._blocks[-1],
                device=Device.GPU))

@staticmethod
def create(allocator_type, num_gpu_blocks, num_cpu_blocks, block_size):
    block_ids = list(range(num_gpu_blocks + num_cpu_blocks))
    gpu_block_ids = block_ids[:num_gpu_blocks]
    cpu_block_ids = block_ids[num_gpu_blocks:]

    if allocator_type == "naive":
        gpu_allocator = NaiveBlockAllocator(
            create_block=NaiveBlock, num_blocks=num_gpu_blocks,
            block_size=block_size, block_ids=gpu_block_ids)
        cpu_allocator = 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)

    return CpuGpuBlockAllocator(
        cpu_block_allocator=cpu_allocator,
        gpu_block_allocator=gpu_allocator)

def __init__(self, block_size, num_gpu_blocks, num_cpu_blocks,
             watermark=0.01, sliding_window=None, enable_caching=False):
    self.block_size = block_size
    self.num_total_gpu_blocks = num_gpu_blocks
    self.num_total_cpu_blocks = num_cpu_blocks
    self.watermark_blocks = int(watermark * num_gpu_blocks)
    self.enable_caching = enable_caching

    self.block_allocator = CpuGpuBlockAllocator.create(
        allocator_type="prefix_caching" if enable_caching else "naive",
        num_gpu_blocks=num_gpu_blocks,
        num_cpu_blocks=num_cpu_blocks,
        block_size=block_size,
    )

    self.block_tables: Dict[SeqId, BlockTable] = {}

def can_append_slots(self, seq_group, num_lookahead_slots) -> bool:
    num_touched_blocks = 0
    for seq in seq_group.get_seqs(status=SequenceStatus.RUNNING):
        block_table = self.block_tables[seq.seq_id]
        num_touched_blocks += (
            block_table.get_num_blocks_touched_by_append_slots(
                token_ids=block_table.get_unseen_token_ids(
                    seq.get_token_ids()),
                num_lookahead_slots=num_lookahead_slots,
            ))

    num_free_gpu_blocks = self.block_allocator.get_num_free_blocks(
        Device.GPU)
    return num_touched_blocks <= num_free_gpu_blocks