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

Request → Scheduler → KVCacheManager.allocate(request)
  → SingleTypeKVCacheManager.get_num_blocks_to_allocate()
  → BlockPool.allocate() → 返回 KVCacheBlock 列表
  → req_to_blocks[req_id] = [blocks]
  → Scheduler 拿到 KVCacheBlocks（含 block_id 列表）
  → Worker 将 block_id 写入 BlockTable

@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

示例: block_size=16, 两个请求共享前缀 "Hello, how are"

时间线:
  T1: Request A 生成 "Hello, how are" 的 KV cache, 写满 block P3
      → P3._block_hash = hash("Hello, how are")
      → P3.ref_cnt = 1

  T2: Request B 以 "Hello, how are" 开头, prefix hash 命中 P3
      → touch(P3): P3.ref_cnt = 2   ← 两个请求共享同一物理 block

  T3: Request A 结束, free_blocks()
      → P3.ref_cnt = 1              ← B 仍在用, P3 不会回到空闲队列

  T4: Request B 结束, free_blocks()
      → P3.ref_cnt = 0              ← 无人引用, 回到 FreeKVCacheBlockQueue

Request 1 (25 tokens, 2 个逻辑块):

  逻辑块 0 (token 0~15)      逻辑块 1 (token 16~24)
  ┌────────────────┬────────────────┐
  │████████████████│█████████░░░░░░░│
  └────────────────┴────────────────┘
     16/16 满          9/16 部分

Request 2 (12 tokens, 1 个逻辑块):

  逻辑块 0 (token 0~11)
  ┌────────────────┐
  │████████████░░░░│
  └────────────────┘
    12/16 部分

GPU KV Cache 物理地址空间 (block_size=16, 每个物理 block 可存 16 个 token 的 KV):

┌────────┬────────┬────────┬────────┬────────┬────────┬────────┬────────┐
│  P0    │  P1    │  P2    │  P3    │  P4    │  P5    │  P6    │  P7    │
│  FREE  │ Req2   │  FREE  │ Req1   │  FREE  │  FREE  │  FREE  │ Req1   │
│        │ T0~T11 │        │ T0~T15 │        │        │        │T16~T24 │
│  [  ]  │ [██▓]  │  [  ]  │ [███]  │  [  ]  │  [  ]  │  [  ]  │ [██▓]  │
└────────┴────────┴────────┴────────┴────────┴────────┴────────┴────────┘
         ↑                  ↑                                   ↑
      被 Req2 占用       被 Req1 占用                       被 Req1 占用
      (逻辑块0)        (逻辑块0, 满)                     (逻辑块1, 部分满)

FreeKVCacheBlockQueue 双向链表:
  HEAD ⇄ P0 ⇄ P2 ⇄ P4 ⇄ P5 ⇄ P6 ⇄ TAIL
         ↑ 这些物理 block 空闲，等待分配

BlockTable (行 = 请求, 列 = 逻辑块序号, 值 = 物理 block id):

            逻辑块0  逻辑块1  逻辑块2  ...  逻辑块N-1
           ┌────────┬────────┬────────┬──┬────────┐
  Req1     │   P3   │   P7   │   -1   │..│   -1   │  num_blocks=2
           ├────────┼────────┼────────┼──┼────────┤
  Req2     │   P1   │   -1   │   -1   │..│   -1   │  num_blocks=1
           ├────────┼────────┼────────┼──┼────────┤
  Req3     │   -1   │   -1   │   -1   │..│   -1   │  (空闲槽位，未分配)
           ├────────┴────────┴────────┴──┴────────┤
  ...                                              │
           └───────────────────────────────────────┘

  -1 = 无效/未分配,  num_blocks = 该请求已分配的逻辑块数量

输入: req_idx=Req1, position=18
                         │
                         ▼
  ┌──────────────────────────────────────────────────┐
  │ Step 1: position → 逻辑块索引 + 块内偏移          │
  │                                                    │
  │   block_index  = 18 // 16 = 1    (第 1 号逻辑块)   │
  │   block_offset = 18 %  16 = 2    (块内第 2 个位置) │
  └──────────────────────┬───────────────────────────┘
                         │
                         ▼
  ┌──────────────────────────────────────────────────┐
  │ Step 2: 查 BlockTable 获取物理 block              │
  │                                                    │
  │   physical_block = BlockTable[Req1][1] = P7        │
  └──────────────────────┬───────────────────────────┘
                         │
                         ▼
  ┌──────────────────────────────────────────────────┐
  │ Step 3: 计算最终 slot_id，访问 GPU 显存            │
  │                                                    │
  │   slot_id = P7 * 16 + 2                           │
  │           = 物理 block 基址 + 块内偏移              │
  │                                                    │
  │   访问: k_cache[slot_id, kv_head, ...]             │
  └──────────────────────────────────────────────────┘

