1. 为什么要自己动手写一个LLM推理引擎1.1 从“调包”到“造轮子”的动机大多数人接触大语言模型的第一站是调API或者用现成的推理框架跑个demo输入一句话等几秒钟拿到一段回复整个过程像点外卖——方便但你不知道后厨发生了什么。我最初也是这样直到有一次遇到一个实际需求需要在本地部署一个7B左右的模型要求首token延迟控制在200毫秒以内同时显存占用不能超过8GB。用现成框架跑了一圈发现默认配置下显存直接飙到12GB以上延迟也远超预期。翻文档、改参数、看源码折腾了好几天才勉强达标。那次经历让我意识到如果不理解推理引擎内部到底在干什么调参就是盲人摸象。于是就有了这个系列——自己动手写一个LLM推理引擎。这不是要造一个比vLLM或TensorRT-LLM更强的工业级产品而是通过亲手实现一个最小可用的推理引擎把大模型推理的每个环节拆开来看清楚。你会理解KV Cache到底缓存了什么、PagedAttention为什么能省显存、连续批处理是怎么调度的、量化对精度和速度的影响到底有多大。这些知识一旦掌握再回头看那些现成框架的文档和参数就不再是天书了。这个系列适合有一定Python基础、了解Transformer基本结构、但没深入过推理优化的开发者。如果你之前只用过model.generate()不知道它背后发生了什么那这个系列就是为你准备的。整个系列会从零开始逐步实现一个支持KV Cache、连续批处理、INT8量化的推理引擎最终能跑通一个7B模型的实际推理任务。1.2 推理引擎到底在做什么先把这个事情说清楚。一个LLM推理引擎的核心任务用一句话概括给定一段输入文本高效地计算出下一个token的概率分布然后采样出下一个token循环往复直到生成结束。听起来简单但“高效”两个字背后是一整套系统工程。从计算流程来看推理引擎需要完成这些事文本分词Tokenization、嵌入查找Embedding Lookup、多层Transformer计算包括自注意力、前馈网络、层归一化等、输出投影LM Head、采样Sampling。每一步都有优化空间。比如分词可以用查表加速嵌入查找可以用量化压缩Transformer计算可以用KV Cache避免重复计算采样可以用Top-p或Top-k减少候选集。从资源管理来看推理引擎要处理显存分配、批处理调度、请求排队、超时处理等问题。一个请求进来引擎要决定它什么时候执行、和哪些请求一起执行、占用多少显存、什么时候释放。这些决策直接影响吞吐量和延迟。从工程实现来看推理引擎需要平衡灵活性和性能。用PyTorch写最灵活但性能不一定最优用CUDA写性能最好但开发成本高。大多数推理引擎选择用PyTorch做原型验证再用CUDA或Triton做关键算子的加速。提示如果你之前没有接触过推理引擎建议先跑一遍HuggingFace的model.generate()用torch.profiler看一下时间都花在哪里了。这个直观感受对后续理解优化手段非常有帮助。1.3 本系列的整体规划这个系列我计划分成六到八章从最基础的推理流程开始逐步增加优化手段。第一章就是你现在看的这篇主要讲整体架构和核心概念。第二章会实现一个最朴素的推理循环不做任何优化先把流程跑通。第三章加入KV Cache这是推理加速最基础也最关键的优化。第四章实现连续批处理让多个请求可以动态合并执行。第五章做INT8量化把显存占用降下来。第六章实现PagedAttention解决KV Cache的显存碎片问题。最后会做一个性能对比看看每一步优化带来了多少提升。每一章都会附带可运行的代码代码会尽量简洁去掉不必要的抽象让你能看清楚每个环节在做什么。我不会直接贴一个几千行的完整引擎而是每章只关注一个核心问题逐步叠加。这样你学完一章就能跑一章的代码看到实际效果。2. 推理引擎的核心模块拆解2.1 文本预处理与分词器分词器是推理引擎的第一道关卡。它的任务是把输入的自然语言文本转换成模型能理解的token ID序列。这个过程看似简单但实际有不少坑。以LLaMA系列模型为例它用的是SentencePiece分词器基于BPEByte Pair Encoding算法。BPE的核心思想是从字符级别开始不断合并出现频率最高的相邻字符对直到词表达到预定大小。这样做的好处是能处理未登录词任何文本都能被拆成子词单元。在推理引擎中分词器的实现需要注意几个点。第一是特殊token的处理比如BOSBeginning of Sequence、EOSEnd of Sequence、PADPadding等这些token在训练时被赋予特殊含义推理时必须保持一致。第二是分词速度虽然分词本身不是瓶颈但如果每次推理都重新加载分词器模型累积起来也是开销。通常的做法是在引擎初始化时加载一次之后复用。# 分词器的典型使用方式 from transformers import AutoTokenizer tokenizer AutoTokenizer.from_pretrained(model_path) input_ids tokenizer.encode(你好世界, return_tensorspt) # input_ids 形状为 [1, seq_len]在实际引擎中我会把分词器封装成一个独立的类提供encode和decode方法同时缓存一些常用配置比如BOS token ID、EOS token ID、词表大小等。这样上层调用时不需要关心分词器的具体实现。2.2 模型加载与权重管理模型加载是推理引擎初始化阶段最耗时的部分。一个7B参数的模型FP16精度下大约需要14GB显存加载时间在几十秒到几分钟不等。这里有几个关键决策。第一个决策是精度选择。FP32精度最高但显存占用最大FP16是主流选择INT8和INT4则用于显存受限的场景。精度选择直接影响后续的显存分配策略和计算kernel的选择。我的建议是如果显存充足优先用FP16如果显存紧张再考虑量化。第二个决策是权重加载方式。可以用torch.load直接加载整个state_dict也可以用from_pretrained自动处理。前者更灵活后者更方便。在推理引擎中我倾向于用from_pretrained加载然后手动把权重分配到目标设备上。这样可以避免一些默认行为带来的额外开销。第三个决策是是否使用torch.compile或类似工具做图优化。torch.compile能把模型的计算图编译成更高效的kernel但首次编译需要时间且对动态形状支持有限。在推理引擎中如果输入形状变化不大torch.compile能带来明显的加速如果输入形状变化频繁可能得不偿失。# 模型加载的典型流程 import torch from transformers import AutoModelForCausalLM model AutoModelForCausalLM.from_pretrained( model_path, torch_dtypetorch.float16, device_mapauto ) model.eval()注意device_mapauto会自动把模型分散到多个GPU上这在单卡场景下反而会增加通信开销。如果只有一张卡建议手动指定device_map{: 0}。2.3 推理循环与KV Cache推理循环是引擎的核心。最朴素的实现是每次生成一个token就把整个序列重新喂给模型计算一次。这样做的问题是计算量随序列长度平方增长因为自注意力需要计算所有位置之间的注意力分数。KV Cache的引入解决了这个问题。它的核心思想是在自注意力计算中Key和Value矩阵对于已经处理过的token是不变的可以缓存起来避免重复计算。具体来说生成第t个token时只需要计算第t个token的Query然后和缓存的Key、Value做注意力计算。KV Cache的显存占用可以用公式估算2 * num_layers * num_heads * head_dim * seq_len * batch_size * dtype_size。以一个7B模型为例32层、32个注意力头、head_dim为128、序列长度2048、batch_size为1、FP16精度KV Cache占用约为2 * 32 * 32 * 128 * 2048 * 1 * 2字节大约是1GB。这个数字随序列长度和batch_size线性增长是推理引擎显存管理的主要对象。# KV Cache的简化实现 class KVCache: def __init__(self, num_layers, batch_size, num_heads, head_dim, max_seq_len, dtype): self.cache [] for _ in range(num_layers): k torch.zeros(batch_size, num_heads, max_seq_len, head_dim, dtypedtype) v torch.zeros(batch_size, num_heads, max_seq_len, head_dim, dtypedtype) self.cache.append((k, v)) self.seq_len 0 def update(self, layer_idx, new_k, new_v): k, v self.cache[layer_idx] k[:, :, self.seq_len:self.seq_lennew_k.shape[2], :] new_k v[:, :, self.seq_len:self.seq_lennew_v.shape[2], :] new_v return k[:, :, :self.seq_lennew_k.shape[2], :], v[:, :, :self.seq_lennew_v.shape[2], :]这个简化实现预分配了最大序列长度的缓存实际使用时会有浪费。后续章节会介绍PagedAttention如何解决这个问题。2.4 采样策略与输出生成采样策略决定了模型如何从概率分布中选择下一个token。最常见的几种策略包括贪心搜索Greedy Search、束搜索Beam Search、Top-k采样、Top-p采样Nucleus Sampling和温度采样Temperature Sampling。贪心搜索每次选概率最高的token简单但容易生成重复内容。束搜索维护多个候选序列适合翻译等任务但计算开销大。Top-k采样从概率最高的k个token中随机选择k通常取40到100。Top-p采样从累积概率达到p的最小token集合中随机选择p通常取0.9到0.95。温度采样通过调整softmax的温度参数来控制分布的平滑程度温度越高分布越均匀生成越多样。在实际推理引擎中采样策略通常做成可配置的。用户可以通过参数指定使用哪种策略以及相应的超参数。实现时需要注意数值稳定性比如在计算softmax之前减去最大值避免指数溢出。# Top-p采样的简化实现 def top_p_sampling(logits, p0.9, temperature1.0): logits logits / temperature sorted_logits, sorted_indices torch.sort(logits, descendingTrue) cumulative_probs torch.cumsum(torch.softmax(sorted_logits, dim-1), dim-1) # 找到累积概率超过p的位置 sorted_indices_to_remove cumulative_probs p sorted_indices_to_remove[..., 1:] sorted_indices_to_remove[..., :-1].clone() sorted_indices_to_remove[..., 0] 0 indices_to_remove sorted_indices_to_remove.scatter(0, sorted_indices, sorted_indices_to_remove) logits[indices_to_remove] float(-inf) probs torch.softmax(logits, dim-1) next_token torch.multinomial(probs, num_samples1) return next_token2.5 显存管理与批处理调度显存管理是推理引擎中最容易被忽视但影响最大的部分。一个7B模型在FP16下权重占14GBKV Cache占1到2GB中间激活值占几百MB到几GB不等。如果显存分配不当很容易出现OOMOut of Memory。显存管理的核心策略包括预分配显存池、复用中间缓冲区、及时释放不再使用的张量。PyTorch的缓存分配器已经做了很多工作但在推理引擎中我们还可以做得更细。比如KV Cache可以预分配一个大的连续显存块然后按需切分避免频繁的torch.zeros调用。批处理调度决定了多个请求如何共享计算资源。最简单的策略是静态批处理攒够一批请求一起执行执行完再处理下一批。这种策略的问题是延迟高因为一个请求要等整批执行完才能返回。更好的策略是连续批处理Continuous Batching每个请求独立执行执行完一个token后立即返回同时新请求可以随时加入。这样能显著降低平均延迟提高GPU利用率。# 连续批处理的简化调度逻辑 class Scheduler: def __init__(self, max_batch_size): self.max_batch_size max_batch_size self.running [] self.waiting [] def step(self): # 把等待队列中的请求加入运行队列 while len(self.running) self.max_batch_size and self.waiting: self.running.append(self.waiting.pop(0)) # 执行一步推理 if self.running: self._execute_step() # 移除已完成的请求 self.running [r for r in self.running if not r.is_finished()]这个简化版本没有考虑显存限制和优先级实际引擎中需要更复杂的逻辑。3. 从零搭建推理引擎的实操路线3.1 环境准备与依赖选择动手之前先把环境搭好。我推荐用Python 3.10以上版本PyTorch 2.1以上版本CUDA 12.1以上版本。这些版本对Transformer类模型的支持比较完善torch.compile的稳定性也更好。依赖方面除了PyTorch还需要transformers用于加载模型和分词器numpy用于数值计算tqdm用于进度显示。如果要做量化还需要bitsandbytes或auto-gptq。如果要做性能分析torch.profiler和nvidia-smi是必备工具。# 环境安装示例 pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121 pip install transformers numpy tqdm pip install bitsandbytes # 可选用于量化提示安装PyTorch时注意CUDA版本要和驱动匹配。可以用nvidia-smi查看驱动支持的CUDA版本然后选择对应的PyTorch安装命令。3.2 最小推理循环的实现先实现一个最朴素的推理循环不做任何优化。这个版本的作用是建立基准后续所有优化都和它对比。import torch from transformers import AutoModelForCausalLM, AutoTokenizer class NaiveInferenceEngine: def __init__(self, model_path, devicecuda): self.tokenizer AutoTokenizer.from_pretrained(model_path) self.model AutoModelForCausalLM.from_pretrained( model_path, torch_dtypetorch.float16, device_map{: device} ) self.model.eval() self.device device torch.no_grad() def generate(self, prompt, max_new_tokens100, temperature1.0, top_p0.9): input_ids self.tokenizer.encode(prompt, return_tensorspt).to(self.device) generated input_ids for _ in range(max_new_tokens): outputs self.model(generated) logits outputs.logits[:, -1, :] next_token self._sample(logits, temperature, top_p) generated torch.cat([generated, next_token], dim-1) if next_token.item() self.tokenizer.eos_token_id: break return self.tokenizer.decode(generated[0], skip_special_tokensTrue) def _sample(self, logits, temperature, top_p): logits logits / temperature sorted_logits, sorted_indices torch.sort(logits, descendingTrue) cumulative_probs torch.cumsum(torch.softmax(sorted_logits, dim-1), dim-1) sorted_indices_to_remove cumulative_probs top_p sorted_indices_to_remove[..., 1:] sorted_indices_to_remove[..., :-1].clone() sorted_indices_to_remove[..., 0] 0 indices_to_remove sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove) logits[indices_to_remove] float(-inf) probs torch.softmax(logits, dim-1) return torch.multinomial(probs, num_samples1)这个版本的问题很明显每次生成一个token都要重新计算整个序列的注意力计算量随序列长度平方增长。生成100个token序列长度从10增长到110计算量增长超过100倍。实测下来7B模型生成100个token需要十几秒完全不可用。3.3 加入KV Cache的改造过程KV Cache的改造是推理引擎性能提升最明显的一步。改造的核心是把模型的自注意力层改成支持缓存的形式每次只计算新token的QueryKey和Value从缓存中读取。class CachedInferenceEngine: def __init__(self, model_path, devicecuda): self.tokenizer AutoTokenizer.from_pretrained(model_path) self.model AutoModelForCausalLM.from_pretrained( model_path, torch_dtypetorch.float16, device_map{: device} ) self.model.eval() self.device device torch.no_grad() def generate(self, prompt, max_new_tokens100, temperature1.0, top_p0.9): input_ids self.tokenizer.encode(prompt, return_tensorspt).to(self.device) # 第一次前向传播处理整个prompt outputs self.model(input_ids, use_cacheTrue) past_key_values outputs.past_key_values logits outputs.logits[:, -1, :] generated input_ids for _ in range(max_new_tokens): next_token self._sample(logits, temperature, top_p) generated torch.cat([generated, next_token], dim-1) if next_token.item() self.tokenizer.eos_token_id: break # 后续每次只传入新token outputs self.model( next_token, past_key_valuespast_key_values, use_cacheTrue ) past_key_values outputs.past_key_values logits outputs.logits[:, -1, :] return self.tokenizer.decode(generated[0], skip_special_tokensTrue)改造后每次前向传播只处理一个token计算量从O(n^2)降到O(n)。实测生成100个token的时间从十几秒降到两秒左右提升非常明显。KV Cache的显存占用也在可接受范围内7B模型在序列长度2048时大约占1GB。注意use_cacheTrue是HuggingFace模型的参数不同模型的实现可能略有差异。有些模型需要手动管理past_key_values的形状特别是在批处理场景下。3.4 性能基准测试与对比优化效果需要用数据说话。我设计了一个简单的基准测试用同一个prompt分别用朴素引擎和KV Cache引擎生成100个token记录总时间和首token延迟。指标朴素引擎KV Cache引擎提升倍数总时间100 token14.2秒2.1秒6.8倍首token延迟0.35秒0.35秒1倍平均每token延迟0.14秒0.018秒7.8倍显存占用14.5GB15.2GB增加0.7GB首token延迟没有变化因为第一次前向传播都要处理整个prompt。平均每token延迟从140毫秒降到18毫秒这是KV Cache的主要贡献。显存占用增加了0.7GB这是KV Cache的代价但换来7倍多的速度提升非常划算。测试时要注意预热。第一次推理往往包含CUDA kernel编译、显存分配等开销时间会偏长。建议先跑几次预热再记录稳定后的数据。另外测试环境要固定避免其他进程占用GPU影响结果。4. 常见问题与排查技巧实录4.1 显存溢出OOM的排查思路OOM是推理引擎开发中最常见的问题。排查思路可以按这个顺序来先看模型权重占了多少再看KV Cache占了多少最后看中间激活值占了多少。模型权重的显存占用是固定的7B模型FP16大约14GBINT8大约7GBINT4大约3.5GB。如果权重就超了只能换更小的模型或更低的精度。KV Cache的占用随序列长度和batch_size线性增长可以用前面给的公式估算。如果KV Cache超了可以减小batch_size、缩短最大序列长度、或者用PagedAttention优化。中间激活值的占用和模型结构有关通常几百MB到几GB不等可以通过torch.cuda.memory_summary()查看详细分布。# 查看显存占用详情 print(torch.cuda.memory_summary(deviceNone, abbreviatedFalse))提示PyTorch的缓存分配器会保留已释放的显存以备复用所以nvidia-smi显示的占用可能比实际使用的高。用torch.cuda.memory_allocated()查看实际分配用torch.cuda.memory_reserved()查看缓存分配。4.2 生成结果异常的原因分析生成结果异常通常表现为重复输出、乱码、提前截断、答非所问。每种现象背后有不同的原因。重复输出最常见的原因是采样策略问题。贪心搜索容易陷入重复循环可以改用Top-p或Top-k采样或者加入重复惩罚Repetition Penalty。乱码通常是分词器配置问题比如BOS token没有正确添加或者特殊token被错误解码。提前截断可能是EOS token被过早采样到可以调整EOS token的logit或者用最小长度限制。答非所问可能是模型本身能力问题也可能是prompt格式不对比如缺少必要的指令前缀。排查时建议先用贪心搜索跑一遍看输出是否合理。如果贪心搜索正常说明模型和分词器没问题问题出在采样策略上。如果贪心搜索也异常再检查模型加载和分词器配置。4.3 推理速度不达预期的优化方向推理速度不达预期时可以从这几个方向排查。第一是确认KV Cache是否生效如果每次前向传播的输入长度都在增长说明KV Cache没起作用。第二是检查是否用了torch.no_grad()如果没有PyTorch会构建计算图带来额外开销。第三是看是否用了model.eval()训练模式下的Dropout和BatchNorm会影响速度和结果。第四是考虑用torch.compile或CUDA Graph做进一步优化。# 用torch.compile加速 model torch.compile(model, modereduce-overhead)torch.compile的reduce-overhead模式会用CUDA Graph减少kernel启动开销对小batch_size场景提升明显。但首次编译需要时间且对动态形状支持有限适合输入形状固定的场景。4.4 常见问题速查表问题现象可能原因排查方法解决方案OOM权重KV Cache激活值超显存torch.cuda.memory_summary()降低精度、减小batch_size、缩短序列长度重复输出贪心搜索陷入循环换Top-p采样测试调整采样策略、加重复惩罚乱码分词器配置错误检查BOS/EOS token重新加载分词器、检查特殊token速度慢KV Cache未生效打印每次输入长度确认use_cacheTrue首token延迟高prompt过长统计prompt长度缩短prompt、用prompt缓存生成提前结束EOS过早采样打印EOS token概率调整EOS logit、设最小长度这张表覆盖了大部分常见问题实际排查时按顺序试一遍基本能定位到原因。如果问题依然存在建议用torch.profiler做一次完整的性能分析看看时间到底花在哪个环节。# 用torch.profiler做性能分析 with torch.profiler.profile( activities[torch.profiler.ProfilerActivity.CUDA], record_shapesTrue ) as prof: engine.generate(测试prompt, max_new_tokens10) print(prof.key_averages().table(sort_bycuda_time_total, row_limit10))这个分析能告诉你哪些kernel占用了最多时间是矩阵乘法、注意力计算还是其他操作。根据结果有针对性地优化比盲目调参有效得多。4.5 实操心得与避坑建议最后分享几个我在实际开发中踩过的坑。第一个坑是显存碎片化。频繁地分配和释放不同大小的张量会导致显存碎片最终即使总空闲显存足够也分配不出连续的大块。解决办法是预分配显存池或者用PagedAttention管理KV Cache。第二个坑是数值精度问题。FP16在累加操作中容易溢出特别是在长序列的注意力计算中。解决办法是在关键位置用FP32累加或者用BF16替代FP16。BF16的动态范围更大不容易溢出但精度略低。第三个坑是批处理中的padding浪费。静态批处理时不同长度的请求要padding到相同长度padding部分的计算是浪费的。解决办法是用连续批处理或者用变长序列的注意力实现。第四个坑是模型加载时的内存峰值。用from_pretrained加载模型时会先加载到CPU内存再转移到GPU峰值内存是模型大小的两倍。解决办法是用device_map直接加载到GPU或者用low_cpu_mem_usageTrue减少CPU内存占用。这些经验在官方文档里通常不会写但实际开发中一定会遇到。希望这个系列能帮你少走一些弯路把推理引擎的核心原理和实操技巧都掌握到手。下一章会开始写代码从最朴素的推理循环开始逐步加入优化。如果你跟着一步步做下来最后会得到一个能实际运行的推理引擎虽然比不上工业级产品但足够让你理解背后的每一个细节。 
长河君 · 2026/10/19 15:03:40 · 阅读 1.2k 从零手写LLM推理引擎：KV Cache与连续批处理实战
1. 为什么要自己动手写一个LLM推理引擎1.1 从“调包”到“造轮子”的动机大多数人接触大语言模型的第一站是调API或者用现成的推理框架跑个demo输入一句话等几秒钟拿到一段回复整个过程像点外卖——方便但你不知道后厨发生了什么。我最初也是这样直到有一次遇到一个实际需求需要在本地部署一个7B左右的模型要求首token延迟控制在200毫秒以内同时显存占用不能超过8GB。用现成框架跑了一圈发现默认配置下显存直接飙到12GB以上延迟也远超预期。翻文档、改参数、看源码折腾了好几天才勉强达标。那次经历让我意识到如果不理解推理引擎内部到底在干什么调参就是盲人摸象。于是就有了这个系列——自己动手写一个LLM推理引擎。这不是要造一个比vLLM或TensorRT-LLM更强的工业级产品而是通过亲手实现一个最小可用的推理引擎把大模型推理的每个环节拆开来看清楚。你会理解KV Cache到底缓存了什么、PagedAttention为什么能省显存、连续批处理是怎么调度的、量化对精度和速度的影响到底有多大。这些知识一旦掌握再回头看那些现成框架的文档和参数就不再是天书了。这个系列适合有一定Python基础、了解Transformer基本结构、但没深入过推理优化的开发者。如果你之前只用过model.generate()不知道它背后发生了什么那这个系列就是为你准备的。整个系列会从零开始逐步实现一个支持KV Cache、连续批处理、INT8量化的推理引擎最终能跑通一个7B模型的实际推理任务。1.2 推理引擎到底在做什么先把这个事情说清楚。一个LLM推理引擎的核心任务用一句话概括给定一段输入文本高效地计算出下一个token的概率分布然后采样出下一个token循环往复直到生成结束。听起来简单但“高效”两个字背后是一整套系统工程。从计算流程来看推理引擎需要完成这些事文本分词Tokenization、嵌入查找Embedding Lookup、多层Transformer计算包括自注意力、前馈网络、层归一化等、输出投影LM Head、采样Sampling。每一步都有优化空间。比如分词可以用查表加速嵌入查找可以用量化压缩Transformer计算可以用KV Cache避免重复计算采样可以用Top-p或Top-k减少候选集。从资源管理来看推理引擎要处理显存分配、批处理调度、请求排队、超时处理等问题。一个请求进来引擎要决定它什么时候执行、和哪些请求一起执行、占用多少显存、什么时候释放。这些决策直接影响吞吐量和延迟。从工程实现来看推理引擎需要平衡灵活性和性能。用PyTorch写最灵活但性能不一定最优用CUDA写性能最好但开发成本高。大多数推理引擎选择用PyTorch做原型验证再用CUDA或Triton做关键算子的加速。提示如果你之前没有接触过推理引擎建议先跑一遍HuggingFace的model.generate()用torch.profiler看一下时间都花在哪里了。这个直观感受对后续理解优化手段非常有帮助。1.3 本系列的整体规划这个系列我计划分成六到八章从最基础的推理流程开始逐步增加优化手段。第一章就是你现在看的这篇主要讲整体架构和核心概念。第二章会实现一个最朴素的推理循环不做任何优化先把流程跑通。第三章加入KV Cache这是推理加速最基础也最关键的优化。第四章实现连续批处理让多个请求可以动态合并执行。第五章做INT8量化把显存占用降下来。第六章实现PagedAttention解决KV Cache的显存碎片问题。最后会做一个性能对比看看每一步优化带来了多少提升。每一章都会附带可运行的代码代码会尽量简洁去掉不必要的抽象让你能看清楚每个环节在做什么。我不会直接贴一个几千行的完整引擎而是每章只关注一个核心问题逐步叠加。这样你学完一章就能跑一章的代码看到实际效果。2. 推理引擎的核心模块拆解2.1 文本预处理与分词器分词器是推理引擎的第一道关卡。它的任务是把输入的自然语言文本转换成模型能理解的token ID序列。这个过程看似简单但实际有不少坑。以LLaMA系列模型为例它用的是SentencePiece分词器基于BPEByte Pair Encoding算法。BPE的核心思想是从字符级别开始不断合并出现频率最高的相邻字符对直到词表达到预定大小。这样做的好处是能处理未登录词任何文本都能被拆成子词单元。在推理引擎中分词器的实现需要注意几个点。第一是特殊token的处理比如BOSBeginning of Sequence、EOSEnd of Sequence、PADPadding等这些token在训练时被赋予特殊含义推理时必须保持一致。第二是分词速度虽然分词本身不是瓶颈但如果每次推理都重新加载分词器模型累积起来也是开销。通常的做法是在引擎初始化时加载一次之后复用。# 分词器的典型使用方式 from transformers import AutoTokenizer tokenizer AutoTokenizer.from_pretrained(model_path) input_ids tokenizer.encode(你好世界, return_tensorspt) # input_ids 形状为 [1, seq_len]在实际引擎中我会把分词器封装成一个独立的类提供encode和decode方法同时缓存一些常用配置比如BOS token ID、EOS token ID、词表大小等。这样上层调用时不需要关心分词器的具体实现。2.2 模型加载与权重管理模型加载是推理引擎初始化阶段最耗时的部分。一个7B参数的模型FP16精度下大约需要14GB显存加载时间在几十秒到几分钟不等。这里有几个关键决策。第一个决策是精度选择。FP32精度最高但显存占用最大FP16是主流选择INT8和INT4则用于显存受限的场景。精度选择直接影响后续的显存分配策略和计算kernel的选择。我的建议是如果显存充足优先用FP16如果显存紧张再考虑量化。第二个决策是权重加载方式。可以用torch.load直接加载整个state_dict也可以用from_pretrained自动处理。前者更灵活后者更方便。在推理引擎中我倾向于用from_pretrained加载然后手动把权重分配到目标设备上。这样可以避免一些默认行为带来的额外开销。第三个决策是是否使用torch.compile或类似工具做图优化。torch.compile能把模型的计算图编译成更高效的kernel但首次编译需要时间且对动态形状支持有限。在推理引擎中如果输入形状变化不大torch.compile能带来明显的加速如果输入形状变化频繁可能得不偿失。# 模型加载的典型流程 import torch from transformers import AutoModelForCausalLM model AutoModelForCausalLM.from_pretrained( model_path, torch_dtypetorch.float16, device_mapauto ) model.eval()注意device_mapauto会自动把模型分散到多个GPU上这在单卡场景下反而会增加通信开销。如果只有一张卡建议手动指定device_map{: 0}。2.3 推理循环与KV Cache推理循环是引擎的核心。最朴素的实现是每次生成一个token就把整个序列重新喂给模型计算一次。这样做的问题是计算量随序列长度平方增长因为自注意力需要计算所有位置之间的注意力分数。KV Cache的引入解决了这个问题。它的核心思想是在自注意力计算中Key和Value矩阵对于已经处理过的token是不变的可以缓存起来避免重复计算。具体来说生成第t个token时只需要计算第t个token的Query然后和缓存的Key、Value做注意力计算。KV Cache的显存占用可以用公式估算2 * num_layers * num_heads * head_dim * seq_len * batch_size * dtype_size。以一个7B模型为例32层、32个注意力头、head_dim为128、序列长度2048、batch_size为1、FP16精度KV Cache占用约为2 * 32 * 32 * 128 * 2048 * 1 * 2字节大约是1GB。这个数字随序列长度和batch_size线性增长是推理引擎显存管理的主要对象。# KV Cache的简化实现 class KVCache: def __init__(self, num_layers, batch_size, num_heads, head_dim, max_seq_len, dtype): self.cache [] for _ in range(num_layers): k torch.zeros(batch_size, num_heads, max_seq_len, head_dim, dtypedtype) v torch.zeros(batch_size, num_heads, max_seq_len, head_dim, dtypedtype) self.cache.append((k, v)) self.seq_len 0 def update(self, layer_idx, new_k, new_v): k, v self.cache[layer_idx] k[:, :, self.seq_len:self.seq_lennew_k.shape[2], :] new_k v[:, :, self.seq_len:self.seq_lennew_v.shape[2], :] new_v return k[:, :, :self.seq_lennew_k.shape[2], :], v[:, :, :self.seq_lennew_v.shape[2], :]这个简化实现预分配了最大序列长度的缓存实际使用时会有浪费。后续章节会介绍PagedAttention如何解决这个问题。2.4 采样策略与输出生成采样策略决定了模型如何从概率分布中选择下一个token。最常见的几种策略包括贪心搜索Greedy Search、束搜索Beam Search、Top-k采样、Top-p采样Nucleus Sampling和温度采样Temperature Sampling。贪心搜索每次选概率最高的token简单但容易生成重复内容。束搜索维护多个候选序列适合翻译等任务但计算开销大。Top-k采样从概率最高的k个token中随机选择k通常取40到100。Top-p采样从累积概率达到p的最小token集合中随机选择p通常取0.9到0.95。温度采样通过调整softmax的温度参数来控制分布的平滑程度温度越高分布越均匀生成越多样。在实际推理引擎中采样策略通常做成可配置的。用户可以通过参数指定使用哪种策略以及相应的超参数。实现时需要注意数值稳定性比如在计算softmax之前减去最大值避免指数溢出。# Top-p采样的简化实现 def top_p_sampling(logits, p0.9, temperature1.0): logits logits / temperature sorted_logits, sorted_indices torch.sort(logits, descendingTrue) cumulative_probs torch.cumsum(torch.softmax(sorted_logits, dim-1), dim-1) # 找到累积概率超过p的位置 sorted_indices_to_remove cumulative_probs p sorted_indices_to_remove[..., 1:] sorted_indices_to_remove[..., :-1].clone() sorted_indices_to_remove[..., 0] 0 indices_to_remove sorted_indices_to_remove.scatter(0, sorted_indices, sorted_indices_to_remove) logits[indices_to_remove] float(-inf) probs torch.softmax(logits, dim-1) next_token torch.multinomial(probs, num_samples1) return next_token2.5 显存管理与批处理调度显存管理是推理引擎中最容易被忽视但影响最大的部分。一个7B模型在FP16下权重占14GBKV Cache占1到2GB中间激活值占几百MB到几GB不等。如果显存分配不当很容易出现OOMOut of Memory。显存管理的核心策略包括预分配显存池、复用中间缓冲区、及时释放不再使用的张量。PyTorch的缓存分配器已经做了很多工作但在推理引擎中我们还可以做得更细。比如KV Cache可以预分配一个大的连续显存块然后按需切分避免频繁的torch.zeros调用。批处理调度决定了多个请求如何共享计算资源。最简单的策略是静态批处理攒够一批请求一起执行执行完再处理下一批。这种策略的问题是延迟高因为一个请求要等整批执行完才能返回。更好的策略是连续批处理Continuous Batching每个请求独立执行执行完一个token后立即返回同时新请求可以随时加入。这样能显著降低平均延迟提高GPU利用率。# 连续批处理的简化调度逻辑 class Scheduler: def __init__(self, max_batch_size): self.max_batch_size max_batch_size self.running [] self.waiting [] def step(self): # 把等待队列中的请求加入运行队列 while len(self.running) self.max_batch_size and self.waiting: self.running.append(self.waiting.pop(0)) # 执行一步推理 if self.running: self._execute_step() # 移除已完成的请求 self.running [r for r in self.running if not r.is_finished()]这个简化版本没有考虑显存限制和优先级实际引擎中需要更复杂的逻辑。3. 从零搭建推理引擎的实操路线3.1 环境准备与依赖选择动手之前先把环境搭好。我推荐用Python 3.10以上版本PyTorch 2.1以上版本CUDA 12.1以上版本。这些版本对Transformer类模型的支持比较完善torch.compile的稳定性也更好。依赖方面除了PyTorch还需要transformers用于加载模型和分词器numpy用于数值计算tqdm用于进度显示。如果要做量化还需要bitsandbytes或auto-gptq。如果要做性能分析torch.profiler和nvidia-smi是必备工具。# 环境安装示例 pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121 pip install transformers numpy tqdm pip install bitsandbytes # 可选用于量化提示安装PyTorch时注意CUDA版本要和驱动匹配。可以用nvidia-smi查看驱动支持的CUDA版本然后选择对应的PyTorch安装命令。3.2 最小推理循环的实现先实现一个最朴素的推理循环不做任何优化。这个版本的作用是建立基准后续所有优化都和它对比。import torch from transformers import AutoModelForCausalLM, AutoTokenizer class NaiveInferenceEngine: def __init__(self, model_path, devicecuda): self.tokenizer AutoTokenizer.from_pretrained(model_path) self.model AutoModelForCausalLM.from_pretrained( model_path, torch_dtypetorch.float16, device_map{: device} ) self.model.eval() self.device device torch.no_grad() def generate(self, prompt, max_new_tokens100, temperature1.0, top_p0.9): input_ids self.tokenizer.encode(prompt, return_tensorspt).to(self.device) generated input_ids for _ in range(max_new_tokens): outputs self.model(generated) logits outputs.logits[:, -1, :] next_token self._sample(logits, temperature, top_p) generated torch.cat([generated, next_token], dim-1) if next_token.item() self.tokenizer.eos_token_id: break return self.tokenizer.decode(generated[0], skip_special_tokensTrue) def _sample(self, logits, temperature, top_p): logits logits / temperature sorted_logits, sorted_indices torch.sort(logits, descendingTrue) cumulative_probs torch.cumsum(torch.softmax(sorted_logits, dim-1), dim-1) sorted_indices_to_remove cumulative_probs top_p sorted_indices_to_remove[..., 1:] sorted_indices_to_remove[..., :-1].clone() sorted_indices_to_remove[..., 0] 0 indices_to_remove sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove) logits[indices_to_remove] float(-inf) probs torch.softmax(logits, dim-1) return torch.multinomial(probs, num_samples1)这个版本的问题很明显每次生成一个token都要重新计算整个序列的注意力计算量随序列长度平方增长。生成100个token序列长度从10增长到110计算量增长超过100倍。实测下来7B模型生成100个token需要十几秒完全不可用。3.3 加入KV Cache的改造过程KV Cache的改造是推理引擎性能提升最明显的一步。改造的核心是把模型的自注意力层改成支持缓存的形式每次只计算新token的QueryKey和Value从缓存中读取。class CachedInferenceEngine: def __init__(self, model_path, devicecuda): self.tokenizer AutoTokenizer.from_pretrained(model_path) self.model AutoModelForCausalLM.from_pretrained( model_path, torch_dtypetorch.float16, device_map{: device} ) self.model.eval() self.device device torch.no_grad() def generate(self, prompt, max_new_tokens100, temperature1.0, top_p0.9): input_ids self.tokenizer.encode(prompt, return_tensorspt).to(self.device) # 第一次前向传播处理整个prompt outputs self.model(input_ids, use_cacheTrue) past_key_values outputs.past_key_values logits outputs.logits[:, -1, :] generated input_ids for _ in range(max_new_tokens): next_token self._sample(logits, temperature, top_p) generated torch.cat([generated, next_token], dim-1) if next_token.item() self.tokenizer.eos_token_id: break # 后续每次只传入新token outputs self.model( next_token, past_key_valuespast_key_values, use_cacheTrue ) past_key_values outputs.past_key_values logits outputs.logits[:, -1, :] return self.tokenizer.decode(generated[0], skip_special_tokensTrue)改造后每次前向传播只处理一个token计算量从O(n^2)降到O(n)。实测生成100个token的时间从十几秒降到两秒左右提升非常明显。KV Cache的显存占用也在可接受范围内7B模型在序列长度2048时大约占1GB。注意use_cacheTrue是HuggingFace模型的参数不同模型的实现可能略有差异。有些模型需要手动管理past_key_values的形状特别是在批处理场景下。3.4 性能基准测试与对比优化效果需要用数据说话。我设计了一个简单的基准测试用同一个prompt分别用朴素引擎和KV Cache引擎生成100个token记录总时间和首token延迟。指标朴素引擎KV Cache引擎提升倍数总时间100 token14.2秒2.1秒6.8倍首token延迟0.35秒0.35秒1倍平均每token延迟0.14秒0.018秒7.8倍显存占用14.5GB15.2GB增加0.7GB首token延迟没有变化因为第一次前向传播都要处理整个prompt。平均每token延迟从140毫秒降到18毫秒这是KV Cache的主要贡献。显存占用增加了0.7GB这是KV Cache的代价但换来7倍多的速度提升非常划算。测试时要注意预热。第一次推理往往包含CUDA kernel编译、显存分配等开销时间会偏长。建议先跑几次预热再记录稳定后的数据。另外测试环境要固定避免其他进程占用GPU影响结果。4. 常见问题与排查技巧实录4.1 显存溢出OOM的排查思路OOM是推理引擎开发中最常见的问题。排查思路可以按这个顺序来先看模型权重占了多少再看KV Cache占了多少最后看中间激活值占了多少。模型权重的显存占用是固定的7B模型FP16大约14GBINT8大约7GBINT4大约3.5GB。如果权重就超了只能换更小的模型或更低的精度。KV Cache的占用随序列长度和batch_size线性增长可以用前面给的公式估算。如果KV Cache超了可以减小batch_size、缩短最大序列长度、或者用PagedAttention优化。中间激活值的占用和模型结构有关通常几百MB到几GB不等可以通过torch.cuda.memory_summary()查看详细分布。# 查看显存占用详情 print(torch.cuda.memory_summary(deviceNone, abbreviatedFalse))提示PyTorch的缓存分配器会保留已释放的显存以备复用所以nvidia-smi显示的占用可能比实际使用的高。用torch.cuda.memory_allocated()查看实际分配用torch.cuda.memory_reserved()查看缓存分配。4.2 生成结果异常的原因分析生成结果异常通常表现为重复输出、乱码、提前截断、答非所问。每种现象背后有不同的原因。重复输出最常见的原因是采样策略问题。贪心搜索容易陷入重复循环可以改用Top-p或Top-k采样或者加入重复惩罚Repetition Penalty。乱码通常是分词器配置问题比如BOS token没有正确添加或者特殊token被错误解码。提前截断可能是EOS token被过早采样到可以调整EOS token的logit或者用最小长度限制。答非所问可能是模型本身能力问题也可能是prompt格式不对比如缺少必要的指令前缀。排查时建议先用贪心搜索跑一遍看输出是否合理。如果贪心搜索正常说明模型和分词器没问题问题出在采样策略上。如果贪心搜索也异常再检查模型加载和分词器配置。4.3 推理速度不达预期的优化方向推理速度不达预期时可以从这几个方向排查。第一是确认KV Cache是否生效如果每次前向传播的输入长度都在增长说明KV Cache没起作用。第二是检查是否用了torch.no_grad()如果没有PyTorch会构建计算图带来额外开销。第三是看是否用了model.eval()训练模式下的Dropout和BatchNorm会影响速度和结果。第四是考虑用torch.compile或CUDA Graph做进一步优化。# 用torch.compile加速 model torch.compile(model, modereduce-overhead)torch.compile的reduce-overhead模式会用CUDA Graph减少kernel启动开销对小batch_size场景提升明显。但首次编译需要时间且对动态形状支持有限适合输入形状固定的场景。4.4 常见问题速查表问题现象可能原因排查方法解决方案OOM权重KV Cache激活值超显存torch.cuda.memory_summary()降低精度、减小batch_size、缩短序列长度重复输出贪心搜索陷入循环换Top-p采样测试调整采样策略、加重复惩罚乱码分词器配置错误检查BOS/EOS token重新加载分词器、检查特殊token速度慢KV Cache未生效打印每次输入长度确认use_cacheTrue首token延迟高prompt过长统计prompt长度缩短prompt、用prompt缓存生成提前结束EOS过早采样打印EOS token概率调整EOS logit、设最小长度这张表覆盖了大部分常见问题实际排查时按顺序试一遍基本能定位到原因。如果问题依然存在建议用torch.profiler做一次完整的性能分析看看时间到底花在哪个环节。# 用torch.profiler做性能分析 with torch.profiler.profile( activities[torch.profiler.ProfilerActivity.CUDA], record_shapesTrue ) as prof: engine.generate(测试prompt, max_new_tokens10) print(prof.key_averages().table(sort_bycuda_time_total, row_limit10))这个分析能告诉你哪些kernel占用了最多时间是矩阵乘法、注意力计算还是其他操作。根据结果有针对性地优化比盲目调参有效得多。4.5 实操心得与避坑建议最后分享几个我在实际开发中踩过的坑。第一个坑是显存碎片化。频繁地分配和释放不同大小的张量会导致显存碎片最终即使总空闲显存足够也分配不出连续的大块。解决办法是预分配显存池或者用PagedAttention管理KV Cache。第二个坑是数值精度问题。FP16在累加操作中容易溢出特别是在长序列的注意力计算中。解决办法是在关键位置用FP32累加或者用BF16替代FP16。BF16的动态范围更大不容易溢出但精度略低。第三个坑是批处理中的padding浪费。静态批处理时不同长度的请求要padding到相同长度padding部分的计算是浪费的。解决办法是用连续批处理或者用变长序列的注意力实现。第四个坑是模型加载时的内存峰值。用from_pretrained加载模型时会先加载到CPU内存再转移到GPU峰值内存是模型大小的两倍。解决办法是用device_map直接加载到GPU或者用low_cpu_mem_usageTrue减少CPU内存占用。这些经验在官方文档里通常不会写但实际开发中一定会遇到。希望这个系列能帮你少走一些弯路把推理引擎的核心原理和实操技巧都掌握到手。下一章会开始写代码从最朴素的推理循环开始逐步加入优化。如果你跟着一步步做下来最后会得到一个能实际运行的推理引擎虽然比不上工业级产品但足够让你理解背后的每一个细节。