3.2 动态页表设计


3.2 动态页表设计

PagedAttention的神经中枢:动态页表管理

3.2.1 动态页表的核心概念

动态页表(在vLLM源码中称为BlockTable)是PagedAttention架构的核心组件,负责管理逻辑token位置与物理内存页面之间的映射关系。与传统操作系统中静态预分配的页表不同,PagedAttention的动态页表能够在推理运行时根据实际生成需求实时增长、调整页面分配,并支持通过Copy-on-Write机制在多个请求之间共享页面映射。

页表的基本功能

四大核心功能:

  1. 地址映射:建立逻辑页面与物理页面的对应关系,使得Attention计算可以通过逻辑索引间接访问分散的物理KV Cache
  2. 动态内存分配:根据序列生成进度按需分配新页面,生成多少token就分配多少页面,不多不少
  3. 共享管理:通过引用计数机制管理共享页面的生命周期,当多个请求引用同一页面时阻止其被回收
  4. 性能桥接:作为逻辑序列空间和物理显存空间之间的桥梁,使得非连续存储对上层计算透明

动态特性的实现

与传统页表的关键区别:

  • 静态页表(操作系统):在进程创建时预分配,大小固定,进程终止时整体释放
  • 动态页表(PagedAttention):运行时随token生成动态增长,新token到来时追加页表项,请求完成时整体释放
  • 智能适配:根据实际生成进度自动扩展,不同请求的页表大小可以不同,完全独立

在vLLM的具体实现中,每个序列的块表是一个整数数组,数组的每个元素是一个物理页面编号。逻辑页面索引就是数组下标,物理页面编号就是数组值。这种极简的一级页表设计牺牲了通用性,但在GPU推理场景下获得了最佳的查询性能。

# vLLM块表的核心数据结构模拟 class BlockTableManager: """模拟vLLM中的块表管理器""" def __init__(self, block_size: int = 16): self.block_size = block_size # 每个序列的块表:{seq_id: [block_number_0, block_number_1, ...]} self.seq_block_tables = {} # 空闲块列表(物理页面池) self.free_block_ids = [] # 块的引用计数 self.block_ref_count = {} def append_token(self, seq_id: int, num_new_tokens: int = 1): """为序列追加token,按需分配新块""" if seq_id not in self.seq_block_tables: self.seq_block_tables[seq_id] = [] table = self.seq_block_tables[seq_id] current_num_tokens = len(table) * self.block_size new_total_tokens = current_num_tokens + num_new_tokens new_num_blocks = (new_total_tokens + self.block_size - 1) // self.block_size # 如果需要新块,从空闲池中分配 while len(table) < new_num_blocks: if not self.free_block_ids: raise RuntimeError( f"OOM: 无法为新token分配块。" f"当前序列 {seq_id} 已有 {len(table)} 块," f"需要 {new_num_blocks} 块" ) block_id = self.free_block_ids.pop() table.append(block_id) self.block_ref_count[block_id] = 1 def get_block_table_gpu_tensor(self, seq_id: int): """获取用于CUDA kernel的块表张量""" return self.seq_block_tables[seq_id] def free_sequence(self, seq_id: int): """释放序列占用的所有块""" if seq_id not in self.seq_block_tables: return for block_id in self.seq_block_tables[seq_id]: self.block_ref_count[block_id] -= 1 if self.block_ref_count[block_id] <= 0: self.free_block_ids.append(block_id) del self.block_ref_count[block_id] del self.seq_block_tables[seq_id] # 初始化:模拟24GB GPU的块池 num_blocks = (24 * 1024**3 * 0.8) // (16 * 2 * 32 * 128 * 2) # 约61440块 btm = BlockTableManager(block_size=16) btm.free_block_ids = list(range(num_blocks)) # 模拟序列生成过程 btm.append_token(seq_id=1, num_new_tokens=500) # 分配32个块 btm.append_token(seq_id=1, num_new_tokens=200) # 再分配12个块 print(f"序列1块表长度: {len(btm.seq_block_tables[1])} 块") print(f"剩余空闲块: {len(btm.free_block_ids)}")

3.2.2 块表的GPU端实现

