Nano-VLLM全代码解析笔记(4)-model_runner

当前笔记顺序Engine(当前model_runner.py)-LayersEngine的运行架构(总结时或者弄不清一些调用顺序时看)model_runner.py最难的一个要理清kv_cache的逻辑与物理分配的关系建议配合总结的图一起看另外这里很多设计会涉及到flash_attention的实现比如cu_seqlen_q/k会在解析里和后面的attention.py讲解作用解析封装模型的执行逻辑负责把调度的 seqs 变成实际的 forward 计算。同时这里会涉及kv_cache的分配需要注意block_table存的是物理块号给 attention 内核读历史 KV 用slot_mapping存的是物理槽位给set_kv_buffer写新 KV 用分配相关逻辑如图import pickle import torch #PyTorch 官方的分布式通信库用于实现多进程 / 多节点间的张量通信、进程同步是大模型张量并行TP/ 数据并行DP的核心依赖支持多种通信后端ncclGPU 间高速通信推荐、glooCPU 通信、mpi import torch.distributed as dist #Python 多进程的同步原语信号量用于进程间的 “信号通知”核心方法 #set()将事件置为 “就绪” 状态 #clear()将事件置为 “未就绪” 状态 #wait()阻塞当前进程直到事件被置为就绪 from multiprocessing.synchronize import Event #Python 多进程的共享内存机制允许多个进程直接访问同一块内存区域无需数据拷贝 #核心优势比管道 / 队列的进程间通信效率高尤其大张量 #代码中用途主进程将推理命令方法名 参数序列化后写入共享内存子进程直接读取避免张量拷贝开销。 from multiprocessing.shared_memory import SharedMemory from nanovllm.config import Config from nanovllm.engine.sequence import Sequence from nanovllm.models.qwen3 import Qwen3ForCausalLM from nanovllm.layers.sampler import Sampler from nanovllm.utils.context import set_context, get_context, reset_context from nanovllm.utils.loader import load_model class ModelRunner: def __init__(self, config: Config, rank: int, event: Event | list[Event]): self.config config hf_config config.hf_config self.block_size config.kvcache_block_size self.enforce_eager config.enforce_eager #这里注意world_size的意思 self.world_size config.tensor_parallel_size #rank是进程编号的意思 self.rank rank #用于同步的信号量 self.event event #此处解析见问题1 dist.init_process_group(nccl, tcp://localhost:2333, world_sizeself.world_size, rankrank) torch.cuda.set_device(rank) default_dtype torch.get_default_dtype() #torch.empty会使用这个默认dtype分配空间 torch.set_default_dtype(hf_config.dtype) torch.set_default_device(cuda) self.model Qwen3ForCausalLM(hf_config) load_model(self.model, config.model) self.sampler Sampler() self.warmup_model() self.allocate_kv_cache() if not self.enforce_eager: self.capture_cudagraph() torch.set_default_device(cpu) torch.set_default_dtype(default_dtype) if self.world_size 1: #多GPU时由主进程创建共享内存来通信 if rank 0: self.shm SharedMemory(namenanovllm, createTrue, size2**20) #等待所有子进程连接前面声明了world_size dist.barrier() else: #等待rank 0 创建完共享内存 dist.barrier() #连接到共享内存 self.shm SharedMemory(namenanovllm) #进入循环等待rank 0 发指令 self.loop() def exit(self): if self.world_size 1: self.shm.close() dist.barrier() if self.rank 0: self.shm.unlink() #此处解析见问题3 if not self.enforce_eager: del self.graphs, self.graph_pool torch.cuda.synchronize() dist.destroy_process_group() #控制子进程循环等待rank 0的命令并执行 def loop(self): while True: #读取命令 method_name, args self.read_shm() #*表示解包args self.call(method_name, *args) if method_name exit: break #子进程读命令 def read_shm(self): assert self.world_size 1 and self.rank 0 #阻塞 self.event.wait() #读取数据有多长 n int.from_bytes(self.shm.buf[0:4], little) #读取数据并反序列化 method_name, *args pickle.loads(self.shm.buf[4:n4]) self.event.clear() return method_name, args #rank0写命令 def write_shm(self, method_name, *args): assert self.world_size 1 and self.rank 0 #序列化以便写入 data pickle.dumps([method_name, *args]) n len(data) #写入数据的长度把数据长度编码为 4 字节的小端整数 self.shm.buf[0:4] n.to_bytes(4, little) #写入数据 self.shm.buf[4:n4] data #通知所有子进程每个子进程有独立的 Event是一个 Event 列表这是LLMengine里的 for event in self.event: event.set() #统一调用入口 def call(self, method_name, *args): if self.world_size 1 and self.rank 0: self.write_shm(method_name, *args) method getattr(self, method_name, None) return method(*args) #预热 CUDA kernels 和 CUDA Graph 捕获避免首次推理的冷启动开销。其他解析见问题5 def warmup_model(self): torch.cuda.empty_cache() torch.cuda.reset_peak_memory_stats() # 构建足够大的假 seq 来预热真实请求通过 LLMEngine.add_request() 添加 max_num_batched_tokens, max_model_len self.config.max_num_batched_tokens, self.config.max_model_len seq_len min(max_num_batched_tokens, max_model_len) num_seqs min(max_num_batched_tokens // seq_len, self.config.max_num_seqs) seqs [Sequence([0] * seq_len) for _ in range(num_seqs)] for seq in seqs: seq.num_scheduled_tokens seq_len # 执行一次 prefill self.run(seqs, True) torch.cuda.empty_cache() #计算可用的显存并分配kv_cache block绑定到模型上 def allocate_kv_cache(self): config self.config hf_config config.hf_config #获取GPU显存空闲 / 总容量 free, total torch.cuda.mem_get_info() used total - free #获取历史最大峰值 peak torch.cuda.memory_stats()[allocated_bytes.all.peak] #获取当前使用的显存量 current torch.cuda.memory_stats()[allocated_bytes.all.current] #每个rank的kv_head数量 num_kv_heads hf_config.num_key_value_heads // self.world_size #每个注意力头对应的隐维度 head_dim getattr(hf_config, head_dim, hf_config.hidden_size // hf_config.num_attention_heads) #计算一个kv_cache block占用的字节 #公式含义是(KV) * 层数 * 一个块的token数 * 注意力头数量 * 隐维度 * dtype字节数 block_bytes 2 * hf_config.num_hidden_layers * self.block_size * num_kv_heads * head_dim * hf_config.dtype.itemsize # 可用块数 (总显存 * 利用率 - 已用 - 峰值 当前) / 单块字节数 #- peak current 的意义是为潜在的碎片/临时分配预留空间。 config.num_kvcache_blocks int(total * config.gpu_memory_utilization - used - peak current) // block_bytes assert config.num_kvcache_blocks 0 #预分配 KV cache tensor[2(K/V), 层数, 块数, block_size, kv_heads, head_dim] self.kv_cache torch.empty(2, hf_config.num_hidden_layers, config.num_kvcache_blocks, self.block_size, num_kv_heads, head_dim) #把预分配的 tensor 绑定到每个 attention 层的 k_cache / v_cache 上 layer_id 0 for module in self.model.modules(): if hasattr(module, k_cache) and hasattr(module, v_cache): module.k_cache self.kv_cache[0, layer_id] module.v_cache self.kv_cache[1, layer_id] layer_id 1 #根据sequence中最长的block_table长度构造block_table #其在运行时被调用见下方prefill和decode def prepare_block_tables(self, seqs: list[Sequence]): max_len max(len(seq.block_table) for seq in seqs) block_tables [seq.block_table [-1] * (max_len - len(seq.block_table)) for seq in seqs] block_tables torch.tensor(block_tables, dtypetorch.int32, pin_memoryTrue).cuda(non_blockingTrue) return block_tables #注意这个函数会把需要prefill的sequences拼成一个一维向量处理记住这个才能理解下面的逻辑 #问题12有更详细的解析 def prepare_prefill(self, seqs: list[Sequence]): input_ids [] #本次要处理的 token positions [] #对应的位置编码 cu_seqlens_q [0] #和k一样都是flash attention需要的边界信息 cu_seqlens_k [0] max_seqlen_q 0 max_seqlen_k 0 slot_mapping [] #每个 token 对应 KV cache 的槽位但必须注意 #slot_mapping存的是kv_cache 物理槽位地址其值为block_table[逻辑块号] × block_size 块内偏移 block_tables None #遍历每个seq构建输入 for seq in seqs: #跳过已经处理的token(已经处理的都会成cache) start seq.num_cached_tokens #seqlen_q是本次处理的token长度 seqlen_q seq.num_scheduled_tokens end start seqlen_q #seqlen_k是完整上下文长度(会加上本次处理的token长度) seqlen_k end #注意这里cu_seqlens_q和k的涉及跟下面这句有关input_ids、positions、slot_mapping 都是扁平化的一维数组 #input_ids把多个seq的token拼在一起 input_ids.extend(seq[start:end]) positions.extend(range(start, end)) cu_seqlens_q.append(cu_seqlens_q[-1] seqlen_q) cu_seqlens_k.append(cu_seqlens_k[-1] seqlen_k) max_seqlen_q max(seqlen_q, max_seqlen_q) max_seqlen_k max(seqlen_k, max_seqlen_k) if not seq.block_table: # warmup continue start_block start // self.block_size #向上取整end_block是不包含的上界 end_block (end self.block_size - 1) // self.block_size #这里会记录每一个block中的每一个token的物理地址加入slot_map for i in range(start_block, end_block): slot_start seq.block_table[i] * self.block_size if i start_block: slot_start start % self.block_size if i ! end_block - 1: slot_end seq.block_table[i] * self.block_size self.block_size #如果是最后一个block else: slot_end seq.block_table[i] * self.block_size end - i * self.block_size #slot_mapping 告诉 attention每个 token 的 KV 应该写到哪里。 slot_mapping.extend(range(slot_start, slot_end)) #为True说明有前缀缓存需要传 block_tables #因为cu_seqlens_q[-1]是本次处理的token数cu_seqlens_k[-1]是完整上下文长度 if cu_seqlens_k[-1] cu_seqlens_q[-1]: # prefix cache block_tables self.prepare_block_tables(seqs) #参数的含义见问题11 input_ids torch.tensor(input_ids, dtypetorch.int64, pin_memoryTrue).cuda(non_blockingTrue) positions torch.tensor(positions, dtypetorch.int64, pin_memoryTrue).cuda(non_blockingTrue) cu_seqlens_q torch.tensor(cu_seqlens_q, dtypetorch.int32, pin_memoryTrue).cuda(non_blockingTrue) cu_seqlens_k torch.tensor(cu_seqlens_k, dtypetorch.int32, pin_memoryTrue).cuda(non_blockingTrue) slot_mapping torch.tensor(slot_mapping, dtypetorch.int32, pin_memoryTrue).cuda(non_blockingTrue) #把这些参数存到一个全局context对象里供 attention kernel 在 forward 时读取 set_context(True, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, slot_mapping, None, block_tables) return input_ids, positions #准备当前已有token的位置信息 def prepare_decode(self, seqs: list[Sequence]): input_ids [] positions [] slot_mapping [] context_lens [] for seq in seqs: #此处是上一步生成的token input_ids.append(seq.last_token) positions.append(len(seq) - 1) context_lens.append(len(seq)) slot_mapping.append(seq.block_table[-1] * self.block_size seq.last_block_num_tokens - 1) input_ids torch.tensor(input_ids, dtypetorch.int64, pin_memoryTrue).cuda(non_blockingTrue) positions torch.tensor(positions, dtypetorch.int64, pin_memoryTrue).cuda(non_blockingTrue) slot_mapping torch.tensor(slot_mapping, dtypetorch.int32, pin_memoryTrue).cuda(non_blockingTrue) context_lens torch.tensor(context_lens, dtypetorch.int32, pin_memoryTrue).cuda(non_blockingTrue) block_tables self.prepare_block_tables(seqs) set_context(False, slot_mappingslot_mapping, context_lenscontext_lens, block_tablesblock_tables) return input_ids, positions def prepare_sample(self, seqs: list[Sequence]): temperatures [seq.temperature for seq in seqs] temperatures torch.tensor(temperatures, dtypetorch.float32, pin_memoryTrue).cuda(non_blockingTrue) return temperatures #这个装饰器的作用见问题15 torch.inference_mode() def run_model(self, input_ids: torch.Tensor, positions: torch.Tensor, is_prefill: bool): #这里512不对应最大sequence数量这三个条件都是跳过CUDA Graphs强制eager模式 #见问题9 if is_prefill or self.enforce_eager or input_ids.size(0) 512: return self.model.compute_logits(self.model(input_ids, positions)) else: #用CUDA Graph的graph.replay()重放录好的图—图的形状固定1/2/4/8/16/32.../512性能最好。 bs input_ids.size(0) context get_context() #选一张能容纳 bs 的预录图 graph self.graphs[next(x for x in self.graph_bs if x bs)] #预分配的显存buffer graph_vars self.graph_vars graph_vars[input_ids][:bs] input_ids graph_vars[positions][:bs] positions graph_vars[slot_mapping].fill_(-1) graph_vars[slot_mapping][:bs] context.slot_mapping graph_vars[context_lens].zero_() graph_vars[context_lens][:bs] context.context_lens graph_vars[block_tables][:bs, :context.block_tables.size(1)] context.block_tables # 重放录好的 CUDA 图比 eager 快很多 graph.replay() # 从 buffer 拿结果 return self.model.compute_logits(graph_vars[outputs][:bs]) #真正的运行代码prefill阶段会建立kv-cache但是prefill返回的token在postprocess阶段会直接被丢弃decode阶段会产生token def run(self, seqs: list[Sequence], is_prefill: bool) - list[int]: input_ids, positions self.prepare_prefill(seqs) if is_prefill else self.prepare_decode(seqs) temperatures self.prepare_sample(seqs) if self.rank 0 else None logits self.run_model(input_ids, positions, is_prefill) token_ids self.sampler(logits, temperatures).tolist() if self.rank 0 else None reset_context() return token_ids #预先为多种batch size把模型forward的执行轨迹录制下来。后续decode阶段只需graph.replay()跳过Python/调度器开销达到接近原生kernel的速度。 torch.inference_mode() def capture_cudagraph(self): config self.config hf_config config.hf_config max_bs min(self.config.max_num_seqs, 512) max_num_blocks (config.max_model_len self.block_size - 1) // self.block_size input_ids torch.zeros(max_bs, dtypetorch.int64) positions torch.zeros(max_bs, dtypetorch.int64) slot_mapping torch.zeros(max_bs, dtypetorch.int32) context_lens torch.zeros(max_bs, dtypetorch.int32) block_tables torch.zeros(max_bs, max_num_blocks, dtypetorch.int32) outputs torch.zeros(max_bs, hf_config.hidden_size) self.graph_bs [1, 2, 4, 8] list(range(16, max_bs 1, 16)) self.graphs {} self.graph_pool None #先录大图占满显存后录小图时复用剩余空间。 for bs in reversed(self.graph_bs): graph torch.cuda.CUDAGraph() set_context(False, slot_mappingslot_mapping[:bs], context_lenscontext_lens[:bs], block_tablesblock_tables[:bs]) #第一次跑触发 kernel 选择、autotune 等热身目的是把那些一次性初始化跑完避免它们被录进图里。 outputs[:bs] self.model(input_ids[:bs], positions[:bs]) # warmup # 第二次跑进入 capture 上下文把 GPU 操作序列录下来 #这是PyTorch 的 capture API进入这个上下文后所有 GPU 操作会被记录到一个 graph 对象中不真正执行——graph.replay() 时才真正执行。 with torch.cuda.graph(graph, self.graph_pool): outputs[:bs] self.model(input_ids[:bs], positions[:bs]) # capture if self.graph_pool is None: # 第一次记录后续图共享这块显存池 self.graph_pool graph.pool() self.graphs[bs] graph torch.cuda.synchronize() reset_context() #保存 buffer 引用 self.graph_vars dict( input_idsinput_ids, positionspositions, slot_mappingslot_mapping, context_lenscontext_lens, block_tablesblock_tables, outputsoutputs, )一些问题1.请介绍一下dist.init_process_group(nccl, tcp://localhost:2333, world_sizeself.world_size, rankrank)做了什么 核心做了这几件事 通信后端初始化指定nccl作为 GPU 间通信后端加载 NCCL 库并初始化通信上下文 进程握手通过tcp://localhost:2333主节点地址完成所有进程的网络握手建立通信通道 元信息分配将world_size总进程数、rank当前进程编号同步给所有进程后续进程间通信依赖这些标识 同步屏障隐式等待所有进程完成初始化确保后续操作的一致性。 2.为什么要先torch.set_default_device(cuda)然后又设置成cpu 设为 CUDA模型初始化、权重加载、KV 缓存分配、CUDA Graph 捕获等操作必须在 GPU 执行设置默认设备为 CUDA 可以避免手动给每个张量加.cuda()简化代码 切回 CPU推理阶段的非核心张量如预处理的临时列表、采样参数默认在 CPU避免不必要的 GPU 内存占用后续仅将推理必需的张量input_ids/positions显式移到 CUDA最大化 GPU 内存利用率。 3.为什么清理CUDA Graph资源后要进行CUDA同步 CUDA 操作是异步执行的CPU 下发指令GPU 后台执行del 只是释放 CPU 侧的 Graph 句柄但 GPU 可能还在执行 Graph 相关的内核操作 torch.cuda.synchronize() 会阻塞 CPU直到 GPU 完成所有未完成的操作此时再释放资源 避免 GPU 侧的 Graph 资源被提前释放导致野指针 避免内存泄漏CUDA 内存未及时回收。 4.详细解析event这个东西它是怎么协调多进程的 代码中event是主从进程的 “信号开关”完整协调流程 plaintext 主进程rank0 子进程rank0 │ │ │ 1. 写入命令到共享内存 │ 1. 调用event.wait() → 阻塞 │ 2. event.set() → 发信号 │ 2. 收到信号退出阻塞 │ │ 3. 读取共享内存的命令 │ │ 4. 执行命令 │ │ 5. event.clear() → 重置信号 │ │ 6. 回到wait()等待下一次信号 核心Event 保证子进程只在主进程写完命令后执行避免 “读空 / 读脏数据” 代码中event是列表对应每个子进程主进程会给所有子进程的event发信号确保同步执行。 这里主卡写入命令后所有从卡都要执行相同命令 5.为什么要预热的具体原因 并且CUDA kernels 第一次执行有 JIT 编译开销 capture_cudagraph() 在 __init__ 末尾调用需要先跑一次 forward 才能捕获计算图 测量峰值显存时需要先清零统计 6.allocate_kv_cache这里似乎会把空闲的显存都用于kv_cache块分配不会导致没有显存加载模型吗 模型权重在 used 里见下面的调用顺序 self.model Qwen3ForCausalLM(hf_config) # 模型加载到 GPU → 占用显存 load_model(self.model, config.model) # 加载权重 → 占用显存 self.warmup_model() # 运行一次 forward → 可能增加临时显存 self.allocate_kv_cache() # ← 此时才计算 KV cache 块数 - peak current 的作用 这解决的是碎片/临时分配的问题。warmup 后显存可能有一些碎片peak 是历史最高点current 是现在实际占用。差值就是峰值回落了多少——给临时分配留的空间。 简单说分配 KV cache 时预留了模型权重 临时计算的空间不会把模型挤掉。 7.allocate_kv_cache是什么时候执行的如果是加载模型后不久此时sequence的block_table是空的 1. 模型加载、warmup、allocate_kv_cache ← 此时 block_table 都是 [] 2. 用户调用 llm.generate() 或 add_request() 3. Scheduler.schedule() → 调用 block_manager.allocate(seq, num_cached_blocks) → 分配物理 Block填充 seq.block_table 4. LLMEngine.step() → model_runner.call(run, seqs, is_prefill) → run() 调用 prepare_prefill() 或 prepare_decode() → prepare_block_tables() ← 此时 block_table 已经有值了 8.prepare_prefill只把sequence需要调度的token加入input_ids那怎么知道它对应哪个sequence并计算 cu_seqlens 就是用来解决这个问题的 完整的流程 假设 batch 中 2 个 seq 要处理 Seq_0: [A, B, C, D] (4个token) Seq_1: [E, F, G, H, I, J] (6个token) 扁平化后 input_ids [A, B, C, D, E, F, G, H, I, J] positions [0, 1, 2, 3, 0, 1, 2, 3, 4, 5] ↑Seq_0 ↑Seq_1 cu_seqlens_q [0, 4, 10] ↑ ↑ Seq_0 Seq_1 在[0,4)在[4,10) cu_seqlens_k [0, 4, 10] (假设都是完整上下文) attention kernel 内部 Flash Attention 等 kernel 看到 cu_seqlens_q 和 cu_seqlens_k就知道 第 0 个 query 序列input_ids[0:4]属于 Seq_0 第 1 个 query 序列input_ids[4:10]属于 Seq_1 但input_ids这里没有sequence的完整上下文attention kernel怎么知道他完整的上下文 这里就是cu_seqlens_k 的作用 cu_seqlens_k.append(cu_seqlens_k[-1] seqlen_k) # 其中 seqlen_k end start seqlen_q 第二次调度时 start 4已缓存了 4 个 seqlen_q 4本次处理 4 个 seqlen_k 8总上下文 缓存 4 新算 4 cu_seqlens_k [0, 8] 这告诉 attention kernelK 的总长度是 8需要从 KV cache 读取 8 个 K/V 值来计算 attention。 9.run_model中的三个条件是什么意思 这三个条件都是跳过 CUDA Graphs、强制用 eager 模式的原因 is_prefill Prefill 阶段一次要处理大量 token整个 promptbatch size 也可能很大形状不固定。CUDA Graph 需要精确捕获输入形状提前录制好一个固定图。prompt 长度人人不同很难用同一张图覆盖所以 prefill 阶段直接 eager 跑。 self.enforce_eager 用户显式传入的参数强制不用 CUDA Graph比如调试、测性能、或者遇到奇怪的 GPU bug。 input_ids.size(0) 512 CUDA Graph 的显存池graph pool是提前按某个 batch size 分配的。超过 512 的 batch size 没被录制过图直接 replay 会越界或拿到错误的预分配显存所以走 eager。 10.简单介绍CUDA GRAPH CAPTURE 捕获阶段 提前为常见 batch size 各录一张 GPU 操作图 共享显存池 预分配输入输出 buffer 后续推理阶段只需把真实数据灌进 buffer → replay → 读结果跳过 Python 调度和 GPU kernel launch 开销。 11.block_tables torch.tensor(block_tables, dtypetorch.int32, pin_memoryTrue).cuda(non_blockingTrue)这里的pin_memory和non_blocking是什么意思 (1) pin_memoryTrue 含义将张量分配到 CPU 的锁页内存page-locked memory操作系统不会把这块内存交换到虚拟内存硬盘中。 优势锁页内存与 GPU 间的拷贝速度比普通内存快避免内存分页 / 交换 场景需频繁从 CPU 拷贝到 GPU 的张量如 input_ids/positions。 (2) non_blockingTrue.cuda(non_blockingTrue) 含义CPU→GPU 的拷贝操作是异步的不阻塞当前 CPU 线程 前提必须配合pin_memoryTrue普通内存无法异步拷贝 优势CPU 可在拷贝的同时执行其他逻辑提升推理吞吐量。 12.prepare_prefill和prepare_decode函数解析 prepare_prefill处理一批序列的“预填充” 对每个序列本次 forward 只计算 [start, end) 这段 token其中 start seq.num_cached_tokens前面这么多 token 的 KV 已经通过 prefix cache 缓存好了不用重算 seqlen_q seq.num_scheduled_tokens调度器这次安排算多少个调度器会受 max_num_batched_tokens 限制一个长 prompt 可能分多次算 end start seqlen_q。 然后构造几样东西 input_ids / positions把所有序列的 [start, end) 段 token 和对应位置拼成一个大一维张量变长拼接不是 padding。 cu_seqlens_q / cu_seqlens_k累积长度前缀和是 varlen flash attention 的标准输入用来告诉 kernel “拼接的大张量里哪一段属于哪个序列”。注意 Q 和 K 长度不同Q 只有本次新算的 seqlen_q 个K 是全长度 end 个——因为注意力要看到全部历史含缓存的前缀所以 seqlen_k end。这也解释了最后那个 if cu_seqlens_k[-1] cu_seqlens_q[-1]: block_tables ...只有存在前缀缓存K 比 Q 长时本次 attention 才需要读 block_table 里的历史 KV纯 prefillstart0时历史 KV 全是本次刚写进去的attention kernel 内部自己就知道边界不需要 block_table。 slot_mapping为 [start, end) 里的每个 token 算出它该把 KV 写进哪个物理槽位。 prepare_decode自回归解码每序列每次只算 1 个 token decode 阶段每步每个序列只 forward 一个 token——刚刚采样出来的那个。所以要构造 input_ids每序列 1 个 tokenseq.last_token positionslen(seq) - 1即这个 token 在序列里的位置最后一个位置 context_lenslen(seq)注意力要看的 KV 总长度包含当前这个正在算的 token——它的 KV 会先被写进 cache 再做 attention slot_mapping每序列 1 个槽位即 最后一个物理块 × block_size last_block_num_tokens - 1。last_block_num_tokens 是最后一个逻辑块里已有的 token 数减 1 就是“下一个空位”在块内的偏移——恰好是位置 len(seq)-1 那个 token 的槽位你可以验证(len(seq)-1) % block_size last_block_num_tokens - 1数学上恒等 block_tablesdecode 一定需要因为每个序列的历史 KV 分散在多个物理块里attention kernel 要靠 block_table 才能找到。 13.另外decode的seq.last_token 是未产生还是已产生 已经产生的 token但它的 KV 还没算过。 时序是这样的见 scheduler.py:84-88每轮 run() 结束、采样出 token_id 后调度器先更新缓存计数然后 seq.append_token(token_id) 把新 token 追加进 token_idslast_token 同步指向它。下一轮 decode 的 prepare_decode 取这个 last_token 作为输入——这个 token 在上一步被采样出来了作为上一次 logits 的输出但它作为输入的 KV 和 logits 还从没算过。本次 forward 就是为它服务的算出它的 K/V 写入 slot_mapping 指向的槽位然后对全部 context_lens 个历史 token 做 attention产出下一个 token 的 logits。 一句话last_token 是“上一步刚生、这一步刚要吃进去”的 token。 14. prefill 里 slot_mapping 循环的算法 先把代码的意图一句话说穿 逻辑位置为 p 的 token其槽位 block_table[p // block_size] × block_size p % block_size。 这个循环只是把上面这个逐 token 的公式按块合并成 range() 批量生成避免 Python 里逐 token 循环。设 bs block_size python start_block start // bs # start 所在的逻辑块号 end_block ceil(end / bs) # end-1 所在的逻辑块号 1左闭右开 for i in range(start_block, end_block): slot_start block_table[i] * bs # 该逻辑块对应物理块的起点槽位 if i start_block: slot_start start % bs # 第一块要从 start 的偏移处开始不是从头 if i ! end_block - 1: slot_end block_table[i] * bs bs # 中间的块整块都属于本次调度 else: slot_end block_table[i] * bs (end - i*bs) # 最后一块只用到逻辑位置 end-1 slot_mapping.extend(range(slot_start, slot_end)) 三个分支的含义 第一块i start_block块内前面的 start % bs 个位置属于历史已缓存跳过所以起点要加偏移 中间的块整个块8 个槽全部属于本次 [start, end) 区间直接整块加入 最后一块i end_block - 1本次只算到逻辑位置 end - 1所以块内只用前 end - i*bs 个槽i*bs 是这块在序列里的起始位置end - i*bs 就是 end 在块内的偏移。 用一个例子手算 block_table [4, 5, 9, 3, 11]block_size 8start 逻辑位置 1第一个块的第二个 tokenend 38最后一个块的倒数第二个 token。 先算块号 start_block 1 // 8 0 end_block (38 7) // 8 5即逻辑块 0~4 都被触及最后一块是 end_block-1 4 然后逐块 i 逻辑块 i 对应物理块 slot_start slot_end 槽位区间 对应的逻辑位置 0 4 4×8 (1%8) 33第一块加偏移 1 4×88 40中间块整块 [33, 40) 共 7 个 1 ~ 7 1 5 5×8 40 408 48 [40, 48) 共 8 个 8 ~ 15 2 9 9×8 72 728 80 [72, 80) 共 8 个 16 ~ 23 3 3 3×8 24 248 32 [24, 32) 共 8 个 24 ~ 31 4最后一块 11 11×8 88非第一块无偏移 11×8 (38 − 4×8) 886 94 [88, 94) 共 6 个 32 ~ 37 注意两个容易懵的点 槽位不连续、甚至不单调是正常的slot_mapping 的顺序跟的是逻辑位置的顺序槽位列表的第 k 个元素对应拼接 input_ids 里的第 k 个 token数值本身才指向物理位置。 最后一块的 slot_end 不是块的末尾而是 88 38 - 32i×bs 32 是逻辑块 4 的起始逻辑位置end - i×bs 6 表示这块内只有前 6 个位置32~37属于本次计算。 验证总数78886 37 end − start 38 − 1 ✓恰好等于本次调度的 token 数。 15.torch.inference_mode()在做什么 torch.inference_mode() 是 PyTorch 专为推理场景设计的装饰器替代旧的torch.no_grad()核心作用 禁用梯度计算不再追踪张量的requires_grad减少显存占用无需存储梯度相关的中间张量 关闭 autograd 校验禁用梯度相关的断言、验证逻辑提升推理速度 优化推理性能PyTorch 会针对推理模式做算子融合、内存复用等底层优化 不可反向传播推理模式下的张量默认requires_gradFalse且无法切换回训练模式除非退出该上下文。 代码中run_model/capture_cudagraph是推理核心函数该装饰器能大幅降低显存消耗、提升推理吞吐量。本系列文章(待写完修正)[系列(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)