nano vllm是简化版的vllm的实现, 里面实现了大模型推理的常见优化. 包括page attention, kv cache, continuous batching.文件结构:E:\NANO-VLLM │ .gitignore │ bench.py │ example.py │ LICENSE │ pyproject.toml │ README.md │ ├───assets │ logo.png │ └───nanovllm │ config.py │ llm.py │ sampling_params.py │ __init__.py │ ├───engine │ block_manager.py │ llm_engine.py │ model_runner.py │ scheduler.py │ sequence.py │ ├───layers │ activation.py │ attention.py │ embed_head.py │ layernorm.py │ linear.py │ rotary_embedding.py │ sampler.py │ ├───models │ qwen3.py │ └───utils context.py loader.py1. 配置文件核心配置文件都在config.py中.llm.py是对外暴露接口的.sampling_params.py是采样的参数配置.engine文件夹定义了调度, runner, block相关的文件.layers是定义模型相关层的文件.models定义模型.utils定义了一些全局contex和加载模型权重的函数.dataclassclassConfig:model:strmax_num_batched_tokens:int16384max_num_seqs:int512max_model_len:int4096gpu_memory_utilization:float0.9tensor_parallel_size:int1enforce_eager:boolFalsehf_config:AutoConfig|NoneNoneeos:int-1kvcache_block_size:int256num_kvcache_blocks:int-1def__post_init__(self):assertos.path.isdir(self.model)assertself.kvcache_block_size%2560assert1self.tensor_parallel_size8self.hf_configAutoConfig.from_pretrained(self.model)self.max_model_lenmin(self.max_model_len,self.hf_config.max_position_embeddings)assertself.max_num_batched_tokensself.max_model_lenmax_num_batched_tokens: 一次请求的最大token数max_num_seqs: 一次请求的最大seq数目max_model_len: 一次请求的seq的最大长度tensor_parallel_size: GPU数目kvcache_block_size: kv cache的block size.num_kvcache_blocks:2. Sequencesequence.py: sequence abstraction用来跟踪一个请求prompt 生成内容的状态、token、以及 KV cache block 的映射关系。Sequence 表示 一个推理请求一个 prompt 它生成的 token。它负责管理当前生成到哪里了token_idsprompt 和生成部分的划分KV cache 对应的 blockblock_table当前状态WAITING / RUNNING / FINISHEDsampling 参数temperature 等classSequenceStatus(Enum):WAITINGauto()RUNNINGauto()FINISHEDauto()表示这个序列在调度器里的状态WAITING还没开始跑排队RUNNING正在 decodeFINISHED已经结束遇到 EOS 或 max_tokensself.seq_id next(Sequence.counter): 全局唯一 ID自增,self.token_ids copy(token_ids): 当前所有 tokenprompt generation,self.last_token token_ids[-1]: 上一个 tokendecode 用self.num_tokenslen(self.token_ids)self.num_prompt_tokenslen(token_ids)初始化时全部都是 prompt后面 append_token 才增加 completionself.num_cached_tokens0self.block_table[]block_table 映射 逻辑 block → 物理 KV cache blocknum_cached_tokens 已经算过 KV cache 的 token 数self.temperaturesampling_params.temperature self.max_tokenssampling_params.max_tokens self.ignore_eossampling_params.ignore_eos每个 sequence 可以有不同 decoding 策略defnum_completion_tokens(self):returnself.num_tokens-self.num_prompt_tokenscompletion token 数defprompt_token_ids(self):returnself.token_ids[:self.num_prompt_tokens]defcompletion_token_ids(self):returnself.token_ids[self.num_prompt_tokens:]prompt / completion 切分block_size256KV cache 按 block 管理类似 vLLM已缓存 block 数, 前多少 block 已经有 KV cachedef num_cached_blocks(self): return self.num_cached_tokens // self.block_size总共需要多少块才能覆盖已有的 token:ceil(num_tokens / block_size)def num_blocks(self): return (self.num_tokens self.block_size - 1) // self.block_sizedef last_block_num_tokens(self): return self.num_tokens - (self.num_blocks - 1) * self.block_size最后一个 block 的 token 数, 用来判断当前 block 是否满了是否需要分配新 blockproperty def num_blocks(self): # 总共需要多少块才能覆盖已有的 token return (self.num_tokens self.block_size - 1) // self.block_size property def last_block_num_tokens(self): return self.num_tokens - (self.num_blocks - 1) * self.block_size def block(self, i): # 第 i 个 block 的 token assert 0 i self.num_blocks return self.token_ids[i*self.block_size: (i1)*self.block_size]defappend_token(self,token_id:int):self.token_ids.append(token_id)self.last_tokentoken_id self.num_tokens1每生成一个 token加入 sequence更新 last_tokentoken 数 1这里只是逻辑 tokenKV cache 分配不在这里做通常在 scheduler / block managerSequence 只是记录用了哪些 blockblock_table哪些 token 已 cachenum_cached_tokenscountercount()全局变量, 记录当前已经完成了多少sequence.2. contextdataclassclassContext:is_prefill:boolFalsecu_seqlens_q:torch.Tensor|NoneNonecu_seqlens_k:torch.Tensor|NoneNonemax_seqlen_q:int0max_seqlen_k:int0slot_mapping:torch.Tensor|NoneNonecontext_lens:torch.Tensor|NoneNoneblock_tables:torch.Tensor|NoneNone_CONTEXTContext()defget_context():return_CONTEXTdefset_context(is_prefill,cu_seqlens_qNone,cu_seqlens_kNone,max_seqlen_q0,max_seqlen_k0,slot_mappingNone,context_lensNone,block_tablesNone):global_CONTEXT _CONTEXTContext(is_prefill,cu_seqlens_q,cu_seqlens_k,max_seqlen_q,max_seqlen_k,slot_mapping,context_lens,block_tables)defreset_context():global_CONTEXT _CONTEXTContext()单例模式, 创建一个全局的context, 发生的事情是Python 读取文件执行整个文件包括 _CONTEXT Context()把模块对象缓存到sys.modules[“my_context_module”]之后再在别的地方import my_context_module, Python 会直接从 sys.modules 里拿, 不会重新执行文件代码, 不会再次执行_CONTEXT Context()这里的context是为了设置一些全局变量, 防止某些变量需要通过函数传参层层传递的麻烦.3. Block和BlockManagerKV cache 分块block: 每block_size256tokens 一个 block, 避免连续内存分配类似分页prefix cache: 如果两个 sequence 有相同前缀seq1: A B C D | E F seq2: A B C D | X Y前 4 个 token 对应的 block 可以共享引用计数ref_count多个 sequence 可以共享一个 block用 ref_count 管理释放3.1 Block 类一个 KV cache 页classBlock:def__init__(self,block_id):self.block_idblock_id# 物理 block ID显存位置self.ref_count0# 被多少 sequence 使用self.hash-1# block 对应 token 的 hash用于 prefix cacheself.token_ids[]# 实际 token用于校验 hash 冲突defupdate(self,hash:int,token_ids:list[int]):# 当 block 写满后才更新 hashself.hashhashself.token_idstoken_idsdefreset(self):# 分配新 block 时调用self.ref_count1self.hash-1self.token_ids[]3.2. BlockManagerself.blocks:所有 block固定大小 pool self.free_block_ids:空闲 blockdeque self.used_block_ids:正在使用的 block self.hash_to_block_id:prefix cache,hash→ block_idhash 设计:compute_hash(token_ids,prefix):# 链式 hashrolling hashh_ihash(h_{i-1}当前block)3.2.1 allocate: 给一个 sequence 分配 blockdef allocate(self, seq: Sequence): 一个seq会被分成多个block, 比如1个seq_len10, block_size4, 那么这个seq会被分成442, 3个block. for 每个 block: 1. 计算 hash 2. 查 cachehash_to_block_id 3. 命中 → 共享 block 4. miss → 分配新 block关键变量h -1 cache_miss False一旦 miss后面全部 miss逐 block 解析step1取 token token_ids seq.block(i) step2算 hash h compute_hash(token_ids, h) if 满block else -1只有满 block 才参与 cachestep3查 cache block_id hash_to_block_id.get(h, -1) step4判断是否命中 if block_id -1 or token_ids 不一致: cache_miss True关键一旦 miss, 因为 prefix 已经不同了if cache_miss: 后面全部 block 都重新分配初始时:self.free_block_ids: deque[int] deque(range(num_blocks)), 从0~N.如果miss了:block_idself.free_block_ids[0]self.free_block_ids.remove(block_id)self.used_block_ids.add(block_id)如果不满block_size, 直接返回-1命中情况: 直接复用 KV cacheseq.num_cached_tokens block_size block.ref_count 1特殊情况if block_id in used_block_ids: block 正在用 → ref_count否则从 free list 重新分配但逻辑上是 reuse最后更新 cacheblock.update(h, token_ids) hash_to_block_id[h] block_id注意, allocate只在prefill阶段执行一次. 后面的decode都是在做append3.2.2 deallocate释放for block_id in reversed(seq.block_table): 从后往前释放类似栈 block.ref_count - 1如果没人用了if ref_count 0: 放回 free_block_ids注意不会删除 hash_to_block_id这是个优化点为什么从后往前释放?从后往前释放是为了符合“后分配先释放LIFO”的访问模式更安全、也更高效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)也就是最后一个 block → 倒数第二个 → … → 第一个 block。核心原因1符合“生成顺序”:Sequence 的 block 是这样增长的block0 → block1 → block2 → … → blockNdecode 时, 先用 block0prompt, 再用 block1, 最后用 blockN最新生成. 释放时最后的 block 最“新”、最可能不被共享. 如果反过来从前往后先处理 block0ref_count2, 再 block1, 逻辑上没错但不符合访问局部性.后面的block对前面的有依赖性, 因此应该先删除依赖最少的block, 这样如果后面的block可以被复用, 则前面的block肯定可以被复用. 相反, 如果前面的block可以被复用, 后面的Block不一定能复用cache locality 更好fragmentation 更少free list 更稳定3.2.3. appenddecode 核心block_tableseq.block_table last_blockself.blocks[block_table[-1]]拿到这个 sequence 的 block 列表找到“当前正在操作的最后一个 block”# 这个 sequence 用到的所有 block按顺序block_tableseq.block_table例如seq.block_table [3, 7, 2]对应block 3 → 前缀 block 7 → 中间 block 2 → 当前正在写block_table[-1]: 取最后一个元素, 当前 sequence 的“最后一个 block”[3,7,2]→2self.blocks[block_id]从全局 block pool 里取出这个 block 对象, 拿到真正的 KV cache blocklast_block self.blocks[block_table[-1]]意思是找到当前 sequence 正在写/刚写完的那个 block完整例子, 假设seq.block_table[3,7]block3:[A B C D]hashh1 block7:[E F _ _]hash-1执行这两行block_table seq.block_table last_block self.blocks[block_table[-1]]得到last_block block 7, 接下来会发生什么如果 append G[E F G _]走else:assertlast_block.hash-1什么都不做如果 append H[E F G H]走eliflen(seq)%block_size0:用的就是last_block.update(...)如果 append I新 block用的还是last_block.hash ! -1def can_append(self, seq): return len(self.free_block_ids) (len(seq) % self.block_size 1)只有在新开一个 block 时才需要分配。因为len(seq) % self.block_size 1成立的时候, 代表新的block开始.(len(seq) % self.block_size 1)的值为1.如果需要新 block → free block 是否 ≥ 1如果不需要 → 永远返回 True3.2.4. may_append重点每生成 一个新 tokendecode step 就会调用一次旧 seq 1 token → 调用 may_appendblock_table: 当前 sequence 用了哪些 block按顺序. 例如block_table [3, 7, 2]last_block: 当前正在写 or 刚写完的 block,block 的状态状态hash写入中-1已完成可复用 ! -1can_append为True才会执行may_append核心逻辑分三种情况情况1刚进入新 block, 触发时机iflen(seq)%block_size1:刚 append 一个 token并且新 token 是这个 block 的第一个 token例子block_size4[A B C D](满)append E →len5-%41成立assertlast_block.hash!-1上一个 block 必须已经是“完成态”block_idself.free_block_ids[0]拿一个空闲 blockself._allocate_block(block_id)将当前block_id标记为使用中通常会设置 ref_count1 等block_table.append(block_id)加到 sequence 的 block 列表里结果:新 block 开始写[E _ _ _]情况2刚好写满 blockeliflen(seq)%self.block_size0:# 该写入block了assertlast_block.hash-1# 上一个block还没写呢token_idsseq.block(seq.num_blocks-1)# 当前block的所有idsprefixself.blocks[block_table[-2]].hashiflen(block_table)1else-1# 上一个block的hash, 前缀hashhself.compute_hash(token_ids,prefix)# 根据当前的tokens和上一个的prefix hash计算当前block的hashlast_block.update(h,token_ids)# 更新hashself.hash_to_block_id[h]last_block.block_id触发时机: block 刚好被填满例子[E F G H](刚写完)assertlast_block.hash-1说明这个 block 之前是“写入中”, 也就是还没有被写入, 但是此时又是满的。因此后面我们需要写入该block, 并计算hash值。token_idsseq.block(seq.num_blocks-1)拿到这个 block 的完整 tokenprefix self.blocks[block_table[-2]].hash if len(block_table) 1 else -1前一个 block 的 hashprefixh self.compute_hash(token_ids, prefix)计算 prefix-aware hash, 本质hashhash(prefix当前 block)last_block.update(h,token_ids)更新 block, 保存 hash, 保存 token_idsself.hash_to_block_id[h]last_block.block_id放进 prefix cache结果:这个 block 从写入中 → 可复用 cachehcompute_hash(...)last_block.update(h,token_ids)hash_to_block_id[h]block_id这一步才让 block 可被复用情况3block 中间,else, 什么都不做核心设计:只有“满 block”才参与 cache, 避免 partial block 混乱rolling hashprefix-aware, 保证 prefix 一致才能共享一旦 miss后面全 miss, 因为 prefix 不一样了ref_count 实现共享, 多 sequence 共用 KV cacheappend 时才 finalize block, decode 阶段逐步构建 cache一个seqence来了之后开始不断建立kv cache, 每256个token分配一个kv cache block, 不足就不分配。当前sequence生成结束后, 对应的block的ref减1.注意: 只有当ref为0的时候, 对应的block才会被释放。kv cache block的服用发生在并发请求的时候。并发请求时间线t1: seq1: A B C D E F (正在生成) t2: seq2: A B C D X Y (新请求进来)allocate seq2 时block0: [A B C D] → 命中 seq1 的 block, 复用发生了block1: miss关键点seq1 还没结束 → block 还在 → 可以复用现实中并发请求非常多最关键比如1000 个用户同时问Write a Python function to..., 前缀一样 所有人共享 KV cachestreaming / batching, 调度器会这样做step1: 收集多个请求 step2: 一起 prefill step3: 一起 decode同时存在 → 可复用beam search / sampling: 一个 prompt → 多个分支, 前缀完全一样prompt → 多条生成路径, KV cache 直接共享
nano vllm代码详解
nano vllm是简化版的vllm的实现, 里面实现了大模型推理的常见优化. 包括page attention, kv cache, continuous batching.文件结构:E:\NANO-VLLM │ .gitignore │ bench.py │ example.py │ LICENSE │ pyproject.toml │ README.md │ ├───assets │ logo.png │ └───nanovllm │ config.py │ llm.py │ sampling_params.py │ __init__.py │ ├───engine │ block_manager.py │ llm_engine.py │ model_runner.py │ scheduler.py │ sequence.py │ ├───layers │ activation.py │ attention.py │ embed_head.py │ layernorm.py │ linear.py │ rotary_embedding.py │ sampler.py │ ├───models │ qwen3.py │ └───utils context.py loader.py1. 配置文件核心配置文件都在config.py中.llm.py是对外暴露接口的.sampling_params.py是采样的参数配置.engine文件夹定义了调度, runner, block相关的文件.layers是定义模型相关层的文件.models定义模型.utils定义了一些全局contex和加载模型权重的函数.dataclassclassConfig:model:strmax_num_batched_tokens:int16384max_num_seqs:int512max_model_len:int4096gpu_memory_utilization:float0.9tensor_parallel_size:int1enforce_eager:boolFalsehf_config:AutoConfig|NoneNoneeos:int-1kvcache_block_size:int256num_kvcache_blocks:int-1def__post_init__(self):assertos.path.isdir(self.model)assertself.kvcache_block_size%2560assert1self.tensor_parallel_size8self.hf_configAutoConfig.from_pretrained(self.model)self.max_model_lenmin(self.max_model_len,self.hf_config.max_position_embeddings)assertself.max_num_batched_tokensself.max_model_lenmax_num_batched_tokens: 一次请求的最大token数max_num_seqs: 一次请求的最大seq数目max_model_len: 一次请求的seq的最大长度tensor_parallel_size: GPU数目kvcache_block_size: kv cache的block size.num_kvcache_blocks:2. Sequencesequence.py: sequence abstraction用来跟踪一个请求prompt 生成内容的状态、token、以及 KV cache block 的映射关系。Sequence 表示 一个推理请求一个 prompt 它生成的 token。它负责管理当前生成到哪里了token_idsprompt 和生成部分的划分KV cache 对应的 blockblock_table当前状态WAITING / RUNNING / FINISHEDsampling 参数temperature 等classSequenceStatus(Enum):WAITINGauto()RUNNINGauto()FINISHEDauto()表示这个序列在调度器里的状态WAITING还没开始跑排队RUNNING正在 decodeFINISHED已经结束遇到 EOS 或 max_tokensself.seq_id next(Sequence.counter): 全局唯一 ID自增,self.token_ids copy(token_ids): 当前所有 tokenprompt generation,self.last_token token_ids[-1]: 上一个 tokendecode 用self.num_tokenslen(self.token_ids)self.num_prompt_tokenslen(token_ids)初始化时全部都是 prompt后面 append_token 才增加 completionself.num_cached_tokens0self.block_table[]block_table 映射 逻辑 block → 物理 KV cache blocknum_cached_tokens 已经算过 KV cache 的 token 数self.temperaturesampling_params.temperature self.max_tokenssampling_params.max_tokens self.ignore_eossampling_params.ignore_eos每个 sequence 可以有不同 decoding 策略defnum_completion_tokens(self):returnself.num_tokens-self.num_prompt_tokenscompletion token 数defprompt_token_ids(self):returnself.token_ids[:self.num_prompt_tokens]defcompletion_token_ids(self):returnself.token_ids[self.num_prompt_tokens:]prompt / completion 切分block_size256KV cache 按 block 管理类似 vLLM已缓存 block 数, 前多少 block 已经有 KV cachedef num_cached_blocks(self): return self.num_cached_tokens // self.block_size总共需要多少块才能覆盖已有的 token:ceil(num_tokens / block_size)def num_blocks(self): return (self.num_tokens self.block_size - 1) // self.block_sizedef last_block_num_tokens(self): return self.num_tokens - (self.num_blocks - 1) * self.block_size最后一个 block 的 token 数, 用来判断当前 block 是否满了是否需要分配新 blockproperty def num_blocks(self): # 总共需要多少块才能覆盖已有的 token return (self.num_tokens self.block_size - 1) // self.block_size property def last_block_num_tokens(self): return self.num_tokens - (self.num_blocks - 1) * self.block_size def block(self, i): # 第 i 个 block 的 token assert 0 i self.num_blocks return self.token_ids[i*self.block_size: (i1)*self.block_size]defappend_token(self,token_id:int):self.token_ids.append(token_id)self.last_tokentoken_id self.num_tokens1每生成一个 token加入 sequence更新 last_tokentoken 数 1这里只是逻辑 tokenKV cache 分配不在这里做通常在 scheduler / block managerSequence 只是记录用了哪些 blockblock_table哪些 token 已 cachenum_cached_tokenscountercount()全局变量, 记录当前已经完成了多少sequence.2. contextdataclassclassContext:is_prefill:boolFalsecu_seqlens_q:torch.Tensor|NoneNonecu_seqlens_k:torch.Tensor|NoneNonemax_seqlen_q:int0max_seqlen_k:int0slot_mapping:torch.Tensor|NoneNonecontext_lens:torch.Tensor|NoneNoneblock_tables:torch.Tensor|NoneNone_CONTEXTContext()defget_context():return_CONTEXTdefset_context(is_prefill,cu_seqlens_qNone,cu_seqlens_kNone,max_seqlen_q0,max_seqlen_k0,slot_mappingNone,context_lensNone,block_tablesNone):global_CONTEXT _CONTEXTContext(is_prefill,cu_seqlens_q,cu_seqlens_k,max_seqlen_q,max_seqlen_k,slot_mapping,context_lens,block_tables)defreset_context():global_CONTEXT _CONTEXTContext()单例模式, 创建一个全局的context, 发生的事情是Python 读取文件执行整个文件包括 _CONTEXT Context()把模块对象缓存到sys.modules[“my_context_module”]之后再在别的地方import my_context_module, Python 会直接从 sys.modules 里拿, 不会重新执行文件代码, 不会再次执行_CONTEXT Context()这里的context是为了设置一些全局变量, 防止某些变量需要通过函数传参层层传递的麻烦.3. Block和BlockManagerKV cache 分块block: 每block_size256tokens 一个 block, 避免连续内存分配类似分页prefix cache: 如果两个 sequence 有相同前缀seq1: A B C D | E F seq2: A B C D | X Y前 4 个 token 对应的 block 可以共享引用计数ref_count多个 sequence 可以共享一个 block用 ref_count 管理释放3.1 Block 类一个 KV cache 页classBlock:def__init__(self,block_id):self.block_idblock_id# 物理 block ID显存位置self.ref_count0# 被多少 sequence 使用self.hash-1# block 对应 token 的 hash用于 prefix cacheself.token_ids[]# 实际 token用于校验 hash 冲突defupdate(self,hash:int,token_ids:list[int]):# 当 block 写满后才更新 hashself.hashhashself.token_idstoken_idsdefreset(self):# 分配新 block 时调用self.ref_count1self.hash-1self.token_ids[]3.2. BlockManagerself.blocks:所有 block固定大小 pool self.free_block_ids:空闲 blockdeque self.used_block_ids:正在使用的 block self.hash_to_block_id:prefix cache,hash→ block_idhash 设计:compute_hash(token_ids,prefix):# 链式 hashrolling hashh_ihash(h_{i-1}当前block)3.2.1 allocate: 给一个 sequence 分配 blockdef allocate(self, seq: Sequence): 一个seq会被分成多个block, 比如1个seq_len10, block_size4, 那么这个seq会被分成442, 3个block. for 每个 block: 1. 计算 hash 2. 查 cachehash_to_block_id 3. 命中 → 共享 block 4. miss → 分配新 block关键变量h -1 cache_miss False一旦 miss后面全部 miss逐 block 解析step1取 token token_ids seq.block(i) step2算 hash h compute_hash(token_ids, h) if 满block else -1只有满 block 才参与 cachestep3查 cache block_id hash_to_block_id.get(h, -1) step4判断是否命中 if block_id -1 or token_ids 不一致: cache_miss True关键一旦 miss, 因为 prefix 已经不同了if cache_miss: 后面全部 block 都重新分配初始时:self.free_block_ids: deque[int] deque(range(num_blocks)), 从0~N.如果miss了:block_idself.free_block_ids[0]self.free_block_ids.remove(block_id)self.used_block_ids.add(block_id)如果不满block_size, 直接返回-1命中情况: 直接复用 KV cacheseq.num_cached_tokens block_size block.ref_count 1特殊情况if block_id in used_block_ids: block 正在用 → ref_count否则从 free list 重新分配但逻辑上是 reuse最后更新 cacheblock.update(h, token_ids) hash_to_block_id[h] block_id注意, allocate只在prefill阶段执行一次. 后面的decode都是在做append3.2.2 deallocate释放for block_id in reversed(seq.block_table): 从后往前释放类似栈 block.ref_count - 1如果没人用了if ref_count 0: 放回 free_block_ids注意不会删除 hash_to_block_id这是个优化点为什么从后往前释放?从后往前释放是为了符合“后分配先释放LIFO”的访问模式更安全、也更高效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)也就是最后一个 block → 倒数第二个 → … → 第一个 block。核心原因1符合“生成顺序”:Sequence 的 block 是这样增长的block0 → block1 → block2 → … → blockNdecode 时, 先用 block0prompt, 再用 block1, 最后用 blockN最新生成. 释放时最后的 block 最“新”、最可能不被共享. 如果反过来从前往后先处理 block0ref_count2, 再 block1, 逻辑上没错但不符合访问局部性.后面的block对前面的有依赖性, 因此应该先删除依赖最少的block, 这样如果后面的block可以被复用, 则前面的block肯定可以被复用. 相反, 如果前面的block可以被复用, 后面的Block不一定能复用cache locality 更好fragmentation 更少free list 更稳定3.2.3. appenddecode 核心block_tableseq.block_table last_blockself.blocks[block_table[-1]]拿到这个 sequence 的 block 列表找到“当前正在操作的最后一个 block”# 这个 sequence 用到的所有 block按顺序block_tableseq.block_table例如seq.block_table [3, 7, 2]对应block 3 → 前缀 block 7 → 中间 block 2 → 当前正在写block_table[-1]: 取最后一个元素, 当前 sequence 的“最后一个 block”[3,7,2]→2self.blocks[block_id]从全局 block pool 里取出这个 block 对象, 拿到真正的 KV cache blocklast_block self.blocks[block_table[-1]]意思是找到当前 sequence 正在写/刚写完的那个 block完整例子, 假设seq.block_table[3,7]block3:[A B C D]hashh1 block7:[E F _ _]hash-1执行这两行block_table seq.block_table last_block self.blocks[block_table[-1]]得到last_block block 7, 接下来会发生什么如果 append G[E F G _]走else:assertlast_block.hash-1什么都不做如果 append H[E F G H]走eliflen(seq)%block_size0:用的就是last_block.update(...)如果 append I新 block用的还是last_block.hash ! -1def can_append(self, seq): return len(self.free_block_ids) (len(seq) % self.block_size 1)只有在新开一个 block 时才需要分配。因为len(seq) % self.block_size 1成立的时候, 代表新的block开始.(len(seq) % self.block_size 1)的值为1.如果需要新 block → free block 是否 ≥ 1如果不需要 → 永远返回 True3.2.4. may_append重点每生成 一个新 tokendecode step 就会调用一次旧 seq 1 token → 调用 may_appendblock_table: 当前 sequence 用了哪些 block按顺序. 例如block_table [3, 7, 2]last_block: 当前正在写 or 刚写完的 block,block 的状态状态hash写入中-1已完成可复用 ! -1can_append为True才会执行may_append核心逻辑分三种情况情况1刚进入新 block, 触发时机iflen(seq)%block_size1:刚 append 一个 token并且新 token 是这个 block 的第一个 token例子block_size4[A B C D](满)append E →len5-%41成立assertlast_block.hash!-1上一个 block 必须已经是“完成态”block_idself.free_block_ids[0]拿一个空闲 blockself._allocate_block(block_id)将当前block_id标记为使用中通常会设置 ref_count1 等block_table.append(block_id)加到 sequence 的 block 列表里结果:新 block 开始写[E _ _ _]情况2刚好写满 blockeliflen(seq)%self.block_size0:# 该写入block了assertlast_block.hash-1# 上一个block还没写呢token_idsseq.block(seq.num_blocks-1)# 当前block的所有idsprefixself.blocks[block_table[-2]].hashiflen(block_table)1else-1# 上一个block的hash, 前缀hashhself.compute_hash(token_ids,prefix)# 根据当前的tokens和上一个的prefix hash计算当前block的hashlast_block.update(h,token_ids)# 更新hashself.hash_to_block_id[h]last_block.block_id触发时机: block 刚好被填满例子[E F G H](刚写完)assertlast_block.hash-1说明这个 block 之前是“写入中”, 也就是还没有被写入, 但是此时又是满的。因此后面我们需要写入该block, 并计算hash值。token_idsseq.block(seq.num_blocks-1)拿到这个 block 的完整 tokenprefix self.blocks[block_table[-2]].hash if len(block_table) 1 else -1前一个 block 的 hashprefixh self.compute_hash(token_ids, prefix)计算 prefix-aware hash, 本质hashhash(prefix当前 block)last_block.update(h,token_ids)更新 block, 保存 hash, 保存 token_idsself.hash_to_block_id[h]last_block.block_id放进 prefix cache结果:这个 block 从写入中 → 可复用 cachehcompute_hash(...)last_block.update(h,token_ids)hash_to_block_id[h]block_id这一步才让 block 可被复用情况3block 中间,else, 什么都不做核心设计:只有“满 block”才参与 cache, 避免 partial block 混乱rolling hashprefix-aware, 保证 prefix 一致才能共享一旦 miss后面全 miss, 因为 prefix 不一样了ref_count 实现共享, 多 sequence 共用 KV cacheappend 时才 finalize block, decode 阶段逐步构建 cache一个seqence来了之后开始不断建立kv cache, 每256个token分配一个kv cache block, 不足就不分配。当前sequence生成结束后, 对应的block的ref减1.注意: 只有当ref为0的时候, 对应的block才会被释放。kv cache block的服用发生在并发请求的时候。并发请求时间线t1: seq1: A B C D E F (正在生成) t2: seq2: A B C D X Y (新请求进来)allocate seq2 时block0: [A B C D] → 命中 seq1 的 block, 复用发生了block1: miss关键点seq1 还没结束 → block 还在 → 可以复用现实中并发请求非常多最关键比如1000 个用户同时问Write a Python function to..., 前缀一样 所有人共享 KV cachestreaming / batching, 调度器会这样做step1: 收集多个请求 step2: 一起 prefill step3: 一起 decode同时存在 → 可复用beam search / sampling: 一个 prompt → 多个分支, 前缀完全一样prompt → 多条生成路径, KV cache 直接共享