块表不仅需要在CPU端管理,更需要高效地在GPU端被CUDA kernel访问。vLLM的PagedAttention kernel在计算注意力时,需要通过块表将每个token的逻辑位置转换为物理内存地址,这一过程必须在GPU上高效完成。

GPU端块表存储格式

在vLLM中,块表以一个一维整数张量的形式存储在GPU显存中。为了支持动态批处理中多个序列的同时查询,vLLM将所有活跃序列的块表拼接成一个大的张量,并通过偏移量快速定位每个序列的块表起始位置。

存储布局设计:

  • max_num_blocks_per_seq:所有活跃序列中块表的最大长度,决定了张量的行数
  • 块表张量形状:[num_seqs, max_num_blocks_per_seq],数据类型为int32
  • 块表序列长度张量:记录每个序列实际使用的块数量,形状为[num_seqs]

这种设计使得CUDA kernel可以通过一次索引操作就获得任意序列任意逻辑页面对应的物理块编号,查询延迟仅为一次全局内存读取。

# 模拟GPU端块表张量的构建 class GPUBlockTable: """模拟GPU端的块表张量管理""" def __init__(self, block_size: int = 16): self.block_size = block_size self.seq_tables = {} # {seq_id: [block_ids...]} def build_batch_tensor(self, seq_ids: list): """构建用于CUDA kernel的批处理块表张量 在实际vLLM中,这会生成一个torch.IntTensor传输到GPU """ max_blocks = max( len(self.seq_tables[sid]) for sid in seq_ids ) num_seqs = len(seq_ids) # 创建块表张量 [num_seqs, max_blocks] # 未使用的位置填充-1 block_table_tensor = [] seq_lengths = [] for sid in seq_ids: table = self.seq_tables[sid] padded = table + [-1] * (max_blocks - len(table)) block_table_tensor.append(padded) seq_lengths.append(len(table)) return { 'block_table': block_table_tensor, 'seq_lengths': seq_lengths, 'shape': (num_seqs, max_blocks) } # 演示 gpu_bt = GPUBlockTable(block_size=16) gpu_bt.seq_tables = { 1: [42, 17, 88, 3], # 序列1: 4个块(64个token) 2: [55, 23, 91], # 序列2: 3个块(48个token) 3: [10, 66, 33, 77, 99] # 序列3: 5个块(80个token) } batch = gpu_bt.build_batch_tensor([1, 2, 3]) print(f"批处理块表形状: {batch['shape']}") print(f"各序列块数量: {batch['seq_lengths']}") print(f"序列1块表: {batch['block_table'][0]}")

PagedAttention Kernel中的块表查询

在vLLM的CUDA kernel实现中,块表查询是注意力计算的第一步。kernel接收块表张量后,对于需要访问的每个历史token,通过以下步骤获取其KV Cache的物理地址:

// 伪代码:PagedAttention CUDA kernel中的块表查询 __global__ void paged_attention_kernel( // ... 其他参数 ... const int* __restrict__ block_tables, // 块表张量 [num_seqs, max_blocks] const int* __restrict__ seq_lengths, // 各序列的块数量 const int block_size, // 每块的token数 const int num_seqs, // 批处理中的序列数 const int max_num_blocks_per_seq // 块表最大列数 ) { // 每个线程处理一个query token对一个key token的注意力计算 int seq_idx = blockIdx.y; // 序列索引 int query_idx = blockIdx.x * blockDim.x + threadIdx.x; int key_idx = threadIdx.y; // 要访问的历史token索引 // 步骤1:计算逻辑块号和块内偏移 int logical_block = key_idx / block_size; int block_offset = key_idx % block_size; // 步骤2:通过块表查询物理块号 int physical_block = block_tables[seq_idx * max_num_blocks_per_seq + logical_block]; // 步骤3:计算物理内存地址 float* key_ptr = key_cache + physical_block * block_size * head_dim + block_offset * head_dim; // 步骤4:读取KV向量,计算注意力分数 float score = dot_product(query_ptr, key_ptr, head_dim); // ... 后续softmax和加权求和 ... }

3.2.3 块表与调度器的协同

块表不是孤立工作的,它与vLLM的调度器(Scheduler)紧密配合,共同实现了高效的GPU资源管理。

调度器驱动的块表操作