# 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

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
        ...

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

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

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

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)]

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]

@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))))

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)

                            Scheduler
                               |
                    KVCacheManager.allocate_slots()
                               |
               +---------------+---------------+
               |                               |
         SingleTypeKVCacheManager    SingleTypeKVCacheManager
          (FullAttention)              (SlidingWindow)
               |                               |
               +---------------+---------------+
                               |
                          BlockPool
                    +-------+-------+
                    |               |
          FreeKVCacheBlockQueue   BlockHashToBlockMap
            (doubly linked, LRU)   (hash index, prefix cache)
                    |
              KVCacheBlock[]
              (physical block metadata)

                    |  allocation result

              KVCacheBlocks
             (Scheduler interface object)
                    |
                    |  get_block_ids()

            BlockTables / BlockTable
         [max_num_reqs, max_num_blocks]  <- GPU int32 tensor
                    |
                    |  compute_slot_mappings()

            slot_mappings[token] = physical_block * block_size + offset
                    |
                    v
            CUDA Kernel reads k_cache[slot_id] / v_cache[slot_id]

q_vecs[THREAD_GROUP_SIZE][NUM_VECS_PER_THREAD]  // 共享内存
每个 thread group 读取同一 query token 的不同部分
线程 0 读取 vec [0,4,8,...]，线程 1 读取 vec [1,5,9,...]

k_vecs[NUM_VECS_PER_THREAD]  // 寄存器
通过 block_table 间接寻址获取物理 block
k_cache[physical_block_number, kv_head, physical_block_offset, vec_offset]
每个 warp 的 thread group 处理 block 中的不同 token

1. QK dot product（thread group 内规约） + scale + ALiBi bias
2. logits[] 存储到共享内存
3. qk_max = warp-reduce(fmax) → block-reduce(fmax) → broadcast
4. exp_sum = sum(exp(logits[i] - qk_max)) → block_reduce
5. logits = exp(logits - qk_max) / exp_sum  (softmax)

1. 按 block 迭代，从 v_cache[physical_block, kv_head, head_size, block_size] 读取
2. v_vec 与 logits_vec 做 dot product → accs[NUM_ROWS_PER_THREAD]
3. Warp 内规约 accs → Warp 间规约 accs
4. 最终输出到 out[seq_idx, head_idx, head_size]

1. V2 Kernel: 每个 partition 计算部分 KV blocks 的 attention
   输出: tmp_out[partition], exp_sums[partition], max_logits[partition]
2. Reduce Kernel: 用 LSE (Log-Sum-Exp) 算法合并各 partition:
   global_max = max(max_logits[:])
   weight[i] = exp_sums[i] * exp(max_logits[i] - global_max)
   out = sum(tmp_out[i] * weight[i]) / sum(weight)

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 写入分页缓存

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, ...)

1. 模型 forward
   → Attention 层计算新 K/V tensor（当前 token）
   → PagedAttention.write_to_paged_cache(key, value, key_cache, value_cache, slot_mapping)
   → reshape_and_cache 将 [num_tokens, heads, dim] → [blocks, block_size, heads, dim]

2. 调度器调度
   → KVCacheManager.get_num_blocks_to_allocate(request)
   → BlockPool.allocate(num_blocks)
   → 返回 KVCacheBlocks(blocks=[KVCacheBlock(block_id=3), ...])

3. Worker 准备输入
   → BlockTables.append_block_ids(req_index, new_block_ids)
   → 构建 input_block_tables: [num_reqs, max_blocks] 的 int32 tensor
   → 传递给 CUDA kernel

4. CUDA kernel 执行 PagedAttention
   → 通过 block_table[seq_idx] 查物理 block_id
   → k_cache[physical_block, kv_head, ...] 读取 K
   → v_cache[physical_block, kv_head, ...] 读取 V
   → 计算 attention 输出

5. 请求完成
   → KVCacheManager.free(request)
   → BlockPool.free(block_ids)
   → 物理块回收复用