vllm 深度解析：一切从 PagedAttention 谈起

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27

@dataclass(slots=True)
class KVCacheBlock:
    """KV-cache block metadata.
    元数据与数据分离: 不存储实际 KV 张量, 仅通过 block_id 计算显存偏移地址.
    """

    block_id: int                                  # 物理 block ID [0, num_gpu_blocks)
    ref_cnt: int = 0                               # 引用计数, =0 时可被回收
    _block_hash: BlockHashWithGroupId | None = None # 写满后的哈希值, prefix caching 用
    prev_free_block: "KVCacheBlock | None" = None   # 双向链表前驱 (FreeKVCacheBlockQueue 内部)
    next_free_block: "KVCacheBlock | None" = None   # 双向链表后继 (FreeKVCacheBlockQueue 内部)
    is_null: bool = False                          # null block 永不参与 prefix caching

    @property
    def block_hash(self) -> BlockHashWithGroupId | None:
        """获取 block 哈希值 (仅写满且被缓存后有效)"""
        return self._block_hash

    @block_hash.setter
    def block_hash(self, value: BlockHashWithGroupId):
        """设置 block 哈希值 (仅允许设置一次)"""
        assert self._block_hash is None, "block already has a hash"
        self._block_hash = value

    def reset_hash(self):
        """驱逐时清空哈希值, 使 block 重新可被分配"""
        self._block_hash = None

1
2
3
4
5

# block_table.py:264-275 (_compute_slot_mappings_kernel)
block_index    = position // block_size          # 逻辑块号
block_offset   = position %  block_size          # 块内偏移
physical_block = block_table[req_idx][block_index]  # 查表 → 物理块号
slot_id        = physical_block * block_size + block_offset

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43

class BlockTables:
    """GPU 端 Block Table: 管理 [max_num_reqs, max_num_blocks] 的 int32 映射张量."""

    def __init__(self, block_sizes: list[int], max_num_reqs: int,
                 max_num_batched_tokens: int, max_num_blocks_per_group: list[int],
                 device: torch.device, kernel_block_sizes: list[int]):
        self.num_kv_cache_groups = len(block_sizes)
        # 核心映射表: 行=请求, 列=逻辑块序号, 值=物理 block_id
        self.block_tables: list[StagedWriteTensor] = []
        for i in range(self.num_kv_cache_groups):
            max_blocks = max_num_blocks_per_group[i]
            self.block_tables.append(
                StagedWriteTensor((max_num_reqs, max_blocks), dtype=torch.int32, device=device)
            )
        # 每请求已分配的逻辑块数量 [num_groups, max_reqs]
        self.num_blocks = UvaBackedTensor(
            (self.num_kv_cache_groups, max_num_reqs), dtype=torch.int32)
        # 模型 forward 用的 block table 副本
        self.input_block_tables: list[torch.Tensor] = [
            torch.zeros_like(b.gpu) for b in self.block_tables]
        # 每个 token 到显存 slot 的直接映射 [num_groups, max_tokens]
        self.slot_mappings = torch.zeros(
            self.num_kv_cache_groups, max_num_batched_tokens,
            dtype=torch.int64, device=device)

    def append_block_ids(self, req_index: int, new_block_ids: tuple[list[int], ...],
                         overwrite: bool):
        """追加新的 block ID 到某请求行"""
        for i in range(self.num_kv_cache_groups):
            start = self.num_blocks.np[i, req_index] if not overwrite else 0
            self.block_tables[i].stage_write(req_index, start, new_block_ids[i])
            self.num_blocks.np[i, req_index] = start + len(new_block_ids[i])

    def compute_slot_mappings(self, idx_mapping: torch.Tensor,
                              query_start_loc: torch.Tensor,
                              positions: torch.Tensor) -> torch.Tensor:
        """计算每个 token 的 slot_id = physical_block * block_size + offset"""
        # Triton kernel 实现, 逻辑等价于:
        # block_index    = position // block_size
        # block_offset   = position %  block_size
        # physical_block = block_table[req_idx][block_index]
        # slot_id        = physical_block * block_size + block_offset
        ...

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20