vLLM的调度器在每次调度决策时,会与块表管理器进行以下交互:

  1. 新请求到达时:调度器检查空闲块数量,决定是否接纳新请求。如果有足够的空闲块来容纳新请求的prompt token(即使只是预估值),则接纳并为其创建空块表
  2. Prefill阶段:将输入prompt的KV Cache写入新分配的物理块,并填充块表
  3. Decode阶段:每生成一个新token,检查是否需要分配新块(当当前最后一个块已满时),如果显存不足则暂停该请求(swapping或preemption)
  4. 请求完成时:释放请求块表中所有块的引用,将空闲块归还到全局池
# 模拟调度器与块表的协同工作 class SimpleScheduler: """简化的vLLM调度器模拟""" def __init__(self, block_table_mgr: BlockTableManager): self.btm = block_table_mgr self.waiting_queue = [] # 等待队列 self.running_seqs = {} # 运行中的序列 {seq_id: num_generated_tokens} def add_request(self, seq_id: int, prompt_length: int): """新请求到达""" self.waiting_queue.append((seq_id, prompt_length)) def schedule(self): """执行一次调度决策""" # 尝试从等待队列中接纳新请求 still_waiting = [] for seq_id, prompt_len in self.waiting_queue: blocks_needed = (prompt_len + self.btm.block_size - 1) // self.btm.block_size if len(self.btm.free_block_ids) >= blocks_needed: self.btm.append_token(seq_id, prompt_len) self.running_seqs[seq_id] = prompt_len print(f"[调度] 接纳序列 {seq_id}(prompt={prompt_len} token, 分配{blocks_needed}块)") else: still_waiting.append((seq_id, prompt_len)) print(f"[调度] 序列 {seq_id} 等待中(需要{blocks_needed}块,空闲{len(self.btm.free_block_ids)}块)") self.waiting_queue = still_waiting def decode_step(self, seq_id: int): """执行一步decode,生成1个token""" try: self.btm.append_token(seq_id, 1) self.running_seqs[seq_id] += 1 except RuntimeError as e: print(f"[OOM] 序列 {seq_id} 被抢占: {e}") # 演示调度流程 btm2 = BlockTableManager(block_size=16) btm2.free_block_ids = list(range(20)) # 模拟只有20个空闲块 scheduler = SimpleScheduler(btm2) scheduler.add_request(1, 100) # prompt=100 token scheduler.add_request(2, 200) # prompt=200 token scheduler.add_request(3, 150) # prompt=150 token scheduler.schedule() print(f"\n接纳后空闲块: {len(btm2.free_block_ids)}")
flowchart TD A[请求到达] --> B{空闲块足够?} B -->|是| C[分配块,创建块表] B -->|否| D[进入等待队列] C --> E[Prefill: 写入KV Cache] E --> F[Decode: 逐token生成] F --> G{当前块已满?} G -->|是| H{有空闲块?} H -->|是| I[分配新块,追加块表] H -->|否| J[暂停/抢占该请求] G -->|否| K[写入当前块] I --> F K --> F F --> L{生成完成?} L -->|是| M[释放所有块,归还块池] L -->|否| F

3.2.4 块表的序列化与传输

在分布式推理场景下,块表需要在不同GPU之间传输。vLLM通过将块表序列化为紧凑的整数列表来实现高效传输。

张量并行下的块表

在张量并行(Tensor Parallelism)模式下,每个GPU只存储部分注意力头的KV Cache,因此每个GPU都有自己独立的物理块池和块表。但所有GPU上的块表逻辑结构保持一致——同一个逻辑页面在不同GPU上映射到各自的物理块号。

这意味着在张量并行中,调度器需要为每个GPU独立管理块分配,但分配决策是同步的:分配一个逻辑块时,所有GPU同时分配各自的物理块,并更新各自的块表。

