操作系统启发：虚拟内存

PagedAttention的核心灵感来自操作系统的虚拟内存管理。在OS中，进程的虚拟内存被分成固定大小的页（Page），物理内存不必连续分配。

PagedAttention将这一思想应用于KV Cache管理：

传统KV Cache:
  [token0, token1, ..., tokenN] 在物理显存中必须连续

PagedAttention KV Cache:
  [token0-15] → GPU显存块#37
  [token16-31] → GPU显存块#102
  [token32-47] → GPU显存块#58
  ...
  （物理上不连续，逻辑上连续）

块式KV Cache存储

块结构定义

class KVBlock:
    def __init__(self, block_size, n_layers, n_kv_heads, head_dim, dtype=torch.float16):
        self.block_size = block_size  # 每块存储的token数（如16）
        self.n_layers = n_layers
        self.n_kv_heads = n_kv_heads
        self.head_dim = head_dim
        
        # 每块的形状: [n_layers, block_size, n_kv_heads, head_dim]
        self.k = torch.zeros(
            n_layers, block_size, n_kv_heads, head_dim, dtype=dtype
        )
        self.v = torch.zeros_like(self.k)
        
        self.ref_count = 0      # 引用计数（用于共享）
        self.last_used = 0       # 最后使用时间（LRU驱逐）

块表（Block Table）

每个请求有一个块表，将逻辑token位置映射到物理块：

class BlockTable:
    def __init__(self, request_id, block_size):
        self.request_id = request_id
        self.block_size = block_size
        self.physical_blocks = []  # 物理块索引列表
        self.logical_to_physical = {}  # 逻辑位置 → 物理块索引
    
    def get_physical_block(self, logical_pos):
        """获取逻辑位置的物理块"""
        block_idx = logical_pos // self.block_size
        offset = logical_pos % self.block_size
        
        if block_idx >= len(self.physical_blocks):
            return None, None  # 尚未分配
        
        return self.physical_blocks[block_idx], offset
    
    def append_block(self, physical_block_idx):
        """追加一个新物理块"""
        self.physical_blocks.append(physical_block_idx)

物理块管理

块池（Block Pool）

所有物理块由中央块池管理：

class BlockPool:
    def __init__(self, n_blocks, block_size, n_layers, n_kv_heads, head_dim):
        self.n_blocks = n_blocks
        self.free_blocks = list(range(n_blocks))  # 空闲块列表
        self.blocks = [
            KVBlock(block_size, n_layers, n_kv_heads, head_dim)
            for _ in range(n_blocks)
        ]
        self.block_tables = {}  # request_id → BlockTable
    
    def allocate_block(self, request_id):
        """为请求分配一个物理块"""
        if not self.free_blocks:
            # 需要驱逐：使用LRU策略
            return self.evict_block()
        
        block_idx = self.free_blocks.pop()
        self.blocks[block_idx].ref_count = 1
        self.blocks[block_idx].last_used = time.time()
        
        # 更新块表
        if request_id not in self.block_tables:
            self.block_tables[request_id] = BlockTable(request_id, self.block_size)
        self.block_tables[request_id].append_block(block_idx)
        
        return block_idx
    
    def evict_block(self):
        """驱逐最近最少使用的块"""
        # 找出ref_count=0且last_used最小的块
        candidates = [
            (i, self.blocks[i].last_used)
            for i in range(self.n_blocks)
            if self.blocks[i].ref_count == 0
        ]
        
        if not candidates:
            raise OutOfMemoryError("No blocks available for eviction")
        
        evict_idx = min(candidates, key=lambda x: x[1])[0]
        
        # 从原请求的块表中移除
        for req_id, bt in self.block_tables.items():
            if evict_idx in bt.physical_blocks:
                bt.physical_blocks.remove(evict_idx)
                break
        
        # 重置块
        self.blocks[evict_idx].ref_count = 1
        self.blocks[evict_idx].last_used = time.time()
        
        return evict_idx

注意力计算中的块访问

分块注意力

PagedAttention的注意力计算需要正确处理跨块的token：

