文章目录vLLM PagedAttention(vLLM 分页注意力)1. 痛点传统 KV Cache 的内存刺客2. 破局思路引入操作系统的虚拟内存3. PagedAttention 的工作机制4. PagedAttention 带来的核心红利SGLang RadixAttention(SGLang基数注意力)1. 核心痛点重复 Prefill 与跨请求浪费2. 核心原理用基数树Radix Tree自动管理 KV Cache3. 典型应用场景4. 对比PagedAttention vs. RadixAttentionvLLM PagedAttention(vLLM 分页注意力)在大语言模型LLM的部署和推理优化中vLLM 提出的PagedAttention分页注意力机制是一个里程碑式的突破。它直接解决了大模型在生成长文本时显存GPU Memory被严重浪费的核心痛点。以下是 PagedAttention 的核心原理及其实战机制的拆解1. 痛点传统 KV Cache 的内存刺客在自回归Autoregressive生成过程中大模型需要保存之前生成的每个 Token 的 Key 和 Value 张量以避免每次生成新词时重复计算。这些被缓存的数据就叫KV Cache。传统推理框架在管理 KV Cache 时面临严重的内存浪费问题预分配导致内部碎片系统通常会为了以防万一预先为每个请求分配一段连续的、达到模型最大允许长度的显存空间例如 2048 或 4096 tokens。但很多请求实际生成的长度远达不到最大值多出来的显存就被白白占用了。无法高效共享在并行采样如 Beam Search中同一个 Prompt 会产生多个不同的输出分支。传统的连续内存分配无法让这些分支共享 Prompt 阶段的 KV Cache导致一份前缀被复制了多份。据统计传统的连续内存管理机制会导致高达60% - 80%的 KV Cache 显存被浪费。2. 破局思路引入操作系统的虚拟内存vLLM 团队从操作系统的虚拟内存分页机制Virtual Memory Paging中汲取了灵感。在操作系统中程序认为自己拥有一块连续的“逻辑内存”但实际上这些内存在物理上是被切分成一个个固定大小的“页Pages”并可以分散存储在物理内存的各个角落。系统通过“页表Page Table”来记录逻辑地址到物理地址的映射关系。3. PagedAttention 的工作机制PagedAttention 将这种分页思想直接移植到了 LLM 的注意力机制计算中KV BlocksKV 块系统不再为每个请求分配一整段连续的 KV 显存而是将其划分为一个个固定大小的块Block。每个 Block 包含固定数量的 Token 的 KV 向量例如每个 Block 只装 16 个 Token。非连续物理存储这些 Block 在 GPU 的显存中不需要连续存放哪里有空闲空间就分配在哪里。Block Table块表vLLM 维护着一个中心化的块表负责将每个请求连续的“逻辑 KV 块”映射到显存中散落的“物理 KV 块”。当模型进行 Attention 计算时PagedAttention 算法会在底层自动查找块表跨越这些非连续的物理块准确地提取出需要的 KV 向量完成矩阵乘法。这个过程对上层模型是完全透明的。4. PagedAttention 带来的核心红利几乎消除显存浪费因为是按需动态分配用完一个 Block 再分配下一个显存的内部浪费被严格控制在最后一个没有装满的 Block 内通常浪费率低于 4%。吞吐量Throughput翻倍节省下来的海量显存可以用来同时容纳更多并发请求的 KV Cache从而显著增大 Batch Size。在相同的硬件下vLLM 的吞吐量通常能达到传统框架的2 - 4 倍。写时复制Copy-on-Write实现内存共享对于多个请求共享同一个 Prompt或者执行 Beam Search 这种有共同前缀的场景不同请求在 Block Table 中只需指向相同的物理块。只有当它们各自生成不同的后续 Token 时系统才会分配新的物理块写时复制。这使得复杂推理场景下的显存占用呈指数级下降。代码实现importtorchfromtypingimportListclassRequest:def__init__(self,request_id:int,prompt_len:int):self.request_idrequest_id self.seq_lenprompt_len self.block_table:List[int][]classKVCacheManager:def__init__(self,num_blocks:int,block_size:int,head_dim:int):self.num_blocksnum_blocks self.block_sizeblock_size self.head_dimhead_dim# TODO 1: 模拟预分配一块大显存池self.physical_kv_cachetorch.zeros(num_blocks,block_size,head_dim)# 跟踪哪些物理块被占用了self.free_blocks:List[int]list(range(num_blocks))defallocate_for_prefill(self,req:Request): 请求刚进来时 (Prefill阶段)为它的 Prompt 长度分配所需的全部 Block # TODO 2: 计算需要的 block 数量向上取整needed_blocks(req.seq_lenself.block_size-1)//self.block_size# TODO 3: 从 free_blocks 中弹出对应数量的 block 索引iflen(self.free_blocks)needed_blocks:raiseRuntimeError(OOM)for_inrange(needed_blocks):block_idself.free_blocks.pop(0)req.block_table.append(block_id)defallocate_for_decode(self,req:Request): 自回归生成时 (Decode阶段)检查序列长度。 如果当前最后一个 Block 满了则按需分配 1 个新 Block。 req.seq_len1# TODO 4: 判断是否需要新的 Blockis_new_block_needed(req.seq_len%self.block_size)1ifis_new_block_needed:ifnotself.free_blocks:raiseRuntimeError(OOM)block_idself.free_blocks.pop(0)req.block_table.append(block_id)defget_physical_cache(self,req:Request)-torch.Tensor: 根据块表把不连续的物理块拼凑成逻辑上连续的 KV Cache # TODO 5: 根据 req.block_table 的索引从物理池中提取对应的块blocks[self.physical_kv_cache[block_id]forblock_idinreq.block_table]cat_blockstorch.cat(blocks,dim0)# 只截取真实 seq_len 长度返回returncat_blocks[:req.seq_len]# 运行此单元格以测试你的实现deftest_paged_attention_manager():try:# Case 1: 典型 Prefill Decode Cache 拼装managerKVCacheManager(num_blocks10,block_size4,head_dim64)print(初始化内存池...)req1Request(request_id1,prompt_len6)manager.allocate_for_prefill(req1)assertlen(req1.block_table)2,长度 6 的请求应分配 2 个 Blockassertlen(manager.free_blocks)8,池中应该剩下 8 个空闲块print(f✅ Prefill 测试通过Req1 分配的块表:{req1.block_table})manager.allocate_for_decode(req1)assertlen(req1.block_table)2,生成第 7 个 token 时不应该分配新块manager.allocate_for_decode(req1)manager.allocate_for_decode(req1)assertlen(req1.block_table)3,生成第 9 个 token 时应当分配了第 3 个新块assertlen(manager.free_blocks)7,池中应该剩下 7 个空闲块print(f✅ Decode 动态分配测试通过Req1 最新块表:{req1.block_table})forblock_id,valueinzip(req1.block_table,[1.0,2.0,3.0]):manager.physical_kv_cache[block_id].fill_(value)cachemanager.get_physical_cache(req1)assertcache.shape(9,64),f拼装出来的连续 Cache 形状不对应为 (9, 64)实为{cache.shape}asserttorch.all(cache[:4]1.0),第 1 个 Block 未正确拼装asserttorch.all(cache[4:8]2.0),第 2 个 Block 未正确拼装asserttorch.all(cache[8:]3.0),第 3 个 Block 的截断拼装不正确print(✅ Cache 拼装测试通过多块物理缓存被正确恢复为逻辑连续序列。)# Case 2: 恰好跨越 block 边界时Decode 应该分配新块并正确截断最后一块manager2KVCacheManager(num_blocks4,block_size4,head_dim8)req2Request(request_id2,prompt_len4)manager2.allocate_for_prefill(req2)assertlen(req2.block_table)1,长度 4 的请求应只分配 1 个 Blockmanager2.allocate_for_decode(req2)assertlen(req2.block_table)2,长度 5 的请求应分配第 2 个 Blockmanager2.physical_kv_cache[req2.block_table[0]].fill_(7.0)manager2.physical_kv_cache[req2.block_table[1]].fill_(8.0)cache2manager2.get_physical_cache(req2)assertcache2.shape(5,8),f拼装出来的连续 Cache 形状不对应为 (5, 8)实为{cache2.shape}asserttorch.all(cache2[:4]7.0),边界块的前 4 个 token 不正确asserttorch.all(cache2[4:]8.0),边界块的最后 1 个 token 不正确print(✅ 边界分配与截断测试通过)# Case 3: OOM 分支必须抛出 RuntimeErroroom_managerKVCacheManager(num_blocks1,block_size4,head_dim8)oom_reqRequest(request_id3,prompt_len5)try:oom_manager.allocate_for_prefill(oom_req)exceptRuntimeErrorase:assertOOMinstr(e),OOM 异常信息不正确print(✅ OOM 测试通过)else:raiseAssertionError(显存池不足时应该抛出 RuntimeError(OOM))print(\n✅ All Tests Passed! PagedAttention 内存管理逻辑验证通过。)exceptNotImplementedError:print(请先完成 TODO 部分的代码)raiseexcept(AttributeError,NameError,TypeError,ValueError,AssertionError,RuntimeError)ase:ifisinstance(e,AttributeError):print(代码未完成无法找到必要的属性)elifisinstance(e,NameError):print(代码可能未完成导致变量为 NoneType。)elifisinstance(e,TypeError):print(代码可能未完成导致变量为 NoneType。)elifisinstance(e,ValueError):print(代码可能未完成导致了张量维度错误)elifisinstance(e,AssertionError):print(代码可能未完成导致了断言失败)elifisinstance(e,RuntimeError):print(代码可能未完成导致了运行时错误)else:print(代码可能未完成导致了断言失败)raiseNotImplementedError(请先完成 TODO 部分的代码)fromeexceptExceptionase:print(f❌ 测试失败:{e})raisetest_paged_attention_manager()结果SGLang RadixAttention(SGLang基数注意力)如果说 vLLM 的PagedAttention解决了“单请求内部Intra-request显存碎片化”的问题那么 SGLang 提出的RadixAttention基数注意力机制则更进一步解决了“跨请求/多轮交互间Inter-requestKV Cache 的自动复用与生命周期管理”的难题。1. 核心痛点重复 Prefill 与跨请求浪费在大模型实际应用场景中大量请求都包含重叠的前缀Overlapping Prefixes多轮对话Multi-turn Chat第 2 轮对话的输入包含第 1 轮的 Prompt 和 Answer。Agent / 复杂工作流在 Tree-of-Thought思维树或 Monte Carlo 搜索中多个分支共享相同的推导历史。固定 System Prompt / Few-shot 示例万级请求共享同一段很长的规则说明或示例。在传统推理引擎中哪怕请求之间有 90% 的 Token 完全相同新请求到来时系统依然要对其前缀重新做一次Prefill预填充计算不仅浪费算力还会导致首包延迟TTFT, Time-To-First-Token居高不下。2. 核心原理用基数树Radix Tree自动管理 KV CacheSGLang 没有采用复杂的全局 Hash 表而是引入了计算机科学中经典的Radix Tree基数树/压缩前缀树来作为 KV Cache 的索引结构。在 RadixTree 中边Edges与节点Nodes保存连续的 Token 序列如You are a helpful assistant...。指针与物理映射节点直接关联底层物理显存中的 KV Cache 块通常结合了类似 PagedAttention 的 Block 机制。动态生命周期自动匹配Prefix Matching新 Prompt 进来时在树中从根节点向下做最长前缀匹配。匹配到的部分直接复用 KV Cache彻底跳过这部分的 Prefill 计算模型只需要对“新后缀”进行计算。动态分裂Split Insert当新请求在某个节点中途出现分叉时原有节点会自动拆分为一个公共父节点和两个子节点。LRU 淘汰LRU Eviction当 GPU 显存满载时系统会根据 LRU最近最少使用策略优先删除树的叶子节点Leaf Nodes对应的 KV Cache并回收显存直到空间足够。3. 典型应用场景RadixAttention 的最大优势在于“零配置全自动”—— 开发者不需要手动维护复杂的 Cache 清单系统会在后台自动识别模式并完成 KV 共享。在以下四种常见模式中RadixAttention 能带来数倍的吞吐与延迟优化Few-shot Learning少样本提示多个请求共享相同的示例 Prompt前缀 Cache 命中率接近 100%。Multi-turn Chat多轮交互随对话轮数增加历史上下文全部在树上后续轮次只需 Prefill 用户刚发送的一句话。Self-consistency采样一致性单 Prompt 产生多个采样分支前缀只计算一次。Tree-of-Thought思维树搜索多路探索任务中所有子分支自动共享根节点与父节点的搜索历史。4. 对比PagedAttention vs. RadixAttention维度PagedAttention (vLLM)RadixAttention (SGLang)解决的核心问题解决单请求内物理显存离散化与碎片问题解决跨请求间自动前缀复用与调度问题索引数据结构扁平的物理块映射表Block Table动态层级基数树Radix Tree前缀匹配机制主要是单请求/显式配置的前缀缓存全自动最长前缀匹配Automatic Prefix Caching缓存回收策略请求结束即立即释放或简单保留基于树结构的LRU 延迟释放作为全局 Cache 池性能优势点极大提高 Batch Size 与 GPU 利用率极大地降低多轮/Agent 场景的TTFT (首字延迟)一句话总结PagedAttention 提供了高效的“物理内存物理块”管理而 RadixAttention 在其之上盖了一层“逻辑前缀索引树”两者结合成为了现代大模型推理引擎如 SGLang、vLLM v1/v2 架构的标准配置。代码实现importtorchclassTreeNode:def__init__(self,key_tokens):self.key_tokenskey_tokens#这条边上的Token序列如[101,532,789]self.children[]#子节点列表self.kv_cache_ptrNone#模拟指向物理KV Cache的指针classSimpleRadixCache:def__init__(self):#根节点是空的self.rootTreeNode([])definsert(self,tokens):nodeTreeNode(tokens)self.root.children.append(node)def_lcp_len(self,cached_tokens,prompt_tokens):match_len0#TODO1:逐个token计算最长公共前缀长度,遇到不相等时立刻停止match_len0whilematch_lenlen(cached_tokens)andmatch_lenlen(prompt_tokens):ifcached_tokens[match_len]prompt_tokens[match_len]:match_len1else:breakreturnmatch_lendefmatch_prefix(self,prompt_tokens):best_match_len0#TODO2:遍历self.root.children,更新最长匹配前缀长度forchildinself.root.children:match_lenself._lcp_len(child.key_tokens,prompt_tokens)ifmatch_lenbest_match_len:best_match_lenmatch_lenreturnbest_match_lendefsplit_prompt(self,prompt_tokens):#TODO3:先找命中长度再拆出前缀和后缀hit_lenself.match_prefix(prompt_tokens)hit_prefixprompt_tokens[:hit_len]miss_suffixprompt_tokens[hit_len:]returnhit_prefix,miss_suffix,hit_len# 测试你的实现deftest_radix_attention():try:cacheSimpleRadixCache()cache.insert([0,1,2,3])cache.insert([0,1,2,3,4])cache.insert([9,9,9])# 1. 基础 LCP 检查assertcache._lcp_len([1,2,3],[1,2,4])2,LCP 计算失败assertcache._lcp_len([7,8],[7,8,9,10])2,完整前缀匹配失败print(✅ 最长公共前缀计算正确)# 2. 多候选路径下应该选择最长命中前缀match_lencache.match_prefix([0,1,2,3,4,5])assertmatch_len5,匹配失败应该命中最长的 5 个 token 前缀。assertcache.match_prefix([7,6,5])0,错误匹配不该匹配到任何东西。print(✅ 多路径前缀命中选择正确)# 3. 前缀拆分验证hit_prefix,miss_suffix,hit_lencache.split_prompt([0,1,2,3,4,5])asserthit_len5,Hit Length 计算错误asserthit_prefix[0,1,2,3,4],可复用前缀拆分错误assertmiss_suffix[5],待重算后缀拆分错误hit_prefix2,miss_suffix2,hit_len2cache.split_prompt([7,6,5])asserthit_len20,无命中时 Hit Length 应为 0asserthit_prefix2[],无命中时前缀应为空assertmiss_suffix2[7,6,5],无命中时后缀应保持原样print(✅ 前缀拆分与回退逻辑正确)print(\n 所有测试通过这正是 SGLang 让大模型推理首字响应飞升 10 倍的底层秘密)exceptNotImplementedError:print(请先完成 TODO 部分的代码)raiseexcept(AttributeError,NameError,TypeError,ValueError,AssertionError,RuntimeError)ase:ifisinstance(e,AttributeError):print(代码未完成无法找到必要的属性)elifisinstance(e,NameError):print(代码可能未完成导致了变量未定义)elifisinstance(e,TypeError):print(代码可能未完成导致了操作错误)elifisinstance(e,ValueError):print(代码可能未完成导致了张量维度错误)elifisinstance(e,AssertionError):print(代码可能未完成导致了断言失败)elifisinstance(e,RuntimeError):print(代码可能未完成导致了运行时错误)else:print(代码可能未完成导致了断言失败)raiseNotImplementedError(请先完成 TODO 部分的代码)fromeexceptExceptionase:print(f❌ 发生未知异常:{e})raisetest_radix_attention()结果
DataWhale组队学习笔记--llm-algo-leetcode(五)
文章目录vLLM PagedAttention(vLLM 分页注意力)1. 痛点传统 KV Cache 的内存刺客2. 破局思路引入操作系统的虚拟内存3. PagedAttention 的工作机制4. PagedAttention 带来的核心红利SGLang RadixAttention(SGLang基数注意力)1. 核心痛点重复 Prefill 与跨请求浪费2. 核心原理用基数树Radix Tree自动管理 KV Cache3. 典型应用场景4. 对比PagedAttention vs. RadixAttentionvLLM PagedAttention(vLLM 分页注意力)在大语言模型LLM的部署和推理优化中vLLM 提出的PagedAttention分页注意力机制是一个里程碑式的突破。它直接解决了大模型在生成长文本时显存GPU Memory被严重浪费的核心痛点。以下是 PagedAttention 的核心原理及其实战机制的拆解1. 痛点传统 KV Cache 的内存刺客在自回归Autoregressive生成过程中大模型需要保存之前生成的每个 Token 的 Key 和 Value 张量以避免每次生成新词时重复计算。这些被缓存的数据就叫KV Cache。传统推理框架在管理 KV Cache 时面临严重的内存浪费问题预分配导致内部碎片系统通常会为了以防万一预先为每个请求分配一段连续的、达到模型最大允许长度的显存空间例如 2048 或 4096 tokens。但很多请求实际生成的长度远达不到最大值多出来的显存就被白白占用了。无法高效共享在并行采样如 Beam Search中同一个 Prompt 会产生多个不同的输出分支。传统的连续内存分配无法让这些分支共享 Prompt 阶段的 KV Cache导致一份前缀被复制了多份。据统计传统的连续内存管理机制会导致高达60% - 80%的 KV Cache 显存被浪费。2. 破局思路引入操作系统的虚拟内存vLLM 团队从操作系统的虚拟内存分页机制Virtual Memory Paging中汲取了灵感。在操作系统中程序认为自己拥有一块连续的“逻辑内存”但实际上这些内存在物理上是被切分成一个个固定大小的“页Pages”并可以分散存储在物理内存的各个角落。系统通过“页表Page Table”来记录逻辑地址到物理地址的映射关系。3. PagedAttention 的工作机制PagedAttention 将这种分页思想直接移植到了 LLM 的注意力机制计算中KV BlocksKV 块系统不再为每个请求分配一整段连续的 KV 显存而是将其划分为一个个固定大小的块Block。每个 Block 包含固定数量的 Token 的 KV 向量例如每个 Block 只装 16 个 Token。非连续物理存储这些 Block 在 GPU 的显存中不需要连续存放哪里有空闲空间就分配在哪里。Block Table块表vLLM 维护着一个中心化的块表负责将每个请求连续的“逻辑 KV 块”映射到显存中散落的“物理 KV 块”。当模型进行 Attention 计算时PagedAttention 算法会在底层自动查找块表跨越这些非连续的物理块准确地提取出需要的 KV 向量完成矩阵乘法。这个过程对上层模型是完全透明的。4. PagedAttention 带来的核心红利几乎消除显存浪费因为是按需动态分配用完一个 Block 再分配下一个显存的内部浪费被严格控制在最后一个没有装满的 Block 内通常浪费率低于 4%。吞吐量Throughput翻倍节省下来的海量显存可以用来同时容纳更多并发请求的 KV Cache从而显著增大 Batch Size。在相同的硬件下vLLM 的吞吐量通常能达到传统框架的2 - 4 倍。写时复制Copy-on-Write实现内存共享对于多个请求共享同一个 Prompt或者执行 Beam Search 这种有共同前缀的场景不同请求在 Block Table 中只需指向相同的物理块。只有当它们各自生成不同的后续 Token 时系统才会分配新的物理块写时复制。这使得复杂推理场景下的显存占用呈指数级下降。代码实现importtorchfromtypingimportListclassRequest:def__init__(self,request_id:int,prompt_len:int):self.request_idrequest_id self.seq_lenprompt_len self.block_table:List[int][]classKVCacheManager:def__init__(self,num_blocks:int,block_size:int,head_dim:int):self.num_blocksnum_blocks self.block_sizeblock_size self.head_dimhead_dim# TODO 1: 模拟预分配一块大显存池self.physical_kv_cachetorch.zeros(num_blocks,block_size,head_dim)# 跟踪哪些物理块被占用了self.free_blocks:List[int]list(range(num_blocks))defallocate_for_prefill(self,req:Request): 请求刚进来时 (Prefill阶段)为它的 Prompt 长度分配所需的全部 Block # TODO 2: 计算需要的 block 数量向上取整needed_blocks(req.seq_lenself.block_size-1)//self.block_size# TODO 3: 从 free_blocks 中弹出对应数量的 block 索引iflen(self.free_blocks)needed_blocks:raiseRuntimeError(OOM)for_inrange(needed_blocks):block_idself.free_blocks.pop(0)req.block_table.append(block_id)defallocate_for_decode(self,req:Request): 自回归生成时 (Decode阶段)检查序列长度。 如果当前最后一个 Block 满了则按需分配 1 个新 Block。 req.seq_len1# TODO 4: 判断是否需要新的 Blockis_new_block_needed(req.seq_len%self.block_size)1ifis_new_block_needed:ifnotself.free_blocks:raiseRuntimeError(OOM)block_idself.free_blocks.pop(0)req.block_table.append(block_id)defget_physical_cache(self,req:Request)-torch.Tensor: 根据块表把不连续的物理块拼凑成逻辑上连续的 KV Cache # TODO 5: 根据 req.block_table 的索引从物理池中提取对应的块blocks[self.physical_kv_cache[block_id]forblock_idinreq.block_table]cat_blockstorch.cat(blocks,dim0)# 只截取真实 seq_len 长度返回returncat_blocks[:req.seq_len]# 运行此单元格以测试你的实现deftest_paged_attention_manager():try:# Case 1: 典型 Prefill Decode Cache 拼装managerKVCacheManager(num_blocks10,block_size4,head_dim64)print(初始化内存池...)req1Request(request_id1,prompt_len6)manager.allocate_for_prefill(req1)assertlen(req1.block_table)2,长度 6 的请求应分配 2 个 Blockassertlen(manager.free_blocks)8,池中应该剩下 8 个空闲块print(f✅ Prefill 测试通过Req1 分配的块表:{req1.block_table})manager.allocate_for_decode(req1)assertlen(req1.block_table)2,生成第 7 个 token 时不应该分配新块manager.allocate_for_decode(req1)manager.allocate_for_decode(req1)assertlen(req1.block_table)3,生成第 9 个 token 时应当分配了第 3 个新块assertlen(manager.free_blocks)7,池中应该剩下 7 个空闲块print(f✅ Decode 动态分配测试通过Req1 最新块表:{req1.block_table})forblock_id,valueinzip(req1.block_table,[1.0,2.0,3.0]):manager.physical_kv_cache[block_id].fill_(value)cachemanager.get_physical_cache(req1)assertcache.shape(9,64),f拼装出来的连续 Cache 形状不对应为 (9, 64)实为{cache.shape}asserttorch.all(cache[:4]1.0),第 1 个 Block 未正确拼装asserttorch.all(cache[4:8]2.0),第 2 个 Block 未正确拼装asserttorch.all(cache[8:]3.0),第 3 个 Block 的截断拼装不正确print(✅ Cache 拼装测试通过多块物理缓存被正确恢复为逻辑连续序列。)# Case 2: 恰好跨越 block 边界时Decode 应该分配新块并正确截断最后一块manager2KVCacheManager(num_blocks4,block_size4,head_dim8)req2Request(request_id2,prompt_len4)manager2.allocate_for_prefill(req2)assertlen(req2.block_table)1,长度 4 的请求应只分配 1 个 Blockmanager2.allocate_for_decode(req2)assertlen(req2.block_table)2,长度 5 的请求应分配第 2 个 Blockmanager2.physical_kv_cache[req2.block_table[0]].fill_(7.0)manager2.physical_kv_cache[req2.block_table[1]].fill_(8.0)cache2manager2.get_physical_cache(req2)assertcache2.shape(5,8),f拼装出来的连续 Cache 形状不对应为 (5, 8)实为{cache2.shape}asserttorch.all(cache2[:4]7.0),边界块的前 4 个 token 不正确asserttorch.all(cache2[4:]8.0),边界块的最后 1 个 token 不正确print(✅ 边界分配与截断测试通过)# Case 3: OOM 分支必须抛出 RuntimeErroroom_managerKVCacheManager(num_blocks1,block_size4,head_dim8)oom_reqRequest(request_id3,prompt_len5)try:oom_manager.allocate_for_prefill(oom_req)exceptRuntimeErrorase:assertOOMinstr(e),OOM 异常信息不正确print(✅ OOM 测试通过)else:raiseAssertionError(显存池不足时应该抛出 RuntimeError(OOM))print(\n✅ All Tests Passed! PagedAttention 内存管理逻辑验证通过。)exceptNotImplementedError:print(请先完成 TODO 部分的代码)raiseexcept(AttributeError,NameError,TypeError,ValueError,AssertionError,RuntimeError)ase:ifisinstance(e,AttributeError):print(代码未完成无法找到必要的属性)elifisinstance(e,NameError):print(代码可能未完成导致变量为 NoneType。)elifisinstance(e,TypeError):print(代码可能未完成导致变量为 NoneType。)elifisinstance(e,ValueError):print(代码可能未完成导致了张量维度错误)elifisinstance(e,AssertionError):print(代码可能未完成导致了断言失败)elifisinstance(e,RuntimeError):print(代码可能未完成导致了运行时错误)else:print(代码可能未完成导致了断言失败)raiseNotImplementedError(请先完成 TODO 部分的代码)fromeexceptExceptionase:print(f❌ 测试失败:{e})raisetest_paged_attention_manager()结果SGLang RadixAttention(SGLang基数注意力)如果说 vLLM 的PagedAttention解决了“单请求内部Intra-request显存碎片化”的问题那么 SGLang 提出的RadixAttention基数注意力机制则更进一步解决了“跨请求/多轮交互间Inter-requestKV Cache 的自动复用与生命周期管理”的难题。1. 核心痛点重复 Prefill 与跨请求浪费在大模型实际应用场景中大量请求都包含重叠的前缀Overlapping Prefixes多轮对话Multi-turn Chat第 2 轮对话的输入包含第 1 轮的 Prompt 和 Answer。Agent / 复杂工作流在 Tree-of-Thought思维树或 Monte Carlo 搜索中多个分支共享相同的推导历史。固定 System Prompt / Few-shot 示例万级请求共享同一段很长的规则说明或示例。在传统推理引擎中哪怕请求之间有 90% 的 Token 完全相同新请求到来时系统依然要对其前缀重新做一次Prefill预填充计算不仅浪费算力还会导致首包延迟TTFT, Time-To-First-Token居高不下。2. 核心原理用基数树Radix Tree自动管理 KV CacheSGLang 没有采用复杂的全局 Hash 表而是引入了计算机科学中经典的Radix Tree基数树/压缩前缀树来作为 KV Cache 的索引结构。在 RadixTree 中边Edges与节点Nodes保存连续的 Token 序列如You are a helpful assistant...。指针与物理映射节点直接关联底层物理显存中的 KV Cache 块通常结合了类似 PagedAttention 的 Block 机制。动态生命周期自动匹配Prefix Matching新 Prompt 进来时在树中从根节点向下做最长前缀匹配。匹配到的部分直接复用 KV Cache彻底跳过这部分的 Prefill 计算模型只需要对“新后缀”进行计算。动态分裂Split Insert当新请求在某个节点中途出现分叉时原有节点会自动拆分为一个公共父节点和两个子节点。LRU 淘汰LRU Eviction当 GPU 显存满载时系统会根据 LRU最近最少使用策略优先删除树的叶子节点Leaf Nodes对应的 KV Cache并回收显存直到空间足够。3. 典型应用场景RadixAttention 的最大优势在于“零配置全自动”—— 开发者不需要手动维护复杂的 Cache 清单系统会在后台自动识别模式并完成 KV 共享。在以下四种常见模式中RadixAttention 能带来数倍的吞吐与延迟优化Few-shot Learning少样本提示多个请求共享相同的示例 Prompt前缀 Cache 命中率接近 100%。Multi-turn Chat多轮交互随对话轮数增加历史上下文全部在树上后续轮次只需 Prefill 用户刚发送的一句话。Self-consistency采样一致性单 Prompt 产生多个采样分支前缀只计算一次。Tree-of-Thought思维树搜索多路探索任务中所有子分支自动共享根节点与父节点的搜索历史。4. 对比PagedAttention vs. RadixAttention维度PagedAttention (vLLM)RadixAttention (SGLang)解决的核心问题解决单请求内物理显存离散化与碎片问题解决跨请求间自动前缀复用与调度问题索引数据结构扁平的物理块映射表Block Table动态层级基数树Radix Tree前缀匹配机制主要是单请求/显式配置的前缀缓存全自动最长前缀匹配Automatic Prefix Caching缓存回收策略请求结束即立即释放或简单保留基于树结构的LRU 延迟释放作为全局 Cache 池性能优势点极大提高 Batch Size 与 GPU 利用率极大地降低多轮/Agent 场景的TTFT (首字延迟)一句话总结PagedAttention 提供了高效的“物理内存物理块”管理而 RadixAttention 在其之上盖了一层“逻辑前缀索引树”两者结合成为了现代大模型推理引擎如 SGLang、vLLM v1/v2 架构的标准配置。代码实现importtorchclassTreeNode:def__init__(self,key_tokens):self.key_tokenskey_tokens#这条边上的Token序列如[101,532,789]self.children[]#子节点列表self.kv_cache_ptrNone#模拟指向物理KV Cache的指针classSimpleRadixCache:def__init__(self):#根节点是空的self.rootTreeNode([])definsert(self,tokens):nodeTreeNode(tokens)self.root.children.append(node)def_lcp_len(self,cached_tokens,prompt_tokens):match_len0#TODO1:逐个token计算最长公共前缀长度,遇到不相等时立刻停止match_len0whilematch_lenlen(cached_tokens)andmatch_lenlen(prompt_tokens):ifcached_tokens[match_len]prompt_tokens[match_len]:match_len1else:breakreturnmatch_lendefmatch_prefix(self,prompt_tokens):best_match_len0#TODO2:遍历self.root.children,更新最长匹配前缀长度forchildinself.root.children:match_lenself._lcp_len(child.key_tokens,prompt_tokens)ifmatch_lenbest_match_len:best_match_lenmatch_lenreturnbest_match_lendefsplit_prompt(self,prompt_tokens):#TODO3:先找命中长度再拆出前缀和后缀hit_lenself.match_prefix(prompt_tokens)hit_prefixprompt_tokens[:hit_len]miss_suffixprompt_tokens[hit_len:]returnhit_prefix,miss_suffix,hit_len# 测试你的实现deftest_radix_attention():try:cacheSimpleRadixCache()cache.insert([0,1,2,3])cache.insert([0,1,2,3,4])cache.insert([9,9,9])# 1. 基础 LCP 检查assertcache._lcp_len([1,2,3],[1,2,4])2,LCP 计算失败assertcache._lcp_len([7,8],[7,8,9,10])2,完整前缀匹配失败print(✅ 最长公共前缀计算正确)# 2. 多候选路径下应该选择最长命中前缀match_lencache.match_prefix([0,1,2,3,4,5])assertmatch_len5,匹配失败应该命中最长的 5 个 token 前缀。assertcache.match_prefix([7,6,5])0,错误匹配不该匹配到任何东西。print(✅ 多路径前缀命中选择正确)# 3. 前缀拆分验证hit_prefix,miss_suffix,hit_lencache.split_prompt([0,1,2,3,4,5])asserthit_len5,Hit Length 计算错误asserthit_prefix[0,1,2,3,4],可复用前缀拆分错误assertmiss_suffix[5],待重算后缀拆分错误hit_prefix2,miss_suffix2,hit_len2cache.split_prompt([7,6,5])asserthit_len20,无命中时 Hit Length 应为 0asserthit_prefix2[],无命中时前缀应为空assertmiss_suffix2[7,6,5],无命中时后缀应保持原样print(✅ 前缀拆分与回退逻辑正确)print(\n 所有测试通过这正是 SGLang 让大模型推理首字响应飞升 10 倍的底层秘密)exceptNotImplementedError:print(请先完成 TODO 部分的代码)raiseexcept(AttributeError,NameError,TypeError,ValueError,AssertionError,RuntimeError)ase:ifisinstance(e,AttributeError):print(代码未完成无法找到必要的属性)elifisinstance(e,NameError):print(代码可能未完成导致了变量未定义)elifisinstance(e,TypeError):print(代码可能未完成导致了操作错误)elifisinstance(e,ValueError):print(代码可能未完成导致了张量维度错误)elifisinstance(e,AssertionError):print(代码可能未完成导致了断言失败)elifisinstance(e,RuntimeError):print(代码可能未完成导致了运行时错误)else:print(代码可能未完成导致了断言失败)raiseNotImplementedError(请先完成 TODO 部分的代码)fromeexceptExceptionase:print(f❌ 发生未知异常:{e})raisetest_radix_attention()结果
