nano-vllm 用千行代码拆解 vLLM 核心是读懂大模型推理最快的捷径。1. 介绍上一篇把张量并行的三种切法weight、vocab、head讲完了。但 TP 这几篇的代码都在单进程里模拟两卡——把tp_size、tp_rank当参数手动传进去。真实的多卡是多进程每张卡一个进程tp_size、tp_rank从dist取。本篇看这套多进程怎么搭——进程怎么起、rank0 怎么把一次方法调用比如每个 step 的run传给每张卡、跑完怎么收摊。2. 总览以 tp2 为例。LLMEngine在主进程里起 1 个子进程连自己共 2 个进程每进程占一张卡、各跑一个ModelRunner。rank0 在主进程是 driver握着Scheduler、Tokenizer跑generate主循环自己也兼一个本卡ModelRunner。其余 rank 在子进程是 worker只有一个ModelRunner平时阻塞在loop()里等命令。两进程靠两条通道协作控制面SharedMemoryEventrank0 单向广播「该跑哪个方法」给 worker。数据面NCCLforward 内部的all_reduce、gather走这条前几篇的张量通信都在这里。核心是SPMD所有卡跑同一段代码。rank0 通过控制面让每张卡都调用同一个方法、同一批参数于是代码跑到all_reduce那行时所有卡正好都到齐——数据面的集合通信才对得上。打个比方rank0 像工头对着对讲机喊「现在都做run」每个工人worker听到同样的指令、对着自己那份料各做各的做到要合料的工序all_reduce时正好都到齐。3. 起进程一个 Python 进程只有一个 GIL全局解释器锁同一时刻只准一个线程执行 Python 代码CPU 端一次只能往一张卡发 kernel。要让多张卡真正并行就得多进程——每进程一个 GIL、独占一张卡。起子进程用 spawn 而非 forkfork 会继承父进程已初始化的 CUDA 上下文而 CUDA 不支持这样跨进程复用spawn 起一个全新的 Python 解释器各自干净地初始化 CUDA。主进程循环起tp_size - 1个子进程rank 1…tp-1每个targetModelRunner、占一张卡rank0 的ModelRunner留在主进程。每个子进程分一个Eventrank0 收下全部Event的列表广播时逐个set。importtorch.multiprocessingasmp# LLMEngine.__init__只留起进程部分def__init__(self,config):self.ps,self.events[],[]ctxmp.get_context(spawn)# 全新解释器CUDA 不能 forkforiinrange(1,config.tensor_parallel_size):# rank 1…tp-1eventctx.Event()pctx.Process(targetModelRunner,args(config,i,event))p.start()# 子进程跑 ModelRunner(config, i, event)self.ps.append(p)self.events.append(event)# rank0 留主进程拿到所有子进程的 event 列表self.model_runnerModelRunner(config,0,self.events)进程起来后每个ModelRunner一启动先dist.init_process_group(nccl, ...)加入 NCCL 组——NCCL 是 NVIDIA 的多卡通信库专管 GPU 之间传张量这就是数据面前几篇的all_reduce、gather都走它。接着各自set_device(rank)、建模型、加载权重、warmup、分配 KV cache。rank0 建好共享内存就返回、回LLMEngine干活worker 打开同一块共享内存后执行loop()阻塞等命令。中间的dist.barrier()是一道栅栏所有进程都到了才放行。rank0 先建好共享内存再过栅栏worker 过栅栏后才去打开——保证它打开时共享内存已存在。# ModelRunner.__init__def__init__(self,config,rank,event):dist.init_process_group(nccl,tcp://localhost:2333,world_sizeself.world_size,rankrank)# 加入数据面# … set_device、建模型、加载权重、warmup、分配 KV cache …ifself.world_size1:ifrank0:self.shmSharedMemory(namenanovllm,createTrue,size2**20)dist.barrier()# 建好共享内存等所有进程到齐else:dist.barrier()self.shmSharedMemory(namenanovllm)# 打开同一块self.loop()# 扎进 loop 阻塞等命令4. 一次 call 怎么传到每张卡控制面用两个多进程原语搭起来SharedMemory一块多个进程都能映射进各自地址空间、直接读写的内存相当于一块公共白板。进程之间内存本来互相隔离有了它 rank0 在白板上写、worker 在白板上读不用拷贝、最快。Event一个跨进程的信号灯有wait/set/clear。光有白板worker 不知道「写好了没」只能空转干等Event让 worker 先阻塞在wait()上rank0 写完set()一下把它叫醒。rank0 要让每张卡都跑某个方法比如每个 step 的run入口是统一的callrank0 调call先write_shm把方法名和参数广播出去再自己执行。worker 在loop()里read_shm收到再call执行worker 不广播否则会循环执行。importpickle# ModelRunner 的控制面四个方法defwrite_shm(self,method_name,*args):# 只有 rank0 调datapickle.dumps([method_name,*args])nlen(data)self.shm.buf[0:4]n.to_bytes(4,little)# 前 4 字节长度 nself.shm.buf[4:n4]data# 其后pickle 字节流foreventinself.event:# 逐个唤醒 workerevent.set()defread_shm(self):# 只有 worker 调self.event.wait()# 阻塞被 set 唤醒nint.from_bytes(self.shm.buf[0:4],little)method_name,*argspickle.loads(self.shm.buf[4:n4])self.event.clear()# 复位等下次returnmethod_name,argsdefcall(self,method_name,*args):# 两边共同入口ifself.world_size1andself.rank0:self.write_shm(method_name,*args)# 只有 rank0 广播methodgetattr(self,method_name,None)returnmethod(*args)# 各卡都执行defloop(self):# worker 主循环whileTrue:method_name,argsself.read_shm()self.call(method_name,*args)ifmethod_nameexit:breakwrite_shm先pickle成字节流pickle就是把 Python 对象和字节互相转换即序列化前 4 字节写长度 n读时才知道截到哪其后写数据最后event.set()逐个唤醒 worker。read_shmevent.wait()一直阻塞被唤醒后先读 4 字节拿 n、再按 n 截出字节流unpickle最后event.clear()复位好等下一次。call是 rank0 和 worker 的共同入口但只有 rank0 会write_shm——worker 的call由loop调再广播就死循环了。两边最后都getattr执行同一个方法这就是 SPMD。loopworker 一辈子在这转——读一条、执行一条直到读到exit才跳出。为什么方法调用走共享内存、不走 NCCLNCCL 只传张量而这里要传的是方法名加seqs这种 Python 对象pickle进共享内存最直接开 2²⁰ 字节够装一批请求的元数据。5. 退出清理generate跑完LLMEngine在atexit里调call(exit)——和run同一条控制面把exit广播给每个 worker。worker 的loop读到exit执行完exit就break出循环、进程结束主进程再join等它们退干净。exit做三件收尾关掉共享内存rank0 还要unlink真正删除、销毁 NCCL 进程组、同步 CUDA。importtorch# ModelRunner.exitdefexit(self):ifself.world_size1:self.shm.close()# 关掉本进程的共享内存映射dist.barrier()# 等所有进程都 close 完ifself.rank0:self.shm.unlink()# rank0 真正删除共享内存ifnotself.enforce_eager:delself.graphs,self.graph_pool# 释放 CUDA Graphtorch.cuda.synchronize()dist.destroy_process_group()# 拆掉 NCCL 组6. 小结多卡的多进程架构到这里清楚了起进程每卡一进程spawnCUDA 不能 forkrank0 留主进程当 driver其余 worker 进loop()阻塞等命令dist.barrier握手保证共享内存建好再打开。控制面SharedMemoryEventrank0write_shm广播方法名参数、event.set()唤醒workerread_shm收下执行。call是两边共同入口只有 rank0 广播。数据面NCCLforward 里的all_reduce、gather走这条。两条通道合起来就是SPMD所有卡跑同一段代码、步调一致集合通信才对得上。退出call(exit)广播worker 跳出loop各自关共享内存、拆 NCCL 组。张量并行至此介绍完毕。下一篇端到端跑通整个推理引擎。