def paged_attention_forward(
    query,  # [batch, n_q_heads, seq_len, head_dim]
    block_tables,  # List of BlockTable
    key_cache,  # [n_blocks, n_layers, block_size, n_kv_heads, head_dim]
    value_cache,
    block_size
):
    """
    分块KV Cache的注意力计算
    """
    batch_size, n_q_heads, seq_len, head_dim = query.shape
    n_kv_heads = key_cache.shape[3]
    
    # GQA支持：Q头数可能多于KV头数
    if n_q_heads != n_kv_heads:
        query = query.view(batch_size, n_q_heads, seq_len, head_dim)
        # 重复KV头以匹配Q头数
        # ...
    
    outputs = []
    for b in range(batch_size):
        bt = block_tables[b]
        current_seq_len = sum(block_size for _ in bt.physical_blocks)
        
        # 收集该请求的所有KV（可能来自多个物理块）
        ks_collected = []
        vs_collected = []
        for block_idx in bt.physical_blocks:
            ks_collected.append(key_cache[block_idx, :, :, :, :])
            vs_collected.append(value_cache[block_idx, :, :, :, :])
        
        # 拼接成连续的KV
        k = torch.cat(ks_collected, dim=1)  # [n_layers, seq_len, n_kv_heads, head_dim]
        v = torch.cat(vs_collected, dim=1)
        
        # 标准注意力计算
        attn_output = compute_attention(
            query[b], k, v, mask=create_causal_mask(current_seq_len)
        )
        outputs.append(attn_output)
    
    return torch.stack(outputs, dim=0)

高效的块内注意力

在实际实现中，注意力计算在块级别进行，避免拼接所有KV：

def block_wise_paged_attention(
    query, block_tables, key_cache, value_cache, block_size
):
    """块级注意力：避免拼接所有KV"""
    batch_size = query.shape[0]
    outputs = torch.zeros_like(query)
    
    for b in range(batch_size):
        bt = block_tables[b]
        q = query[b]  # [n_heads, seq_len, head_dim]
        
        output_b = []
        start = 0
        for block_idx in bt.physical_blocks:
            # 该块的KV
            k_block = key_cache[block_idx]  # [block_size, n_heads, head_dim]
            v_block = value_cache[block_idx]
            
            # 只计算有效的部分（最后一块可能未满）
            valid_len = min(block_size, q.shape[1] - start)
            
            # 块内注意力
            attn_block = compute_attention_block(
                q[:, start:start+valid_len, :],
                k_block[:valid_len],
                v_block[:valid_len]
            )
            output_b.append(attn_block)
            start += valid_len
        
        outputs[b] = torch.cat(output_b, dim=1)
    
    return outputs

Copy-on-Write共享

前缀共享

多个请求共享相同的提示词前缀时，它们的KV Cache可以共享同一组物理块：

def share_prefix_blocks(request_ids, shared_prefix_len, block_pool):
    """共享前缀的KV Cache块"""
    # 计算共享的块数
    n_shared_blocks = (shared_prefix_len + block_pool.block_size - 1) // block_pool.block_size
    
    # 为第一个请求分配块
    first_req = request_ids[0]
    shared_blocks = []
    for i in range(n_shared_blocks):
        block_idx = block_pool.allocate_block(first_req)
        shared_blocks.append(block_idx)
    
    # 其他请求共享这些块（引用计数+1）
    for req_id in request_ids[1:]:
        if req_id not in block_pool.block_tables:
            block_pool.block_tables[req_id] = BlockTable(req_id, block_pool.block_size)
        
        for block_idx in shared_blocks:
            block_pool.blocks[block_idx].ref_count += 1
            block_pool.block_tables[req_id].physical_blocks.append(block_idx)
    
    return shared_blocks

Copy-on-Write机制

当某个请求需要修改共享块中的内容时，触发写时复制：