class FreeKVCacheBlockQueue:
    """空闲 block 双向链表, 按 LRU 排序 (队头 = 最久未用).
    用 fake head/tail 哨兵节点减少边界判断分支, 支持 O(1) 中间删除.
    """

    def __init__(self, blocks: list[KVCacheBlock]):
        self.num_free_blocks = len(blocks)
        self.fake_free_list_head = KVCacheBlock(block_id=-1)  # 哨兵头, 永不弹出
        self.fake_free_list_tail = KVCacheBlock(block_id=-1)  # 哨兵尾, 永不弹出
        # 初始化: 将所有 block 串联为双向链表
        for i in range(self.num_free_blocks):
            if i > 0:
                blocks[i].prev_free_block = blocks[i - 1]
            if i < self.num_free_blocks - 1:
                blocks[i].next_free_block = blocks[i + 1]
        # 连接哨兵
        self.fake_free_list_head.next_free_block = blocks[0]
        blocks[0].prev_free_block = self.fake_free_list_head
        self.fake_free_list_tail.prev_free_block = blocks[-1]
        blocks[-1].next_free_block = self.fake_free_list_tail

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54

class BlockPool:
    """Block 池管理器: 负责分配, 释放, 缓存 KVCacheBlock."""

    def __init__(self, num_gpu_blocks: int, enable_caching: bool, hash_block_size: int):
        self.num_gpu_blocks = num_gpu_blocks
        self.enable_caching = enable_caching
        self.hash_block_size = hash_block_size
        # 所有 kv-cache blocks, 索引即为 block_id
        self.blocks: list[KVCacheBlock] = [
            KVCacheBlock(idx) for idx in range(num_gpu_blocks)
        ]
        # 空闲 block 双向链表 (含驱逐候选), LRU 排序
        self.free_block_queue = FreeKVCacheBlockQueue(self.blocks)
        # prefix caching 哈希索引: block_hash → KVCacheBlock
        self.cached_block_hash_to_block = BlockHashToBlockMap()
        # 占位 block (block_id=0), 滑动窗口等场景用于替换跳过的 block
        self.null_block = self.free_block_queue.popleft()
        self.null_block.is_null = True

    def get_new_blocks(self, num_blocks: int) -> list[KVCacheBlock]:
        """从空闲池分配新 block, 若启用 caching 则可能驱逐 LRU 缓存块"""
        if num_blocks > self.get_num_free_blocks():
            raise ValueError(f"Not enough free blocks: need {num_blocks}")
        ret = self.free_block_queue.popleft_n(num_blocks)
        for block in ret:
            if self.enable_caching:
                self._maybe_evict_cached_block(block)  # 驱逐 LRU 缓存块
            block.ref_cnt += 1
        return ret

    def touch(self, blocks: Sequence[KVCacheBlock]):
        """增加引用计数, prefix cache 命中时使用; ref_cnt=0 的 block 会从空闲队列移除"""
        for block in blocks:
            if block.ref_cnt == 0 and not block.is_null:
                self.free_block_queue.remove(block)
            block.ref_cnt += 1

    def free_blocks(self, ordered_blocks: Iterable[KVCacheBlock]):
        """释放 block 列表 (按驱逐优先级排序, 首块最先被驱逐)"""
        blocks_list = list(ordered_blocks)
        for block in blocks_list:
            block.ref_cnt -= 1
        # 仅回收 ref_cnt 归零且非 null 的 block
        self.free_block_queue.append_n(
            [b for b in blocks_list if b.ref_cnt == 0 and not b.is_null]
        )

    def get_num_free_blocks(self) -> int:
        return self.free_block_queue.num_free_blocks

    def get_usage(self) -> float:
        """返回 KV cache 利用率 [0.0, 1.0]"""
        total = self.num_gpu_blocks - 1  # 减去 null_block
        return 1.0 - (self.get_num_free_blocks() / total) if total else 0.0

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45

