-block_manager)
当前笔记顺序Engine(当前block_manager.py)-LayersEngine的运行架构(总结时或者弄不清一些调用顺序时看)block_manager.py作用解析管理kv_cache Block实现Block的分配、释放、复用这里需要理解一个句子能利用到的kv_cache Block需要是相同链计算出来的两个Block即使存的token一致也不一定能复用这是本文件的核心逻辑比如下图中的Block苹果它来自对不同前缀的计算。from collections import deque #给block生成唯一的hash用于快速查找和比对 import xxhash import numpy as np from nanovllm.engine.sequence import Sequence class Block: def __init__(self, block_id): self.block_id block_id #记录该块被多少个sequence使用只有运行态的sequence才会贡献ref_count self.ref_count 0 self.hash -1 self.token_ids [] def update(self, hash: int, token_ids: list[int]): self.hash hash self.token_ids token_ids def reset(self): #reset设为1而不是0的原因见问题1 self.ref_count 1 self.hash -1 self.token_ids [] class BlockManager: def __init__(self, num_blocks: int, block_size: int): self.block_size block_size #按递增初始化一个block资源表 self.blocks: list[Block] [Block(i) for i in range(num_blocks)] #用于实现「前缀缓存」相同 token 序列复用 block self.hash_to_block_id: dict[int, int] dict() #以队列形式记录空闲的block有哪些 self.free_block_ids: deque[int] deque(range(num_blocks)) #用过的block必不能重复 self.used_block_ids: set[int] set() classmethod def compute_hash(cls, token_ids: list[int], prefix: int -1): #h不是数字,是一个有状态的对象初始内部状态完全一样 h xxhash.xxh64() #如果不是第一个块,就把前一个块的哈希值(64位整数)以小端序转成8字节,塞进当前哈希器中 #这是构造链式哈希(rolling hash),使得每个块的哈希不仅依赖自己的 token,还依赖整个前缀链。这样做可以避免内容一样位置不同的块使用相同HASH if prefix ! -1: h.update(prefix.to_bytes(8, little)) #此处看can_allocate的逻辑是把一个块的token进行转换而不是输入的全部token h.update(np.array(token_ids).tobytes()) #返回int类型的64位摘要 return h.intdigest() #从空闲池取一个 Block 并重置 def _allocate_block(self) - int: block_id self.free_block_ids.popleft() block self.blocks[block_id] assert block.ref_count 0 #此处逻辑见问题2 if block.hash ! -1 and self.hash_to_block_id.get(block.hash) block_id: del self.hash_to_block_id[block.hash] block.reset() self.used_block_ids.add(block_id) return block_id #把一个 Block 放回空闲池只设置引用数为0不清空token和hash以便再次利用 def _deallocate_block(self, block_id: int): #当释放block时必须确保没有sequence在使用它 assert self.blocks[block_id].ref_count 0 self.used_block_ids.remove(block_id) self.free_block_ids.append(block_id) #Prefill-问:这个序列我能缓存多少 Block?。返回已经缓存的kv_cache block数量这些块里的token不用重新算也不用分配新的block直接让seq.block_table引用即可 def can_allocate(self, seq: Sequence) - int: #第一个块没有前缀因此h-1 h -1 num_cached_blocks 0 #num_blocks是实时计算的 num_new_blocks seq.num_blocks #最后一个 Block 永远是正在填的状态即使空的也会预分配因此不会缓存命中需要跳过 for i in range(seq.num_blocks - 1): token_ids seq.block(i) h self.compute_hash(token_ids, h) block_id self.hash_to_block_id.get(h, -1) #必须缓存命中且内容一致如果不一致就不需要继续因为kv_cache的利用必须前面的所有token一致见transformer原理 if block_id -1 or self.blocks[block_id].token_ids ! token_ids: break num_cached_blocks 1 if block_id in self.used_block_ids: num_new_blocks - 1 #如果需要分配的新block数不能够满足就return if len(self.free_block_ids) num_new_blocks: return -1 return num_cached_blocks #Prefill-把上面函数返回的N个缓存的Block拿来,其余新分配 def allocate(self, seq: Sequence, num_cached_blocks: int): assert not seq.block_table h -1 for i in range(num_cached_blocks): token_ids seq.block(i) h self.compute_hash(token_ids, h) block_id self.hash_to_block_id[h] block self.blocks[block_id] if block_id in self.used_block_ids: block.ref_count 1 else: block.ref_count 1 self.free_block_ids.remove(block_id) self.used_block_ids.add(block_id) seq.block_table.append(block_id) for i in range(num_cached_blocks, seq.num_blocks): seq.block_table.append(self._allocate_block()) seq.num_cached_tokens num_cached_blocks * self.block_size #还回序列的所有token此函数在序列到完成态再执行 def deallocate(self, seq: Sequence): for block_id in reversed(seq.block_table): block self.blocks[block_id] block.ref_count - 1 if block.ref_count 0: self._deallocate_block(block_id) seq.num_cached_tokens 0 seq.block_table.clear() #Decode-此处是问再decode一个token时,要不要新的Block? def can_append(self, seq: Sequence) - bool: #此处逻辑有点难由于token是一个个生成的我们需要新的Block时肯定是当我们token数为n *block_size 1的时候 return len(self.free_block_ids) (len(seq) % self.block_size 1) #Decode-答:需要就分配,不需要就不动 def may_append(self, seq: Sequence): if len(seq) % self.block_size 1: seq.block_table.append(self._allocate_block()) #把Block的hash写进表里已经缓存的肯定写过了毕竟都是靠hash查的缓存这里是把新分配的Block写入发生在seq调度后 def hash_blocks(self, seq: Sequence): start seq.num_cached_tokens // self.block_size end (seq.num_cached_tokens seq.num_scheduled_tokens) // self.block_size if start end: return h self.blocks[seq.block_table[start - 1]].hash if start 0 else -1 for i in range(start, end): block self.blocks[seq.block_table[i]] token_ids seq.block(i) h self.compute_hash(token_ids, h) block.update(h, token_ids) self.hash_to_block_id[h] block.block_id一些问题1.reset为什么把self.ref_count 1而不是0。 reset 仅在「分配 block」时调用见 _allocate_block 方法此时该 block 已经被某个 sequence 占用引用计数初始化为 1 是合理的。 2.if block.hash ! -1 and self.hash_to_block_id.get(block.hash) block_id: del self.hash_to_block_id[block.hash]此处逻辑是否重复 简而言之我要重置自己啦,但如果全局哈希表还指向我,就把那条记录删掉;如果指向别人,就千万别动,免得把别人的缓存给废了。 此处背景是Block不会被销毁,而是回收复用即使引用计数归0也是保留缓存了的token_ids和hash放入 free_block_ids空闲池。如果有新内容要放进这个块需要清空self.hash_to_block_id[block.hash]避免下次有人用相同token_ids查询缓存时被误判缓存命中。这是左边的逻辑 右边的逻辑是为了处理多个block有相同hash的情况比如 Block 5空闲中,内容 [我爱学习],hash12345 hash_to_block_id {12345: 5} 后来 Block 5 被释放,Block 7 也释放了同样的内容,但... 实际上:Block 7 先被分配、并被某个活跃序列使用: hash_to_block_id {12345: 7} ← 7 覆盖了 5(因为 7 现在是持有者) → 现在轮到分配 block 5 给另一个新序列 → 检查 hash_to_block_id[12345] 5 ? ✗ (现在是 7,不是 5) → 不删除! ← 关键!因为如果删了,block 7 的缓存就丢了 3.classmethod和staticmethod有什么区别。 特性 classmethod staticmethod 第一个参数 必须是类本身通常命名为 cls自动传参无需调用时传入 无强制参数和普通函数一样 访问类 / 实例属性 可通过 cls 访问类属性、创建实例 无法直接访问类 / 实例属性需显式传参 用途 类级别的操作如工厂方法、通用计算 工具函数逻辑属于类但无类 / 实例依赖 4.h self.compute_hash(token_ids, h) if len(token_ids) self.block_size else -1是什么意思compute_hash不是有cls参数吗为什么他调用的时候没填写。 见问题3 5.block_id self.hash_to_block_id.get(h, -1)字典类型他为什么传入了(h,-1)不应该是单参数吗 dict.get(key, default) 是 Python 字典的内置方法不是单参数 第一个参数 key要查找的键这里是哈希值 h 第二个参数 default可选若键不存在时返回的默认值这里是 - 1 代码中 block_id self.hash_to_block_id.get(h, -1) 的含义 如果哈希 h 存在于字典中返回对应的 block_id 如果不存在返回 - 1表示未命中缓存 对比 dict[key]get 不会因键不存在抛异常更安全。 6.关于「不满的 block 是否会缓存命中」 结论不可能。代码里有严格的逻辑限制不满的 block 不会被缓存自然也不会命中。 7.can_allocate vs allocate (Prefill 阶段)同can_append vs may_append 这是经典的 两阶段提交(two-phase): 第一阶段:只读不写(can_allocate)、无副作用,可以安全地试探 第二阶段:拿到第一阶段的结果再真正动手(allocate) 为什么要分两步?因为显存可能不够!调度器要先问能不能,再决定要不要做。如果 can_allocate 返回 -1,就不该调用 allocate。 8.整体顺序 can_allocate - 算出已有的缓存Block - allocate - 把block缓存了的直接写没缓存的分配再写写入seq.block_table - can_append - 显存够不够 - may_append - 可能调用_allocate_block - 生成结束调用deallocate - hash_blocks本系列文章(待写完修正)[系列(1) 开篇](https://blog.csdn.net/xxx/article/details/xxxxxx)[系列(2) 核心原理](https://blog.csdn.net/xxx/article/details/xxxxxx)上一篇[系列(1)开篇](https://blog.csdn.net/xxx/article/details/xxxxxx)下一篇[系列(3)实战演练](https://blog.csdn.net/xxx/article/details/xxxxxx)