【AI Infra 核心】从零剖析大模型服务框架：如何榨干 GPU 算力实现极致推理吞吐？

【AI Infra 核心】从零剖析大模型服务框架如何榨干 GPU 算力实现极致推理吞吐摘要上一篇我们通过 PagedAttention 解决了大模型推理时的“显存爆炸”危机。但在实际的生产环境中光有显存是不够的。老板花重金买的 A100/H100如果 GPU 利用率只有 20%那无异于暴殄天物。今天我们深入剖析 LLM 服务框架如 vLLM、TGI、TensorRT-LLM的核心调度机制用代码手写一个Continuous Batching连续批处理调度器带你彻底搞懂如何榨干 GPU 的每一滴算力一、 传统 Static Batching 的原罪“木桶效应”在传统的深度学习比如 ResNet 图像分类中为了提高 GPU 的吞吐量我们通常会把多个请求拼成一个 Batch 一起送进显卡。这叫Static Batching静态批处理。但在大语言模型LLM的世界里这套玩法行不通了。为什么因为 LLM 生成文本的长度是不可预测的。假设我们把 3 个请求打成一个 Batch 送去推理请求 A只需要生成 5 个 Token 就结束了比如回答“北京的首都是哪”。请求 B需要生成 50 个 Token。请求 C需要生成 200 个 Token比如写一篇小作文。在静态批处理下整个 Batch 必须等待最长的“请求 C”全部生成完毕才能整体返回并接收下一批任务。结果就是请求 A 的计算在第 5 步就结束了在剩下的 195 步里它对应的 GPU 算力完全处于闲置Padding/Idle状态这就是典型的“木桶效应”导致 GPU 算力被白白浪费。二、 破局之道Continuous Batching (连续批处理)为了打破木桶效应Orca 论文率先提出了Iteration-level Scheduling迭代级调度也就是后来工业界熟知的Continuous Batching或In-flight Batching。它的核心思想非常暴力且有效打破 Batch 的静态边界把调度粒度细化到每一个 Token每一次 Iteration系统维护一个“正在运行”Running的请求池。在每一个 Token 生成的 Step即一次 Forward Pass结束后检查池子里的请求。如果请求 A 生成完毕遇到了 EOS token立即把它踢出 Batch返回给用户。立即从等待队列中拉取一个新的请求 D把它塞进刚刚空出来的 Batch 位置。GPU 继续执行下一个 Step 的计算。通过这种“动态进出”的机制GPU 永远处于满载状态没有一滴算力被浪费在无意义的 Padding 上。三、 纸上得来终觉浅Python 徒手实现 Continuous Batching 调度器为了让大家直观理解底层的调度逻辑我们剥离掉复杂的 CUDA 和网络通信代码用纯 Python 面向对象的方式写一个极简版的 Continuous Batching 核心引擎。1. 定义请求对象 (Request)首先我们需要一个类来记录每个请求的状态是等待中、正在运行还是已完成。fromenumimportEnumclassStatus(Enum):WAITING0# 在队列中等待RUNNING1# 正在 GPU 中生成FINISHED2# 生成完毕classRequest:def__init__(self,req_id:int,prompt_len:int,max_output_len:int):self.req_idreq_id self.prompt_lenprompt_len self.max_output_lenmax_output_len self.generated_tokens0self.statusStatus.WAITINGdefis_finished(self):# 简单模拟达到最大输出长度或随机遇到终止符视为完成returnself.generated_tokensself.max_output_lendef__repr__(self):returnfReq(id{self.req_id}, tokens{self.generated_tokens}/{self.max_output_len})2. 构建核心调度器 (Scheduler)这是整个框架的大脑它决定了每一个 Forward Step 中哪些请求上 GPU哪些请求下车。classContinuousBatchingScheduler:def__init__(self,max_batch_size:int):self.max_batch_sizemax_batch_size self.waiting_queue[]# 等待处理的请求self.running_batch[]# 正在 GPU 上跑的请求defadd_request(self,request:Request):接收上游网关发来的新请求self.waiting_queue.append(request)print(f[API] 接收到新请求{request.req_id}加入等待队列。)defstep(self): 模拟一次 GPU 的 Forward Pass (生成一个 Token) # 1. 踢出已完成的请求finished_reqs[reqforreqinself.running_batchifreq.is_finished()]forreqinfinished_reqs:req.statusStatus.FINISHED self.running_batch.remove(req)print(f [-] 请求{req.req_id}已完成离开 Batch。)# 2. 从等待队列拉取新请求填补 Batch 空缺whilelen(self.running_batch)self.max_batch_sizeandself.waiting_queue:new_reqself.waiting_queue.pop(0)new_req.statusStatus.RUNNING self.running_batch.append(new_req)print(f [] 动态拉取请求{new_req.req_id}进入 Batch。)# 3. 模拟 GPU 推理为当前 Batch 中的所有请求生成 1 个 Tokenifnotself.running_batch:print(GPU 闲置中...)returnFalse# 没有任务了print(f [GPU Compute] 当前 Batch 状态:{self.running_batch})forreqinself.running_batch:req.generated_tokens1returnTrue3. 运行测试见证奇迹的时刻我们来模拟一个真实的请求流看看它与静态 Batching 有什么不同。# 初始化一个最大 Batch Size 为 3 的调度器schedulerContinuousBatchingScheduler(max_batch_size3)# 模拟并发到来 5 个长度不一的请求scheduler.add_request(Request(1,prompt_len10,max_output_len2))scheduler.add_request(Request(2,prompt_len12,max_output_len5))scheduler.add_request(Request(3,prompt_len8,max_output_len3))scheduler.add_request(Request(4,prompt_len20,max_output_len4))scheduler.add_request(Request(5,prompt_len15,max_output_len2))# 驱动 GPU 进行推理直到所有任务完成step_count1whilescheduler.waiting_queueorscheduler.running_batch:print(f\n GPU Step{step_count})scheduler.step()step_count1运行结果解析在 Step 3 结束时请求 1 已经达到了最大长度2被立刻踢出 Batch。同时处于等待队列中的请求 4 立刻被补位拉入 Batch 中参与 Step 4 的计算。没有任何一个 GPU 时钟周期被用来做无意义的 Padding 等待整体吞吐量在真实场景下可以翻 20 倍以上。四、 进阶与博弈Prefill 与 Decode 的分离在上述代码中我们做了简化把每个 Step 视作等价的。但在真实的 LLM 推理中阶段分为两步Prefill 阶段预填充处理用户的 Prompt是计算密集型的矩阵乘法GEMM。Decode 阶段逐字生成生成后续的 Token是访存密集型的矩阵向量乘法GEMV。在高端的工程落地中如 vLLM 中的 Chunked Prefill 技术为了防止新加入请求的超长 Prompt 把正在生成的 Decode 请求卡住导致卡顿会将长的 Prefill 切割成小块Chunk与 Decode 任务混合在一个 Batch 里执行。五、 总结大模型推理架构优化的本质就是一场“压榨计算资源”与“对抗显存墙”的战争。通过结合上篇博客的PagedAttention解决空间碎片和本篇的Continuous Batching解决时间碎片现代 AI Infra 终于能够将 GPU 的利用率推向极致让千亿参数模型的大规模商用化成为可能。懂模型结构只能决定你能否跑通一个 Demo而懂底层 Infra 调度才决定了你能否扛住双十一级别的千万级高并发。