1. 批处理Batching到底是什么为什么能“榨干”GPU大家好我是老K一个在AI大模型和智能硬件领域摸爬滚打了十来年的工程师。今天咱们接着聊大模型推理优化这次的主角是批处理Batching。如果你曾经部署过在线问答或者对话服务肯定遇到过这样的烦恼用户一个接一个地提问服务器上的GPU好像永远在“睡大觉”利用率低得可怜但响应速度却快不起来。这感觉就像你开着一辆八缸跑车却只在小区里以20码的速度遛弯既浪费油又跑不快。批处理就是解决这个问题的“油门”。简单来说批处理就是把多个用户的请求比如多个问题句子打包成一个“包裹”一次性塞给GPU进行计算而不是来一个算一个。这背后的道理其实很直观现代GPU无论是NVIDIA的A100、H100还是消费级的4090它们强大的算力核心CUDA Core/Tensor Core生来就是为了并行计算而设计的。让这些核心同时处理多个任务远比让它们排队处理单个任务要高效得多。你可以想象一下厨房里只有一个厨师GPU如果客人点一份炒饭他就做一份那大部分时间他都在等锅热、等油开效率极低。但如果他能同时处理五份炒饭的订单统筹安排下锅、翻炒、装盘那么单位时间内他能端出去的炒饭总量吞吐量就会大大增加。在实际的模型推理中一次前向传播Forward Pass的计算开销是固定的主要包括矩阵乘法和注意力机制的计算。当batch size为1时GPU的很多计算单元可能处于闲置状态。而当我们把batch size增加到4、8甚至更大时这些计算单元就能被更充分地利用起来分摊掉每次前向传播的固定开销从而显著提升吞吐量Throughput也就是每秒能处理多少token。我实测过一个7B参数的模型在A10 GPU上batch size从1提升到8吞吐量能从每秒几十个token暴涨到接近300个token效果非常明显。当然天下没有免费的午餐。批处理在提升吞吐量的同时往往会增加延迟Latency也就是单个请求从发出到收到第一个token所需的时间。因为GPU需要等凑够一批请求才开始计算。这就引出了我们今天要深入探讨的核心矛盾如何在吞吐量和延迟之间找到最佳平衡点这就像经营一家餐厅你是选择让每个顾客稍等片刻凑够一桌再一起上菜高吞吐高延迟还是坚持“随到随炒”低延迟低吞吐答案取决于你的业务场景。接下来我们就从代码层面开始一步步拆解批处理的实现细节和其中的权衡艺术。2. 从零开始手把手实现一个批处理推理引擎光说不练假把式咱们直接上代码。我会用一个简化但完整的例子带你走一遍批处理推理的全流程。这里我们以Hugging Face的transformers库和GPT-2模型为例因为它的结构清晰适合教学。你完全可以把这里的思路迁移到LLaMA、ChatGLM等任何你正在使用的模型上。2.1 环境准备与单条推理的回顾首先我们得把场子搭起来。确保你安装了PyTorch和transformers。pip install torch transformers我们先快速回顾一下不使用批处理时模型是如何生成一个token的。这能帮助我们理解批处理到底改变了什么。import torch from transformers import AutoTokenizer, AutoModelForCausalLM # 加载模型和分词器 model_name gpt2 tokenizer AutoTokenizer.from_pretrained(model_name) model AutoModelForCausalLM.from_pretrained(model_name) model.eval() # 设置为评估模式 def generate_token_with_past(inputs): 基于KV Cache生成单个token with torch.no_grad(): # 推理时不计算梯度节省内存 outputs model(**inputs) logits outputs.logits # 形状: [batch_size, seq_len, vocab_size] # 只取每个序列最后一个token的logits last_logits logits[:, -1, :] # 贪心解码取概率最大的token next_token_id last_logits.argmax(dim-1) return next_token_id, outputs.past_key_values # 单条输入的例子 prompt The quick brown fox inputs tokenizer(prompt, return_tensorspt) # 得到 input_ids, attention_mask generated prompt for _ in range(10): # 生成10个token next_token_id, past_key_values generate_token_with_past(inputs) # 把新生成的token拼接到输入中用于下一步生成 inputs { input_ids: next_token_id.unsqueeze(0), # 保持batch维度 attention_mask: torch.cat([inputs[attention_mask], torch.ones(1, 1)], dim1), past_key_values: past_key_values } generated tokenizer.decode(next_token_id) print(generated)这段代码就是典型的自回归生成每次调用模型只生成下一个token然后把这个token当作新的输入循环往复。它的缺点是显而易见的GPU强大的并行能力被白白浪费每次前向传播只处理一个非常小的序列。2.2 批处理的核心挑战处理不定长序列当我们想把多个句子比如[Hello world, How are you today?, AI is amazing]一起送入模型时第一个拦路虎就是它们的长度不一样。Transformer模型要求输入是一个规整的张量Tensor就像Excel表格一样每一行一个句子的列数token数必须相同。解决方案是填充Padding。我们把短的句子后面加上特殊的[PAD]token让所有句子变得一样长。但是我们肯定不希望模型去“理解”这些没意义的[PAD]token这就需要用到注意力掩码Attention Mask。# 设置填充token和策略 tokenizer.pad_token tokenizer.eos_token # 通常用结束符[EOS]作为填充符 model.config.pad_token_id model.config.eos_token_id # 关键设置采用左侧填充left padding tokenizer.padding_side left tokenizer.truncation_side left # 超长时从左边截断 # 准备一批输入 prompts [ The quick brown fox jumped over the, The rain in Spain falls, What comes up must, ] # 分词并自动填充 batch_inputs tokenizer(prompts, paddingTrue, return_tensorspt) print(input_ids 形状:, batch_inputs[input_ids].shape) # 例如 torch.Size([3, 10]) print(input_ids 内容:\n, batch_inputs[input_ids]) print(\nattention_mask 内容:\n, batch_inputs[attention_mask])运行这段代码你会看到input_ids是一个3行N列的矩阵短句前面被填满了pad_token_id。对应的attention_mask矩阵中真实token的位置是1填充token的位置是0。这个mask会在模型计算注意力时发挥作用让填充位置的权重为0从而被忽略。这里选择左侧填充至关重要。因为在自回归生成中我们是在序列的右侧不断添加新token。如果采用右侧填充那么新生成的token就会和填充符混在一起导致位置信息错乱。左侧填充保证了所有有效token都向右对齐新token可以无缝追加在右侧。2.3 实现批处理生成函数现在我们来改造之前的单条生成函数让它能处理一个批次batch。这是整个批处理推理的核心。def generate_batch_token_with_past(batch_inputs): 批处理版本一次为整个batch生成下一个token with torch.no_grad(): outputs model(**batch_inputs) logits outputs.logits # 形状: [batch_size, seq_len, vocab_size] # 关键变化我们同时取batch中每个序列的最后一个token的logits last_logits logits[:, -1, :] # 形状: [batch_size, vocab_size] # 为batch中的每个样本选择下一个token (贪心策略) next_token_ids last_logits.argmax(dim-1) # 形状: [batch_size] return next_token_ids, outputs.past_key_values def generate_batch(batch_inputs, max_new_tokens20): 批处理文本生成循环 batch_size batch_inputs[input_ids].shape[0] # 为batch中的每个样本单独维护一个生成结果列表 generated_texts [[] for _ in range(batch_size)] # 初始化position_ids这对于某些模型是必须的 attention_mask batch_inputs[attention_mask] # 计算每个token的绝对位置填充部分设为1或0取决于模型 position_ids attention_mask.long().cumsum(-1) - 1 position_ids.masked_fill_(attention_mask 0, 1) # 准备第一次模型输入的字典 current_inputs { input_ids: batch_inputs[input_ids], attention_mask: attention_mask, position_ids: position_ids, } for step in range(max_new_tokens): # 生成当前步的token (整个batch) next_token_ids, past_key_values generate_batch_token_with_past(current_inputs) # 解码这一步生成的token new_tokens tokenizer.batch_decode(next_token_ids) # 更新每个样本的生成结果 for i in range(batch_size): generated_texts[i].append(new_tokens[i]) # 准备下一步的输入 # 1. input_ids 更新为刚生成的token # 2. attention_mask 右侧添加一列1表示新token有效 # 3. position_ids 更新为最后一个位置1 current_inputs { input_ids: next_token_ids.unsqueeze(-1), # 形状 [batch_size, 1] attention_mask: torch.cat( [current_inputs[attention_mask], torch.ones(batch_size, 1, deviceattention_mask.device)], dim1 ), position_ids: current_inputs[position_ids][:, -1:] 1, # 取最后一个位置并加1 past_key_values: past_key_values, } # 将每个样本生成的token列表拼接成字符串 final_outputs [] original_prompts tokenizer.batch_decode(batch_inputs[input_ids], skip_special_tokensTrue) for i in range(batch_size): prompt original_prompts[i] continuation .join(generated_texts[i]) final_outputs.append(prompt continuation) return final_outputs # 运行批处理生成 results generate_batch(batch_inputs, max_new_tokens15) for i, res in enumerate(results): print(f样本 {i}: {res}\n)仔细对比一下单条生成的generate_token_with_past和批处理的generate_batch_token_with_past你会发现最核心的变化在于last_logits的索引和next_token_ids的获取。单条版本是logits[0, -1, :]和argmax()输出一个标量批处理版本是logits[:, -1, :]和argmax(dim1)输出一个包含batch_size个token ID的向量。模型的一次前向传播同时计算了batch中所有样本的下一个token概率这就是并行计算的威力。另一个需要特别注意的细节是past_key_values。在批处理中KV Cache也变成了批量的。模型返回的past_key_values是一个嵌套的元组或列表其中每个元素都对应一个注意力层并且其Key和Value张量的形状都包含了batch维度例如[batch_size, num_heads, seq_len, head_dim]。我们的代码不需要特殊处理它只需要在每一步将它正确地传递回模型即可。3. 吞吐量 vs. 延迟寻找那个甜蜜点代码跑通了但作为工程师我们更关心数字批处理到底带来了多少性能提升代价又是什么我们来做个简单的基准测试。3.1 设计一个性能测试实验我们将测量不同batch size下生成固定数量token所需的总时间duration、吞吐量throughput tokens/sec和平均延迟avg latency sec/token。这里为了简化我们用“平均延迟”来近似表示每个请求的响应时间实际上第一个token的延迟Time To First Token, TTFT和后续token的延迟Time Per Output Token, TPOT在批处理场景下会有不同的表现。import time def benchmark_batch_performance(batch_size_list, prompt_length10, max_new_tokens50): 基准测试不同batch size下的性能指标 # 准备一个固定的长句然后复制成不同的batch base_prompt The future of artificial intelligence is * 5 base_input tokenizer(base_prompt, return_tensorspt) seq_len base_input[input_ids].shape[1] results [] for bs in batch_size_list: # 构造batch复制相同的输入bs次 input_ids base_input[input_ids].repeat(bs, 1) attention_mask base_input[attention_mask].repeat(bs, 1) # 预热避免第一次运行慢 _ generate_batch({input_ids: input_ids, attention_mask: attention_mask}, max_new_tokens2) # 正式计时 start_time time.perf_counter() outputs generate_batch({input_ids: input_ids, attention_mask: attention_mask}, max_new_tokensmax_new_tokens) end_time time.perf_counter() duration end_time - start_time total_tokens_generated bs * max_new_tokens throughput total_tokens_generated / duration avg_latency duration / total_tokens_generated # 平均每个token的生成时间 results.append({ batch_size: bs, duration: duration, throughput: throughput, avg_latency: avg_latency }) print(fBatch Size: {bs:3d} | Duration: {duration:.4f}s | Throughput: {throughput:8.2f} tokens/s | Avg Latency: {avg_latency*1000:.2f} ms/token) return results # 测试不同的batch size batch_sizes_to_test [1, 2, 4, 8, 16, 32] performance_data benchmark_batch_performance(batch_sizes_to_test)跑完这个测试你大概率会得到类似下面这样的数据趋势Batch Size总耗时 (秒)吞吐量 (tokens/秒)平均延迟 (毫秒/token)10.10595.210.520.146136.87.340.178224.64.580.189422.62.4160.249643.21.6320.334957.31.0640.5051268.30.83.2 解读性能曲线与权衡看到这些数据我们可以画出两条关键的曲线吞吐量随batch size增长曲线和延迟随batch size增长曲线。吞吐量曲线通常是一个先快速上升后逐渐平缓甚至下降的曲线。在batch size较小时比如1到8增加batch size能极大地“摊薄”每次模型前向传播的固定开销如kernel启动、数据搬运GPU的SM流多处理器被更充分利用因此吞吐量几乎线性增长。当batch size达到一定规模后比如32以上增长会变慢因为GPU的计算资源如寄存器、共享内存可能成为瓶颈或者内存带宽跟不上数据供给的速度。延迟曲线则是一条单调上升的曲线。batch size越大GPU需要计算的数据量就越多完成一次前向传播的时间就越长。对于排在最前面的那个请求来说它必须等待整个batch计算完成才能得到结果所以它的延迟尤其是TTFT会随着batch size增大而增加。这就把我们带到了工程决策的核心你的服务更看重吞吐量还是延迟高吞吐优先场景离线任务处理、批量内容生成、模型微调后的推理验证。这些场景下任务没有实时交互需求我们追求在单位时间内处理尽可能多的数据。这时你可以选择一个较大的batch size比如32、64让GPU满载运行把硬件性价比拉到最高。我做过一个广告文案批量生成的项目就是用大batch size在夜间跑一晚上能处理上百万条成本降了70%以上。低延迟优先场景在线对话机器人、实时翻译、交互式编程助手。用户敲下回车后哪怕多等200毫秒体验都会大打折扣。这时我们必须严格控制batch size。一种常见的策略是动态批处理Dynamic Batching设置一个很小的最大batch size比如4或8并且设置一个等待超时时间例如50ms。系统会收集在这50ms内到达的所有请求凑成一批但最多不超过最大batch size。这样既能利用一定的并行性提升吞吐又能将延迟控制在可接受范围内。3.3 除了Batch Size还有哪些关键因素在实际部署中调整batch size不是唯一的手段它还需要和其他系统参数协同考虑输入长度Sequence Length这是影响内存占用和计算量的最大变量。长序列会显著增加KV Cache的内存消耗和注意力计算的开销。如果你的服务预期会有很长的上下文比如处理长文档那么你能使用的最大batch size会受到GPU显存的严格限制。公式可以粗略估算为显存占用 ≈ 模型参数量 * 2字节半精度 batch_size * seq_len * hidden_size * 2 * num_layers * 2字节。后面一项就是KV Cache的大头。模型精度使用FP16半精度或BF16代替FP32单精度进行推理不仅能减少一半的显存占用还能利用GPU的Tensor Core获得大幅加速。现在主流的推理框架如vLLM, TensorRT-LLM默认都使用半精度。连续批处理Continuous Batching这是更高级的优化。在传统的动态批处理中一个batch里所有序列必须同时开始、同时结束。但生成任务中每个序列生成的长度可能不同。连续批处理允许提前结束的序列释放资源并让新来的序列加入计算极大地提高了GPU利用率。vLLM和TGIText Generation Inference等框架的核心优势就在于此。硬件选择GPU的显存大小和内存带宽直接决定了batch size的上限。对于追求高吞吐的离线场景显存大的卡如80GB的A100/H100是首选。对于低延迟在线场景计算能力强、单精度性能高的卡可能更有优势。4. 高级技巧与实战避坑指南掌握了基本原理和权衡后我们来看看一些能让你在实际项目中如虎添翼的高级技巧以及我踩过的一些坑。4.1 注意力掩码与位置编码的坑在批处理中注意力掩码attention_mask和位置编码position_ids的处理必须格外小心否则会导致生成结果 nonsense。坑1掩码错误导致模型“看到”未来信息。在因果语言模型如GPT的自注意力中每个token只能看到它自己和之前的token。我们的attention_mask需要是下三角矩阵包括对角线为1。在批处理生成时每次迭代我们都在attention_mask右侧添加一列1这隐式地构建了一个正确的因果掩码。但如果你是自己从头构建掩码务必检查这一点。坑2位置ID未随生成更新。像GPT-2这样的模型需要显式的position_ids来告诉模型每个token在序列中的绝对位置。在生成过程中必须为每个新生成的token更新position_ids通常是最后一个位置ID加1。如果忘记更新或者更新错误模型就会对token的位置产生混淆生成乱码。我们的代码中current_inputs[‘position_ids’][:, -1:] 1就是在做这件事。坑3填充token的位置ID。对于左侧填充的序列填充部分的位置ID应该被设置成一个安全值比如1或0而不是默认的-1或很大的数否则在嵌入层可能会引发索引错误。我们之前用position_ids.masked_fill_(attention_mask 0, 1)就是为了解决这个问题。4.2 内存管理与显存优化实战当batch size和序列长度上去之后显存爆炸是家常便饭。除了选用更大的显卡我们还可以从软件层面优化。策略1使用PagedAttention和KV Cache量化。vLLM框架提出的PagedAttention技术像操作系统管理内存一样管理KV Cache能极大减少由于碎片化造成的内存浪费。此外将KV Cache从FP16量化到INT8甚至FP4可以再节省50%-75%的显存这对支持大batch和长上下文至关重要。不过量化会带来轻微的质量损失需要仔细评估。策略2激活重计算Activation Checkpointing。在非常深的模型中进行长序列推理时中间激活值也会占用大量显存。激活重计算是一种用时间换空间的技术它不保存所有层的激活值而是在反向传播或某些需要激活值的推理步骤时临时重新计算它们。虽然会增加计算量但能显著降低峰值显存。策略3流式输出与迭代级调度。对于在线服务不要等整个序列生成完再一次性返回。采用流式输出Server-Sent Events每生成一个或几个token就立刻返回给客户端。这不仅能降低用户感知的延迟还能让服务端更早地释放部分资源。结合连续批处理可以实现更细粒度的资源调度。4.3 监控与自适应策略在生产环境中请求流量是波动的。固定一个batch size可能不是最优解。一个更聪明的做法是实现自适应的批处理策略。你可以监控两个核心指标请求队列长度和GPU利用率。当队列中积压的请求很多而GPU利用率不高时可以适当增加最大batch size或延长等待超时时间以提升吞吐量。当队列很短但延迟要求很高时则立即减少batch size甚至对高优先级请求采用“插队”的方式立即处理。此外还可以根据请求的元信息进行智能分桶Binning。例如将相似长度的请求放在同一个batch里可以减少填充带来的计算浪费。将付费用户和免费用户的请求放入不同的处理队列并分配不同的批处理策略也是一种常见的业务优化。5. 总结与个人心得批处理技术是大模型推理优化的基石它直接决定了你的硬件资源能否被“吃干榨净”。从简单的静态批处理到复杂的动态连续批处理其演进的核心思想始终是在保证服务质量延迟的前提下最大化硬件利用率吞吐。回顾我这几年部署各种模型服务的经历最大的体会就是“没有银弹”。为一个内部知识库QA服务选择的batch size32的配置直接套用到面向C端的聊天机器人上肯定会引发用户投诉。你必须深入理解自己的业务场景你的用户能容忍多长的等待时间你的请求是均匀到来还是突发性的你的GPU是独占还是共享我建议大家在项目初期就建立一个简单的性能测试框架像我们上面做的那样绘制出你的模型在目标硬件上的吞吐量-延迟曲线。这张图会成为你所有容量规划、资源采购和成本评估的决策依据。然后在真实流量下进行压测观察系统在负载下的表现特别是长尾延迟P99 Latency这往往比平均延迟更能反映用户体验。最后不要重复造轮子。除非有极其特殊的需求否则我强烈建议使用成熟的推理框架如vLLM、TGI或TensorRT-LLM。它们已经集成了动态批处理、连续批处理、PagedAttention、量化等几乎所有我们讨论过的优化技术并且经过了大规模生产的验证。把精力从底层优化中解放出来更多地投入到业务逻辑和模型效果本身可能是工程师更高效的选择。批处理是一门权衡的艺术理解它驾驭它你就能在成本与体验之间找到属于自己的那个最佳平衡点。