# 模拟张量并行下的块表同步 class TPBlockTableSync: """张量并行下的块表同步管理""" def __init__(self, num_gpus: int, blocks_per_gpu: int, block_size: int = 16): self.num_gpus = num_gpus self.block_size = block_size # 每个GPU有独立的块管理器 self.gpu_managers = [ BlockTableManager(block_size=block_size) for _ in range(num_gpus) ] for mgr in self.gpu_managers: mgr.free_block_ids = list(range(blocks_per_gpu)) def allocate_sync(self, seq_id: int, num_tokens: int): """在所有GPU上同步分配块""" # 先检查所有GPU是否有足够的空闲块 blocks_needed = (num_tokens + self.block_size - 1) // self.block_size for i, mgr in enumerate(self.gpu_managers): if len(mgr.free_block_ids) < blocks_needed: raise RuntimeError(f"GPU {i} 空闲块不足: " f"需要{blocks_needed}, 可用{len(mgr.free_block_ids)}") # 所有GPU都有足够空间,同步分配 for i, mgr in enumerate(self.gpu_managers): mgr.append_token(seq_id, num_tokens) print(f"[TP同步] 在{self.num_gpus}个GPU上为序列{seq_id}分配了{blocks_needed}块") # 打印各GPU的物理块映射(不同GPU物理块号不同) for i, mgr in enumerate(self.gpu_managers): print(f" GPU {i} 块表: {mgr.seq_block_tables[seq_id][:5]}...") # 演示:4卡张量并行 tp_sync = TPBlockTableSync(num_gpus=4, blocks_per_gpu=50, block_size=16) tp_sync.allocate_sync(seq_id=1, num_tokens=256)

3.2.5 前缀缓存中的块表优化

vLLM从0.3.0版本开始引入了自动前缀缓存(Automatic Prefix Caching, APC)功能,这是块表设计的一个重要扩展。APC利用块表的映射能力,自动检测多个请求之间相同的输入前缀,并复用对应的物理页面和已计算的KV Cache。

前缀匹配的哈希机制

APC通过计算每个页面的内容哈希值来实现前缀匹配。具体来说,当一个请求的输入token被填充到物理页面后,系统会计算每个页面的KV Cache内容的哈希值,并将(哈希值, 物理页面号)的映射关系存入一个全局的哈希表中。当新请求到来时,系统将其prompt的页面哈希与全局哈希表进行比对,找到最长匹配的前缀长度。

这个机制的关键在于:块表使得前缀匹配可以在页面粒度上高效进行。系统不需要逐token比较,而是逐页面比较哈希值,将比较次数从O(n)降低到O(n/block_size),极大地提高了前缀匹配的效率。

# 模拟前缀缓存的块表复用 class PrefixCache: """基于块表的前缀缓存管理器""" def __init__(self, allocator: BlockAllocator): self.allocator = allocator self.hash_to_block = {} # {page_hash: physical_block_id} self.block_refcount = {} # {block_id: ref_count} def compute_page_hash(self, tokens: list) -> int: """计算页面内容的哈希值(简化模拟)""" return hash(tuple(tokens)) def try_match_prefix(self, prompt_tokens: list, block_size: int = 16) -> tuple: """尝试匹配已缓存的前缀页面 返回: (matched_blocks, new_tokens_start_idx) """ matched_blocks = [] for i in range(0, len(prompt_tokens), block_size): page_tokens = prompt_tokens[i:i+block_size] page_hash = self.compute_page_hash(page_tokens) if page_hash in self.hash_to_block: matched_blocks.append(self.hash_to_block[page_hash]) else: break unmatched_start = len(matched_blocks) * block_size return matched_blocks, unmatched_start def allocate_with_prefix_cache(self, seq_id: int, prompt_tokens: list, block_size: int = 16): """使用前缀缓存分配页面""" matched_blocks, unmatched_start = self.try_match_prefix(prompt_tokens, block_size) # 共享已匹配的前缀块 if matched_blocks: self.allocator.allocated_blocks[seq_id] = matched_blocks.copy() for block_id in matched_blocks: self.block_refcount[block_id] = self.block_refcount.get(block_id, 0) + 1 # 为未匹配部分分配新块 remaining_tokens = prompt_tokens[unmatched_start:] if remaining_tokens: self.allocator.allocate(seq_id, len(remaining_tokens)) print(f"前缀缓存: 匹配{len(matched_blocks)}页(跳过{unmatched_start}token), " f"新分配{(len(remaining_tokens)+block_size-1)//block_size}页") # 演示:两个请求共享prompt alloc = BlockAllocator(num_blocks=100) prefix_cache = PrefixCache(alloc) prompt = list(range(160)) # 160 token的prompt prefix_cache.allocate_with_prefix_cache(1, prompt) print(f"请求1块表: {alloc.allocated_blocks[1][:5]}...") # 请求2有相同的前120 token + 不同的后40 token prompt2 = list(range(120)) + [999]*40 prefix_cache.allocate_with_prefix_cache(2, prompt2) print(f"请求2块表: {alloc.allocated_blocks[2][:5]}...")

