第3章 KV Block的设计

3. KV Block的设计：固定大小的逻辑块

好，咱们接着聊PagedAttention的核心——KV Block。说白了，这就是把KV Cache切成一块一块的，每块存固定数量token的Key和Value。你想想看，如果不切块，一整块连续内存存所有token的KV，那内存碎片和浪费问题会非常严重。

我个人习惯把KV Block比作「乐高积木」。每个积木大小一样，可以灵活拼接。模型推理时，需要多少token就拼多少块，不需要了就拆掉还给内存池。这个思路，其实跟操作系统的分页管理如出一辙。

3.1 为什么非要固定大小？

我在项目中遇到过一个问题：早期尝试用变长Block，每个序列按实际长度分配。结果呢？内存碎片化严重，分配和释放的开销也大。后来改成固定大小，问题迎刃而解。

固定大小的好处很明显：

分配简单：内存池里全是同样大小的槽位，分配时直接取一个，释放时直接还回去，O(1)复杂度
无外部碎片：所有Block大小一致，不会出现「大块空闲被切成小块」的情况
便于管理：Block的地址计算、索引查找都非常直接

核心要点：每个Block存储固定数量token的Key和Value。这个「固定数量」通常称为block_size，常见取值是16、32或64。block_size越小，内存利用率越高，但管理开销也越大；block_size越大，管理开销小，但可能浪费更多内存。

3.2 Block的内部结构

一个KV Block到底长什么样？咱们直接看代码：

// 伪代码：KV Block的数据结构
struct KVBlock {
    // Key的存储，形状: [num_layers, num_heads, block_size, head_dim]
    float* key_data;
    
    // Value的存储，形状: [num_layers, num_heads, block_size, head_dim]
    float* value_data;
    
    // 逻辑块ID，用于索引
    int block_id;
    
    // 引用计数，用于判断是否可以回收
    int ref_count;
    
    // 下一个空闲Block的指针（用于空闲链表）
    KVBlock* next_free;
};

嗯，这里要注意：实际实现中，key_data和value_data通常是一块连续内存，而不是分开的。为什么呢？因为GPU访问连续内存效率更高。我曾经踩过这个坑——把Key和Value分开存储，结果访存带宽利用率掉了30%。

每个Block存储的KV数据量是固定的：

参数	说明	典型值
block_size	每个Block存储的token数	16
num_heads	注意力头数	32
head_dim	每个头的维度	128
num_layers	Transformer层数	32

所以一个Block占用的显存大小 = block_size × num_heads × head_dim × num_layers × 2（Key和Value各一份）× 2字节（FP16）。算下来，block_size=16时，一个Block大约占用 16×32×128×32×2×2 = 8,388,608 字节，约8MB。嗯，这个数字不小，所以Block数量不能太多。

3.3 逻辑块 vs 物理块

这里有个关键概念：逻辑块和物理块是分开的。逻辑块是「虚拟地址」，物理块是「实际内存」。你想想看，一个序列的KV Cache在逻辑上是连续的，但在物理内存中可能分散在不同的物理Block里。

为什么会这样？因为PagedAttention的核心思想就是「非连续存储」。序列在生成过程中，前面的token可能被淘汰，后面的token继续追加。如果要求物理连续，那每次淘汰都要移动数据，代价太大了。

我的经验：逻辑块到物理块的映射，用一个简单的数组就能搞定。block_table[逻辑块ID] = 物理块ID。查找时直接索引，O(1)复杂度。我曾经试过用哈希表，但GPU上哈希冲突处理太慢，最后还是用数组最稳。

3.4 Block的分配与回收

Block的管理，说白了就是「谁用谁拿，用完归还」。vLLM里维护了一个空闲Block链表：

// 空闲Block管理
class BlockManager {
    KVBlock* free_list_head;  // 空闲链表头
    int total_blocks;         // 总Block数
    int used_blocks;          // 已使用Block数
    
    // 分配一个Block
    KVBlock* allocate() {
        if (free_list_head == nullptr) {
            // 内存不足，触发淘汰策略
            evict_blocks();
        }
        KVBlock* block = free_list_head;
        free_list_head = block->next_free;
        used_blocks++;
        return block;
    }
    
    // 释放一个Block
    void free(KVBlock* block) {
        block->next_free = free_list_head;
        free_list_head = block;
        used_blocks--;
    }
};

分配时，直接从链表头部取一个；释放时，插回链表头部。这个操作在GPU上可以做到非常高效，因为只涉及指针操作，没有内存拷贝。

避坑指南：我曾经在实现时忘了处理「引用计数」。多个序列可能共享同一个物理Block（比如beam search场景），如果直接释放，其他序列就崩了。所以每个Block必须维护ref_count，只有ref_count降到0才能回收。

3.5 Block大小怎么选？

block_size的选择是个权衡。我建议这样考虑：

小Block（8或16）：内存利用率高，适合短序列或内存紧张的场景。但管理开销大，Block数量多，查找表也大。
大Block（32或64）：管理开销小，适合长序列。但可能浪费内存——如果一个序列只有5个token，却要分配一个能存16个token的Block。

我个人习惯默认用16。为什么？因为大多数场景下，序列长度在几百到几千之间，16个token一个Block，既不会太碎，也不会太浪费。当然，如果你做的是超长文本生成（比如几万token），可以考虑用32或64。

3.6 核心流程图

下面这张图展示了KV Block从分配到使用的完整流程：

这张图展示了三个核心组件：空闲Block池、已分配Block、以及Block Table映射表。分配时从空闲池取，释放时还回去，映射表负责把逻辑地址翻译成物理地址。整个流程清晰明了。

3.7 实际项目中的坑

最后分享几个我踩过的坑：

Block对齐问题：GPU对内存对齐很敏感。Block的起始地址最好128字节对齐，否则访存性能会下降。我一开始没注意，结果性能差了15%。
并发分配：多线程同时分配Block时，空闲链表需要加锁。我建议用原子操作或者无锁队列，避免锁竞争。
Block预分配：不要等到需要时才分配。可以在初始化时预分配一批Block，减少运行时开销。

小技巧：如果你用CUDA，可以用cudaMallocPitch来分配Block内存，它会自动处理对齐。我后来一直这么用，省心不少。

好了，KV Block的设计就聊到这儿。记住一句话：固定大小、逻辑物理分离、引用计数管理。这三个点抓住了，PagedAttention的内存管理你就掌握了八成。

📚 课程目录

3. KV Block的设计：固定大小的逻辑块

3.1 为什么非要固定大小？

3.2 Block的内部结构

3.3 逻辑块 vs 物理块

3.4 Block的分配与回收

3.5 Block大小怎么选？

3.6 核心流程图

3.7 实际项目中的坑