1. 为什么要自己动手写一个LLM推理引擎市面上已经有这么多推理框架了为什么还要自己写一个这个问题我被问过不下十次。每次我的回答都一样因为你不自己写一遍你永远不知道那些框架里到底在干什么。我用过不少开源推理方案也读过一些框架的源码。说实话大部分时候你调用一个接口就能跑起来但一旦遇到性能瓶颈、显存不够、输出结果异常这些问题如果你不理解底层的推理流程排查起来就跟盲人摸象一样。你不知道KV Cache是怎么管理的不知道采样策略在哪里生效不知道batch是怎么拼的那你调参就纯粹是在碰运气。自己动手写一个推理引擎哪怕只是一个最小可用的版本能让你把整个链路彻底打通。从模型权重加载、tokenizer处理、前向计算、KV Cache管理、采样策略、到最终输出解码每一个环节你都要亲手实现一遍。这个过程里踩过的坑比你看十篇论文都管用。这篇文章是系列的第一章主要做整体概览。我会把推理引擎的核心模块拆开讲清楚每个模块解决什么问题、模块之间怎么协作、有哪些常见的实现方案和取舍。后续章节会逐个模块深入带代码实现。适合有一定Python基础、了解Transformer基本结构、想深入理解LLM推理过程的开发者。如果你只是想让模型跑起来那直接用现成框架就行但如果你想搞清楚“为什么这么慢”“为什么显存炸了”“为什么输出不稳定”那这个系列就是写给你的。2. 推理引擎的整体架构拆解2.1 从一次完整的推理请求说起先看一个最朴素的场景用户输入一段prompt模型生成一段回复。这中间到底发生了什么整个过程可以拆成几个阶段。第一阶段是预处理包括tokenizer把文本转成token id序列然后组装成模型需要的输入格式。第二阶段是前向计算模型根据输入token逐个或批量地计算出下一个token的概率分布。第三阶段是采样根据概率分布选出下一个token。第四阶段是后处理把生成的token id序列解码回文本。但这里有个关键点LLM的生成是自回归的也就是说每生成一个token都要把它拼接到输入序列后面再跑一次前向计算。如果每次都把整个序列重新算一遍那计算量会随着序列长度平方级增长。这就是为什么需要KV Cache——把之前算过的Key和Value缓存下来每次只算新token的那一部分。所以一个推理引擎的核心工作其实就是围绕“如何高效地管理KV Cache”和“如何高效地调度多个请求”来展开的。前者决定了单次推理的速度后者决定了整体吞吐量。2.2 核心模块划分与职责一个完整的推理引擎我把它拆成六个核心模块模型加载器负责从磁盘加载权重文件初始化模型结构把参数放到正确的设备上。这里涉及权重格式解析、数据类型转换、设备分配等细节。Tokenizer文本和token id之间的双向转换。看起来简单但不同模型的tokenizer行为差异很大特殊token的处理、padding策略、截断逻辑都需要仔细处理。推理调度器管理多个请求的排队、批处理、优先级调度。这是区分玩具实现和生产级引擎的关键模块。KV Cache管理器负责缓存空间的分配、回收、复用。显存够不够用很大程度上取决于这个模块的设计。采样器根据logits选出下一个token。贪心、beam search、top-k、top-p、温度调节等策略都在这里实现。输出解码器把生成的token序列转回文本处理流式输出、停止条件等。这六个模块里调度器和KV Cache管理器是最复杂的也是最能体现一个推理引擎设计水平的。模型加载和tokenizer相对标准化采样器逻辑清晰但参数调优需要经验输出解码器主要是工程细节。2.3 为什么选择从零实现而不是改现成框架有人可能会问我直接去改开源框架的源码不行吗当然行但有几个问题。第一开源框架为了通用性往往做了大量抽象代码路径很长。你改一个地方可能影响到你完全没想到的模块。第二框架的文档通常不会告诉你“为什么这么设计”你只能靠读代码去猜。第三也是最关键的改框架你永远是在别人的思路上打补丁很难形成自己的整体认知。从零实现的好处是你可以完全控制复杂度。第一版不需要支持所有特性先把最核心的链路跑通单请求、单批次、贪心采样。跑通之后再逐步加功能批处理、KV Cache优化、多种采样策略、流式输出。每一步你都知道自己在加什么、为什么加。我自己的经验是第一版代码量控制在500行以内比较合适。太少了跑不起来太多了容易迷失在细节里。后续每加一个功能代码量增长控制在100行左右这样你能清楚地感知到每个功能的成本。3. 核心模块的详细设计与实现要点3.1 模型加载权重解析与设备分配模型加载看起来是最没技术含量的部分但实际上坑不少。首先是权重格式常见的有safetensors、PyTorch的bin格式、GGUF等。不同格式的读取方式不同safetensors可以直接内存映射加载速度快且安全bin格式需要反序列化有安全风险且速度慢一些。加载的时候要注意数据类型。训练时通常用fp32或bf16推理时可以根据硬件选择fp16、bf16、int8甚至int4。数据类型直接影响显存占用和计算速度。以7B模型为例fp16需要约14GB显存int8需要约7GBint4只需要约3.5GB。但量化会带来精度损失需要根据任务类型权衡。设备分配也是个问题。如果只有一张卡直接全部放上去就行。但如果是多卡就要考虑模型并行策略。最简单的做法是按层切分每张卡放一部分层。但这样通信开销会比较大因为每层计算完都要把结果传给下一张卡。另一种做法是张量并行把每个层的矩阵运算切分到多张卡上通信更频繁但计算更均衡。注意加载权重时一定要检查参数名称是否和模型结构完全匹配。我遇到过好几次因为key名称差一个前缀导致加载失败的情况排查起来很费时间。建议加载后打印一下模型参数的名称列表和权重文件的key做一次比对。3.2 Tokenizer不只是文本转IDTokenizer的工作看似简单但细节很多。首先是特殊token的处理。比如BOS token、EOS token、PAD token不同模型的定义不同有的模型甚至没有BOS token。生成的时候EOS token是停止条件如果处理不当会导致生成无法停止。其次是padding策略。批处理的时候不同请求的输入长度不同需要padding到相同长度。但padding token不应该参与注意力计算所以需要attention mask。如果mask写错了模型会关注到padding部分输出质量会明显下降。还有一个容易被忽略的点是截断。当输入超过模型最大长度时需要截断。但截断哪里截断开头还是结尾对于对话模型通常保留最近的对话历史截断较早的内容。但有些任务需要保留开头的信息比如系统提示词。这个策略需要根据具体场景来定。我自己的做法是在tokenizer模块里维护一个配置对象把所有和tokenizer相关的参数都放在里面最大长度、截断策略、padding方向、特殊token id等。这样后续换模型的时候只需要改配置不需要改代码。3.3 推理调度器从单请求到批处理单请求的推理流程很直接输入进来跑模型输出结果。但实际服务中请求是并发来的。如果每个请求都单独跑一次模型GPU利用率会非常低因为大部分时间都在等内存搬运。批处理是提升吞吐量的关键。把多个请求拼成一个batch一次前向计算就能处理多个请求。但这里有个问题不同请求的生成长度不同有的可能生成10个token就结束了有的要生成500个。如果等所有请求都结束再返回那短请求的延迟会很高。解决方案是连续批处理continuous batching。当一个请求生成结束时立刻把新的请求加进来不用等整个batch结束。这样GPU始终处于忙碌状态吞吐量和延迟都能兼顾。但实现起来复杂度高很多需要动态管理batch内的序列、动态调整attention mask、动态分配KV Cache空间。另一个问题是优先级。有些请求需要低延迟有些可以容忍高延迟。调度器需要支持优先级队列高优先级的请求优先处理。这个在实现上可以用多个队列每个队列对应一个优先级调度时从高优先级队列开始取。实操心得第一版实现建议先做静态批处理也就是攒一批请求一起跑跑完再跑下一批。虽然延迟高一些但实现简单容易调试。等静态批处理跑通了再改成连续批处理。我当初就是直接上连续批处理结果调试了两周才跑通中间一度想放弃。3.4 KV Cache管理显存优化的核心战场KV Cache是推理引擎里最核心的数据结构。每生成一个token都要把当前层的Key和Value缓存下来供后续token计算注意力时使用。缓存的大小是层数 × 2 × 注意力头数 × 头维度 × 序列长度 × 数据类型大小。以7B模型为例假设32层、32个注意力头、头维度128、fp16数据类型序列长度2048那么KV Cache大小约为32 × 2 × 32 × 128 × 2048 × 2字节 ≈ 1GB。这还只是一个请求的缓存。如果同时处理10个请求就是10GB。加上模型本身的14GB显存很快就满了。所以KV Cache管理的关键是复用和压缩。复用是指不同请求之间共享缓存空间比如用分页管理的方式把缓存切成固定大小的块按需分配。压缩是指用更小的数据类型存储缓存比如int8量化或者用一些稀疏化技术丢弃不重要的缓存。分页管理是目前主流方案。把KV Cache切成固定大小的page每个page存储一定数量token的缓存。请求需要缓存时从空闲page列表中分配请求结束时回收page。这样不同请求的缓存可以交错存放显存利用率更高。注意分页管理需要考虑page大小和序列长度的关系。page太小管理开销大page太大内部碎片多。经验值是page大小设为16或32个token比较合适。另外注意力计算时需要根据page索引找到对应的缓存块这个索引结构的设计直接影响计算效率。3.5 采样策略从贪心到核采样采样器决定了生成的多样性。最简单的贪心采样就是每次选概率最大的token结果确定但缺乏多样性。温度采样是在softmax之前把logits除以温度值温度大于1分布更平坦小于1分布更尖锐。Top-k采样是只保留概率最高的k个token其余置零后重新归一化。Top-p采样也叫核采样是保留累积概率达到p的最小token集合。这两种方法都能避免采样到低概率的奇怪token。实际使用中通常是组合使用先按温度调整分布再按top-k或top-p截断最后从剩余token中采样。参数的选择需要根据任务来定。代码生成通常用低温加top-p创意写作可以用高温加top-k。还有一个细节是重复惩罚。模型有时候会陷入循环反复生成相同的短语。重复惩罚就是在计算logits时对已经出现过的token降低其概率。惩罚力度需要调太强会导致生成不流畅太弱又起不到作用。3.6 输出解码流式与非流式输出解码就是把token id序列转回文本。非流式就是等所有token生成完一次性解码返回。流式是每生成一个token就解码并返回用户可以看到逐字输出的效果。流式输出的难点在于多字节字符的处理。一个中文字符可能对应多个token如果逐个token解码可能会出现乱码。解决方案是维护一个缓冲区当解码结果包含不完整字符时先缓存起来等后续token补齐后再输出。停止条件也需要仔细处理。除了EOS token还需要支持自定义停止词。比如用户指定“\n\n”作为停止词那生成过程中要检测是否出现了这个序列。检测逻辑要考虑跨token的情况因为停止词可能被切分到多个token里。4. 实操过程从零搭建一个最小可用引擎4.1 环境准备与依赖选择先明确目标实现一个支持单请求、贪心采样、fp16推理的最小引擎。不需要批处理不需要量化不需要流式输出。先把链路跑通。依赖方面PyTorch是必须的用来做张量计算和模型加载。Transformers库可以用来加载tokenizer和模型权重但前向计算我们自己写不用它的generate方法。safetensors用来读取权重文件。numpy用来做数值计算辅助。环境配置上CUDA版本要和PyTorch版本匹配。我遇到过好几次因为CUDA版本不匹配导致的各种奇怪报错排查起来很痛苦。建议先用一个简单的张量运算测试一下GPU是否正常工作。import torch # 测试GPU是否可用 print(torch.cuda.is_available()) print(torch.cuda.get_device_name(0)) # 测试基本运算 x torch.randn(1000, 1000).cuda() y torch.randn(1000, 1000).cuda() z x y print(z.shape)如果这段代码能跑通环境基本没问题。4.2 模型结构的最小实现我们以LLaMA结构为例实现一个最小版本。核心组件包括RMSNorm比LayerNorm更简单没有均值中心化只有缩放。旋转位置编码把位置信息编码到Query和Key中通过旋转矩阵实现。多头注意力标准的QKV计算加上因果mask。前馈网络SwiGLU激活函数三个线性层。代码结构上我建议把每个组件写成独立的类然后组合成Transformer层再堆叠成完整模型。这样调试的时候可以逐层验证。class RMSNorm(nn.Module): def __init__(self, dim, eps1e-6): super().__init__() self.eps eps self.weight nn.Parameter(torch.ones(dim)) def forward(self, x): norm x * torch.rsqrt(x.pow(2).mean(-1, keepdimTrue) self.eps) return norm * self.weight这是RMSNorm的实现非常简洁。注意eps的位置是在均方值上加eps再开根号不是先开根号再加eps。这个细节会影响数值稳定性。4.3 前向计算的完整流程前向计算的核心是注意力机制。给定输入hidden states先计算Q、K、V然后做注意力计算最后过输出投影。def forward(self, x, kv_cacheNone, start_pos0): bsz, seq_len, _ x.shape # 计算QKV q self.q_proj(x) k self.k_proj(x) v self.v_proj(x) # 重塑为多头格式 q q.view(bsz, seq_len, self.n_heads, self.head_dim).transpose(1, 2) k k.view(bsz, seq_len, self.n_kv_heads, self.head_dim).transpose(1, 2) v v.view(bsz, seq_len, self.n_kv_heads, self.head_dim).transpose(1, 2) # 应用旋转位置编码 q, k apply_rotary_emb(q, k, start_pos) # 更新KV Cache if kv_cache is not None: k torch.cat([kv_cache[0], k], dim2) v torch.cat([kv_cache[1], v], dim2) kv_cache (k, v) # 注意力计算 attn (q k.transpose(-2, -1)) / math.sqrt(self.head_dim) attn attn self.causal_mask[:, :, start_pos:start_posseq_len, :start_posseq_len] attn F.softmax(attn, dim-1) out attn v # 输出投影 out out.transpose(1, 2).reshape(bsz, seq_len, -1) return self.o_proj(out), kv_cache这段代码里KV Cache的拼接是关键。每次前向计算只算新token的K和V然后拼接到缓存后面。注意力计算时Query只有新token的Key和Value是完整的序列。实操心得第一次实现的时候我忘了在注意力计算时对Key做转置结果形状对不上报错信息又很模糊排查了很久。建议每步计算后都打印一下张量形状确认符合预期。4.4 生成循环与停止条件生成循环就是不断调用前向计算每次取最后一个位置的logits采样出下一个token拼接到输入序列后面继续下一轮。def generate(model, input_ids, max_new_tokens100, eos_token_idNone): generated input_ids kv_cache None for _ in range(max_new_tokens): with torch.no_grad(): logits, kv_cache model(generated[:, -1:], kv_cache) next_token_logits logits[:, -1, :] next_token torch.argmax(next_token_logits, dim-1, keepdimTrue) generated torch.cat([generated, next_token], dim-1) if eos_token_id is not None and next_token.item() eos_token_id: break return generated注意这里每次只传入最后一个token因为前面的token已经通过KV Cache缓存了。如果传入完整序列KV Cache会重复计算结果会出错。停止条件除了EOS token还应该支持最大长度限制。另外如果生成了PAD token或者其他特殊token也应该考虑是否停止。5. 常见问题与排查技巧实录5.1 显存溢出从现象到根因显存溢出是最常见的问题。现象是程序报CUDA out of memory但原因可能有很多种。第一种是模型本身太大加载就爆了。这时候需要检查模型参数量和显存容量是否匹配。7B模型fp16需要约14GB加上KV Cache和中间激活值实际需要16GB以上。如果显存不够可以考虑量化或者模型并行。第二种是KV Cache增长过快。如果生成长度很长KV Cache会线性增长。这时候需要限制最大生成长度或者实现缓存淘汰策略。第三种是中间激活值占用过多。前向计算过程中会产生大量中间张量如果batch size太大或者序列太长这些张量会占用大量显存。解决方案是减小batch size或者用梯度检查点技术推理时不需要梯度但可以用类似思路减少中间张量。排查的时候可以用torch.cuda.memory_summary()查看显存分配情况定位是哪个部分占用了最多显存。5.2 输出乱码或重复采样与解码的坑输出乱码通常和tokenizer有关。如果tokenizer的编码解码不对称或者特殊token处理不当就会出现乱码。排查方法是单独测试tokenizer编码一段文本再解码看是否和原文一致。输出重复通常和采样策略有关。如果温度太低模型会倾向于选概率最高的token容易陷入循环。解决方案是提高温度或者加重复惩罚。重复惩罚的实现是在logits上减去一个和出现次数相关的值。还有一种情况是KV Cache更新错误导致模型看到了错误的上下文。这种问题比较隐蔽现象是输出逻辑混乱但又不是完全乱码。排查方法是关掉KV Cache每次重新计算完整序列看输出是否正常。如果正常说明KV Cache实现有问题。5.3 生成速度慢性能瓶颈定位生成速度慢的原因可能有很多。首先用profiler看一下时间花在哪里。PyTorch自带的profiler可以给出每个操作的耗时。常见瓶颈包括注意力计算没有用Flash Attention导致显存访问频繁KV Cache拼接用了torch.cat每次都会重新分配内存采样时用了CPU同步操作导致GPU等待。优化方向用Flash Attention替换标准注意力用预分配的缓存池代替动态拼接把采样逻辑放到GPU上避免CPU-GPU同步。注意优化之前一定要先测量不要凭感觉猜。我见过有人花了一周优化注意力计算结果发现瓶颈其实在tokenizer上。用数据说话不要用直觉。5.4 常见问题速查表问题现象可能原因排查方法解决方案显存溢出模型太大/KV Cache过长/中间激活过多memory_summary查看分配量化/限制长度/减小batch输出乱码tokenizer编解码不对称单独测试tokenizer检查特殊token处理输出重复温度过低/缺少重复惩罚调整采样参数提高温度/加重复惩罚生成速度慢注意力未优化/KV Cache拼接低效profiler定位瓶颈Flash Attention/预分配缓存生成不停止EOS token未正确处理检查EOS token id正确设置停止条件结果不一致KV Cache更新错误关闭KV Cache对比检查缓存拼接逻辑6. 后续章节规划与扩展方向第一章主要是搭框架、理思路。后续章节我会逐个模块深入带完整的代码实现。第二章讲模型加载和权重解析包括不同格式的读取、数据类型转换、设备分配策略。第三章讲tokenizer的细节特别是特殊token和padding的处理。第四章讲KV Cache的实现从最简单的拼接版本到分页管理。第五章讲调度器从静态批处理到连续批处理。第六章讲采样策略的数学原理和代码实现。第七章讲性能优化包括Flash Attention、量化、算子融合等。每个章节都会遵循同样的结构先讲清楚问题和需求再讲设计方案和取舍然后给代码实现最后分享实操中遇到的坑和解决技巧。这个系列不会涉及太前沿的研究内容主要是工程实现层面的经验分享。目标是让你看完之后能自己动手写出一个可用的推理引擎并且理解每个设计决策背后的原因。我自己在实现过程中最大的体会是不要一开始就追求完美。第一版能跑通就行哪怕慢一点、占显存多一点。跑通之后你才有基础去优化才知道瓶颈在哪里。如果一开始就想做一个生产级的引擎很容易陷入细节出不来最后什么都没做成。先跑通再优化这个顺序不能反。