3.2.6 块表内存开销分析

块表本身也需要占用一定的内存,虽然相比KV Cache的开销微不足道,但在大规模并发场景下仍值得分析。

块表内存开销计算:

  • 每个块表项(int32):4字节
  • 每个序列的平均块数:假设平均生成长度1024 token,每块16 token,则约64个块
  • 每个序列的块表开销:64 × 4 = 256字节
  • 1000个并发序列的块表总开销:256KB

相比KV Cache本身的GB级开销,块表仅占约0.01%-0.1%的额外内存,几乎可以忽略不计。这也是一级页表设计的一大优势——结构简单,元数据开销极低。

graph TB subgraph 块表开销 vs KV Cache开销 A[块表元数据<br/>256字节/序列<br/>1000序列=256KB] B[KV Cache数据<br/>256KB/token<br/>1024 token/序列<br/>1000序列=256MB] end A ---|"开销比: 1:1000"| B

3.2.7 块表性能优化技巧

虽然块表查询本身的开销很小,但在高频调用的Attention kernel中,任何微小的优化都会被放大数百万倍。

关键优化技术:

  1. 块表缓存到共享内存:在CUDA kernel启动时,将当前序列的块表预加载到共享内存中。由于块表通常只有几十到几百个int32值,完全可以放入共享内存(通常为48-96KB),从而将每次查询从全局内存延迟(约400-800个时钟周期)降低到共享内存延迟(约20-30个时钟周期),性能提升超过10倍。在vLLM的实际实现中,这个优化是PagedAttention kernel达到接近连续KV Cache性能的关键因素之一

  2. 合并块表访问:同一warp中的线程通常访问同一个序列的连续逻辑块,利用GPU的合并访问特性一次加载多个块表项

  3. 块表预取:在计算当前块的注意力时,通过__ldg()内置函数预取下一个块表项到缓存中,隐藏内存访问延迟

  4. 块表压缩:对于短序列(块表只有1-2项),可以绕过块表直接硬编码物理地址,避免间接寻址开销

# 模拟共享内存中的块表缓存优化 # 在实际CUDA代码中,这通过 __shared__ int block_table_cache[MAX_BLOCKS] 实现 def simulate_shared_memory_cache_hit_rate(seq_length: int, block_size: int = 16): """模拟块表在共享内存中的缓存命中率""" num_blocks = (seq_length + block_size - 1) // block_size cache_size_kb = 48 # 典型的共享内存大小 int32_size = 4 max_cacheable_blocks = (cache_size_kb * 1024) // int32_size # 12288个块 if num_blocks <= max_cacheable_blocks: hit_rate = 1.0 strategy = "完全缓存" else: hit_rate = max_cacheable_blocks / num_blocks strategy = f"部分缓存(缓存{max_cacheable_blocks}块中的前{max_cacheable_blocks}块)" print(f"序列长度: {seq_length} token") print(f"块数量: {num_blocks}") print(f"可缓存块数: {max_cacheable_blocks}") print(f"缓存策略: {strategy}") print(f"命中率: {hit_rate*100:.1f}%") return hit_rate # 测试不同序列长度 for length in [256, 1024, 4096, 16384, 65536]: simulate_shared_memory_cache_hit_rate(length) print()

通过本节的学习,我们深入理解了动态页表(块表)设计的核心技术和实现细节。从块表的基本数据结构到GPU端的高效查询,从调度器协同到张量并行下的同步机制,从内存开销分析到共享内存缓存优化,我们掌握了完整的块表管理体系。块表作为PagedAttention的"神经中枢",虽然结构简单,却在连接逻辑序列空间和物理显存空间中发挥着不可替代的作用,是理解vLLM推理引擎工作原理的关键所在。</arg_value>


作者与出处
来源:灏天文库
整理: 灏天文库整理
由灏天文库平台收录,内容或由平台用户上传,仅供学习交流
发布者: 作者: 内存溢出警告的小龙虾 转发
评论区 (0)
U