class BlockHashToBlockMap:
    """Prefix caching 哈希索引.
    单 block 时直接存 KVCacheBlock; 多 block 同哈希时退化为 dict 以减少 GC 开销.
    """

    def __init__(self):
        self._cache: dict[
            BlockHashWithGroupId,
            KVCacheBlock | dict[int, KVCacheBlock]
        ] = {}

    def get_one_block(self, key: BlockHashWithGroupId) -> KVCacheBlock | None:
        """根据哈希键获取任意一个匹配的 block"""
        blocks = self._cache.get(key)
        if blocks is None:
            return None
        if isinstance(blocks, KVCacheBlock):
            return blocks
        return next(iter(blocks.values()))  # 多 block 时返回第一个

    def insert(self, key: BlockHashWithGroupId, block: KVCacheBlock):
        """插入 block 到缓存"""
        existing = self._cache.get(key)
        if existing is None:
            self._cache[key] = block                           # 首次: 直接挂 block
        elif isinstance(existing, KVCacheBlock):
            self._cache[key] = {existing.block_id: existing,
                                block.block_id: block}         # 第二次: 升级为 dict
        else:
            existing[block.block_id] = block                   # 后续: 追加到 dict

    def pop(self, key: BlockHashWithGroupId, block_id: int) -> KVCacheBlock | None:
        """驱逐时按 block_id 精确弹出"""
        blocks = self._cache.pop(key, None)
        if blocks is None:
            return None
        if isinstance(blocks, KVCacheBlock):
            if blocks.block_id == block_id:
                return blocks
            self._cache[key] = blocks  # ID 不匹配, 放回去
            return None
        block = blocks.pop(block_id, None)
        if blocks:                     # dict 中还有剩余 block
            self._cache[key] = blocks
        return block

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66

class BlockTable:
    """Worker 端 Block Table: CPU/GPU 双缓冲的二维映射表."""

    def __init__(self, block_size: int, max_num_reqs: int,
                 max_num_blocks_per_req: int, max_num_batched_tokens: int,
                 pin_memory: bool, device: torch.device,
                 kernel_block_size: int):
        # Hybrid block: 当 block_size != kernel_block_size 时拆分
        if kernel_block_size == block_size:
            self.block_size = block_size
            self.blocks_per_kv_block = 1
            self.use_hybrid_blocks = False
        else:
            self.block_size = kernel_block_size
            self.blocks_per_kv_block = block_size // kernel_block_size
            self.use_hybrid_blocks = True

        self.max_num_blocks_per_req = max_num_blocks_per_req * self.blocks_per_kv_block
        # CPU/GPU 双缓冲映射表: [max_reqs, max_blocks_per_req] int32
        self.block_table = CpuGpuBuffer(
            max_num_reqs, self.max_num_blocks_per_req,
            dtype=torch.int32, device=device, pin_memory=pin_memory)
        # 每行有效 block 数量
        self.num_blocks_per_row = np.zeros(max_num_reqs, dtype=np.int32)
        # token → slot 映射
        self.slot_mapping = CpuGpuBuffer(
            max_num_batched_tokens, dtype=torch.int64,
            device=device, pin_memory=pin_memory)

    def append_row(self, block_ids: list[int], row_idx: int):
        """向指定请求行追加 block IDs"""
        if not block_ids:
            return
        if self.use_hybrid_blocks:
            block_ids = self.map_to_kernel_blocks(block_ids)
        start = self.num_blocks_per_row[row_idx]
        self.num_blocks_per_row[row_idx] += len(block_ids)
        self.block_table.np[row_idx, start:start + len(block_ids)] = block_ids

    def add_row(self, block_ids: list[int], row_idx: int):
        """覆盖写入指定请求行 (先清零再追加)"""
        self.num_blocks_per_row[row_idx] = 0
        self.append_row(block_ids, row_idx)

    def clear_row(self, row_idx: int):
        """清空指定请求行"""
        n = self.num_blocks_per_row[row_idx]
        if n > 0:
            self.block_table.np[row_idx, :n] = 0
        self.num_blocks_per_row[row_idx] = 0

    def commit_block_table(self, num_reqs: int):
        """将 CPU 端 block table 拷贝到 GPU"""
        self.block_table.copy_to_gpu(num_reqs)

    @staticmethod
    def map_to_kernel_blocks(kv_block_ids: list[int],
                             blocks_per_kv_block: int) -> list[int]:
        """将 KV manager block ID 展开为 kernel block ID.
        例: block_size=32, kernel_block_size=16, blocks_per_kv_block=2
            [0, 1] → [0, 1, 2, 3]
        """
        if blocks_per_kv_block == 1:
            return kv_block_ids
        return [b * blocks_per_kv_block + k
                for b in kv_block_ids for k in range(blocks_per_kv_block)]

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31