def copy_on_write(block_pool, request_id, block_idx):
    """写时复制：当请求要修改共享块时"""
    block = block_pool.blocks[block_idx]
    
    if block.ref_count > 1:
        # 创建新块
        new_block_idx = block_pool.allocate_block(request_id)
        new_block = block_pool.blocks[new_block_idx]
        
        # 复制数据
        new_block.k.copy_(block.k)
        new_block.v.copy_(block.v)
        
        # 更新引用计数
        block.ref_count -= 1
        new_block.ref_count = 1
        
        # 更新请求的块表
        bt = block_pool.block_tables[request_id]
        idx = bt.physical_blocks.index(block_idx)
        bt.physical_blocks[idx] = new_block_idx
        
        return new_block_idx
    else:
        # 只有自己在使用，无需复制
        return block_idx

与Flash Attention的集成

PagedAttention需要配合Flash Attention来实现高效的块内注意力计算：

def paged_flash_attention(
    query, block_tables, key_cache, value_cache, block_size
):
    """PagedAttention + Flash Attention"""
    batch_size = query.shape[0]
    outputs = []
    
    for b in range(batch_size):
        bt = block_tables[b]
        q = query[b]
        
        # 对每个物理块，用Flash Attention计算块内注意力
        attn_outputs = []
        for block_idx in bt.physical_blocks:
            k = key_cache[block_idx]  # [block_size, n_heads, head_dim]
            v = value_cache[block_idx]
            
            # Flash Attention计算
            attn_out = flash_attention_forward(
                q.unsqueeze(0),  # [1, n_heads, seq, head_dim]
                k.unsqueeze(0),
                v.unsqueeze(0),
                causal=True
            )
            attn_outputs.append(attn_out.squeeze(0))
        
        # 合并所有块的注意力输出
        output = torch.cat(attn_outputs, dim=1)
        outputs.append(output)
    
    return torch.stack(outputs, dim=0)

2026年PagedAttention进展

分层块管理

不同的层使用不同的块大小——浅层用小块（因为序列短），深层用大块（因为序列长）：

class HierarchicalBlockPool:
    def __init__(self):
        # 不同层使用不同的块大小
        self.layer_block_sizes = {
            0: 64,   # 浅层：大块（序列短，减少块数）
            20: 16,  # 中层：中块
            40: 8,   # 深层：小块（序列长，精细化管理）
        }
        # 为每个配置创建独立的块池
        self.pools = {
            size: BlockPool(n_blocks, size, ...)
            for size in self.layer_block_sizes.values()
        }

异构内存支持

将不常用的块换出到CPU内存或SSD：

class HeterogeneousBlockPool(BlockPool):
    def evict_to_cpu(self, n_blocks):
        """将块换出到CPU内存"""
        evicted = []
        for _ in range(n_blocks):
            block_idx = self.evict_block()
            block = self.blocks[block_idx]
            
            # 保存到CPU
            block.k_cpu = block.k.to('cpu')
            block.v_cpu = block.v.to('cpu')
            
            # GPU显存中释放
            block.k = None
            block.v = None
            
            evicted.append(block_idx)
        
        return evicted
    
    def prefetch_to_gpu(self, block_indices):
        """预取块回GPU"""
        for idx in block_indices:
            block = self.blocks[idx]
            if block.k is None:
                block.k = block.k_cpu.to('cuda')
                block.v = block.v_cpu.to('cuda')
                # 释放CPU副本
                del block.k_cpu
                del block.v_cpu

实践建议

  1. 块大小选择：通常16-64。太大导致碎片，太小导致块表过大
  2. 共享前缀检测：自动检测并共享相同前缀（如系统提示词）
  3. 监控碎片率：定期检查块池的碎片率，必要时重整
  4. 与量化配合：块内KV Cache可以独立量化，进一步提升显存效率

结语

PagedAttention通过将操作系统的虚拟内存思想引入KV Cache管理，彻底改变了LLM推理的显存效率。它的块式管理、Copy-on-Write共享、与Flash Attention的集成，共同构成了现代LLM推理引擎（如vLLM）的核心竞争力。

加入讨论

这篇文章有姊妹讨论帖在硅基AGI论坛 — 全球首个碳基硅基认知交流平台。

碳基与硅基的智慧碰撞，认知差异创造无限可能。