class MultiGroupBlockTable:
    """多组 KV cache 的 Block Table 聚合.
    为混合注意力模型 (full attention + sliding window) 设计,
    每个 KV cache group 独立管理自己的 block table.
    """

    def __init__(self, block_sizes: list[int], kernel_block_sizes: list[int],
                 max_num_reqs: int, max_model_len: int,
                 max_num_batched_tokens: int, pin_memory: bool,
                 device: torch.device):
        max_num_blocks = [cdiv(max_model_len, bs) for bs in block_sizes]
        self.block_tables: list[BlockTable] = [
            BlockTable(block_size, max_num_reqs, blocks_per_req,
                       max_num_batched_tokens, pin_memory, device, kbs)
            for block_size, kbs, blocks_per_req
            in zip(block_sizes, kernel_block_sizes, max_num_blocks)
        ]

    def append_row(self, block_ids: tuple[list[int], ...], row_idx: int):
        """向所有 group 的指定行追加 block IDs"""
        for i, bt in enumerate(self.block_tables):
            bt.append_row(block_ids[i], row_idx)

    def clear_row(self, row_idx: int):
        """清空所有 group 的指定行"""
        for bt in self.block_tables:
            bt.clear_row(row_idx)

    def __getitem__(self, idx: int) -> BlockTable:
        """获取第 i 个 KV cache group 的 BlockTable"""
        return self.block_tables[idx]

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21

@dataclass
class KVCacheBlocks:
    """Scheduler 与 KVCacheManager 之间的接口对象.
    隐藏 KVCacheManager 内部数据结构, Scheduler 只通过此对象交互.
    blocks[i][j] = 第 i 个 kv_cache_group 的第 j 个 block.
    """

    blocks: tuple[Sequence[KVCacheBlock], ...]

    def __add__(self, other: "KVCacheBlocks") -> "KVCacheBlocks":
        """合并两组 blocks (跨调度步骤时拼接)"""
        return KVCacheBlocks(
            tuple(list(a) + list(b) for a, b in zip(self.blocks, other.blocks)))

    def get_block_ids(self) -> tuple[list[int], ...]:
        """提取所有 block_id, 用于写入 Block Table"""
        return tuple([blk.block_id for blk in group] for group in self.blocks)

    def new_empty(self) -> "KVCacheBlocks":
        """创建同结构但无 block 的空对象"""
        return KVCacheBlocks(tuple(() for _ in range(len(self.blocks))))

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14

BlockHash = NewType("BlockHash", bytes)                     # 32 字节哈希
BlockHashWithGroupId = NewType("BlockHashWithGroupId", bytes)  # BlockHash + 4 字节 group_id

def make_block_hash_with_group_id(block_hash: BlockHash, group_id: int) -> BlockHashWithGroupId:
    """打包: BlockHash(32B) + group_id(4B big-endian)"""
    return BlockHashWithGroupId(block_hash + group_id.to_bytes(4, "big", signed=False))

def get_block_hash(key: BlockHashWithGroupId) -> BlockHash:
    """解包: 取前 32 字节"""
    return BlockHash(key[:-4])

def get_group_id(key: BlockHashWithGroupId) -> int:
    """解包: 取后 4 字节"""
    return int.from_bytes(key[-4:], "big", signed=False)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11

class PagedAttention:
    @staticmethod
    def split_kv_cache(kv_cache, num_kv_heads, head_size):
        # 将 [2, num_blocks, block_size, num_kv_heads, head_size]
        # 拆成 key_cache 和 value_cache
        # key: view(num_blocks, num_kv_heads, head_size//x, -1, x)
        # value: view(num_blocks, num_kv_heads, head_size, -1)

    @staticmethod
    def write_to_paged_cache(key, value, key_cache, value_cache, slot_mapping, ...):
        # 调用 ops.reshape_and_cache 将新 K/V 写入分页缓存

1
2
3
4

paged_attention_v1(out, query, key_cache, value_cache, num_kv_heads, scale,
                   block_tables, seq_lens, block_size, max_seq_len, ...)
paged_attention_v2(out, exp_sum, max_logits, tmp_out, query, key_cache, ...)
paged_attention_rocm(out, exp_sum, max_logits, tmp_out, query, ...)