自己动手写大模型推理引擎:从Token到Transformer再到批处理

发布时间:2026/9/29 11:08:07
自己动手写大模型推理引擎:从Token到Transformer再到批处理 1. 先搞清楚我们到底在写什么推理引擎的双层身份如果你去搜推理引擎这个词会发现一堆看起来很吓人的名词LLM推理引擎、GPU推理引擎、高效推理框架……但实际上当我说自己动手写一个推理引擎的时候指的并不是去造一个全新的、能打败vLLM或者TensorRT-LLM的东西——那是工业级团队干的事。自己动手写的意义在于把大模型从调API的黑盒变成能看透每一层计算的白盒。这个过程中你会被迫搞清楚很多平时根本不会注意到的问题而这些问题的答案恰恰是理解大模型推理的一把钥匙。我给推理引擎一个非常朴素的定义让它输入一段文本经过一系列计算输出另一段文本的程序。看起来简单但这背后其实叠了两层身份。第一层身份是计算执行者。模型本质上是一堆数字参数和一堆计算规则算子输入文本变成数字之后要在这些参数上执行大量的矩阵乘法、归一化、激活函数计算。这一层的核心问题是怎么算得快、算得省。GPU为什么能跑大模型因为它擅长并行做大量矩阵运算。你的推理引擎首先要解决的就是怎么把自己脑子里的模型计算图翻译成GPU上真实跑的算子指令。第二层身份是资源调度者。模型不是凭空在GPU上跑起来的它要占显存、要占计算单元、要跟其他任务抢资源。还记得热词里那句llm request failed: provider rejected the request schema or tool payload吗这种报错很多人第一反应是API接口是不是写错了但实际上在自建推理引擎的场景下很多类似的问题是资源管理和请求调度层面的——引擎没有把并发请求排好队或者上下文长度超了模型的承受范围。所以推理引擎得像个酒店的客房调度系统知道每间房GPU显存块住着谁、什么时候退房、新客人来了怎么安排偶尔还要应付几个客人突然要加床动态申请更大上下文的情况。这两层身份决定了你是从哪个维度去理解和构建它。如果只关注计算执行你会把注意力全部放在算子实现和显存分配上如果只关注资源调度你会更在意批处理、排队、抢占这些逻辑。但一个真正能跑起来的引擎两者缺一不可。我自己动手写的这个推理引擎定位是教育教学版本。它不会是最快的也不追求支持几十种模型架构但它必须做到一点每一行代码都能被理解每一步计算都能被审计。我选择从零实现而不是直接调框架是因为我相信只有把token怎么变成数字、数字怎么变成矩阵、矩阵怎么算出概率、概率怎么挑出下一个词这个完整链路亲手走一遍你对大模型的理解才会真正落地。这里也回答一个常见的疑问我已经会调HuggingFace的transformers库了还需要自己写推理引擎吗我的看法是完全不一样。调transformers库就像你会开车自己写推理引擎就像你打开引擎盖看里面的活塞和曲轴——你可以一辈子不开引擎盖也很爽但一旦你想改装动力、优化油耗、排查异响你至少得知道发动机的基本结构。推理引擎就是大模型的发动机结构。聊完定位第一章接下来要做三件事拆清推理引擎的模块构成、走一遍一次生成的全流程、说清楚写一个能跑的引擎和写一个跑得快的引擎之间的巨大鸿沟。这三件事看完你对整条路线的轮廓就基本有数了。2. 推理引擎的模块地图每一个件都是干什么的先上一张我脑子里的模块地图不是画图是文字描述版的架构清单。后面每一章都会围绕其中一些模块展开所以这个地图值得多看两遍。一个最小的可运行推理引擎至少包含以下模块我按数据从进到出的流动顺序来列模块核心职责类比Tokenizer把文本字符串拆成token序列再映射成数字ID翻译官把人类语言翻译成模型能读的数字语言Embedding层把token ID查表变成稠密向量查字典每个字/词对应一行固定长度的语义向量位置编码给每个token加上位置信息给每个词贴一个第几个位置的标签Transformer层核心多层堆叠每层含注意力前馈网络模型的大脑皮层反复处理信息采样策略从最后一层输出的概率分布中挑一个token决策者根据概率决定下一个生成什么KV Cache缓存历史token的Key/Value矩阵做会议纪要的人避免每次翻旧账显存管理分配/释放张量内存管理缓存仓库管理员批处理调度器把多个请求组合成一批跑得更高效拼车平台顺路的一起走采样参数控制温度、top-p、top-k、惩罚项等超参调随机性旋钮这里面最容易被忽略、但最推理引擎自己的事的是最后三项KV Cache、显存管理、批处理调度器。你会发现模型本身Embedding、Transformer层、采样其实HuggingFace都帮你做好了直接调model.generate()就行。但KV Cache怎么管理、显存怎么复用、请求怎么批次化——这些恰恰是通用框架封装得最深、也最影响性能的东西。自己写引擎主战场就在这三块。也有很多人问写推理引擎需不需要先懂Python我的答案是懂基本语法就够起步真正编程难度不在语法而在数据流向和生命周期管理。你写的是一个状态机数据流水线而不是一首50行的Python小脚本。在动手写之前先搞清楚每个模块之间传递的数据类型是什么Tokenizer输出的是List[int]Embedding输出的是Tensor[batch, seq_len, hidden_dim]注意力输出还是Tensor[batch, seq_len, hidden_dim]最后采样层级输出的是Tensor[batch, seq_len, vocab_size]的概率分布。所有模块之间的接口都是张量形状的对齐形状对不上整个流水线就崩。这就是为什么我在设计引擎时会优先定义一张张量形状对照表——每个模块输入什么形状、输出什么形状提前写死后面出bug的时候查这张表比看报错快得多。接着看模块之间的依赖关系。Tokenizer不依赖GPU纯CPU就能跑Embedding层开始就要访问模型参数得管好内存Transformer层是整个引擎里计算密度最高的部分也是优化的核心对象采样策略虽然计算量很小但对生成质量的影响极大KV Cache则像一个横切模块它跟Transformer层强耦合——因为注意力计算需要读它、更新它。这个模块地图的价值在于后续你写代码的时候脑子里始终有一个数据流图。比如当你想加一个多轮对话功能你只知道加个prompt模板是不够的你得搞清楚历史对话的token序列怎么往前走、KV Cache怎么增量维护当你想做流式输出你得改的是采样层之后的输出逻辑而不是动Transformer层内部的任何东西。模块划清楚了改哪里、不动哪里就一目了然。3. 一次生成背后的全链路从文本到token再到文本聊完模块接下来走一遍完整流程。我用一个具体例子来带着过用户输入今天天气怎么样模型输出今天天气很好适合出门。第一步是Tokenizer。它干的事就是拆词和映射。中文场景下今天天气怎么样这个字符串可能被切成今天、天气、怎么样、具体切法取决于词表。每个token对应一个数字ID假设今天是1024天气是2048怎么样是512是816。这串ID就是模型的输入。这里有一个新手很容易懵的点token不是词它可能是半个词、一个词、甚至一个标点。像llm的token三个点key我是谁、query我在找什么、value我能提供什么这不是我说的是热词里捡来的不过它倒是歪打正着点到了后面要讲的注意力机制——token的形式取决于训练时怎么切的跟你觉得该怎么切没有关系。第二步是Embedding和位置编码。数字ID进入模型后先通过查表变成向量。比如1024查表得到一个768维的向量这个向量试图编码今天这个token的语义。但这里有个问题单独看今天两个字它不知道自己在句子里的位置所以还要加位置编码。位置编码的目的是让模型知道今天在第1个位置、天气在第2个位置这样一来同一个词在不同位置得到不同的输入表示。位置编码有几种实现方式最经典的是用sin/cos函数生成固定位置向量还有现在大模型里更常见的可学习位置编码以及ROPE这种能外推到更长文本的新方法。自己写引擎的时候一开始用最简单的那种就行先让流程跑通回头再升级。第三步是整个引擎的重头戏Transformer层堆叠。我们假设模型有32层每一层都做完全相同的事情先做自注意力计算再做前馈网络计算。自注意力这一步就是上面那个热词歪打正着说的Key、Query、Value——只不过它说的是语义上的理解Key代表我能提供什么、Query代表我在找什么、Value代表我实际给出的信息。具体到计算上每个token生成三个向量Key、Query、Value是通过三个权重矩阵乘出来的然后拿每个token的Query去跟所有token的Key做匹配点积得到注意力分数再归一化成权重最后用权重去加权所有token的Value。这个过程本质上就是让每个token看一遍整句话里的其他token决定自己应该重点关注谁。今天的Query去匹配天气的Key时匹配度高说明今天预测后面出现天气的概率大。来回这么32层每一层都让表示变得更抽象、更语义化。这个步骤的计算量是最大的。假设你的输入长度是100个tokenhidden_size是4096那一层注意力就要做100x100的注意力矩阵计算、再加上四五个大矩阵乘法。32层下来计算量大概在万亿次级别。这就是为什么GPU必不可少——没有GPU并行CPU硬算的话生成一个token可能要几秒钟甚至更久。第四步是采样。32层跑完得到的是最后一个token位置或所有token位置的隐藏表示再经过一个输出投影矩阵把它映射到词表大小比如32000个候选。得到每个候选token的分数再经过softmax变成概率分布。这个时候采样策略开始工作如果选择贪心解码就直接取概率最大的那个token如果选择采样就按概率分布随机抽一个温度高就抽得随机温度低就抽得更确定。这一步直接决定生成内容的风格——同一个问题温度0.1时的回答保守稳定温度0.9时的回答脑洞大开。第五步是把新生成的token拼接进输入序列然后回到第三步继续循环直到生成结束符或者达到最大长度。循环一次生成一个token这就是自回归生成这个名字的来源。整个过程走完今天天气很好适合出门就这么蹦出来了。这版全链路讲完你应该能回答一个问题大模型生成时每一秒都在干什么答案是拉一次Tokenizer、查一次Embedding表、跑32层Transformer、算一次采样概率然后反复循环。听起来不难但每一步的数据量都很大。自己动手写引擎的价值就在于你能亲眼看到这些循环是怎么发生、怎么串联的而不是只在model.generate()这一行代码里看效果。4. 真刀真枪跑起来才发现的事显存、KV Cache与批处理等到我第一版引擎真的能跑起来、能一个token一个token地蹦出句子的时候我以为大功告成了一半。然后我用它的API同时压了3个请求结果直接OOM显存爆了。那一刻我才意识到让模型能推理和让引擎能服务是完全不同的两个World。先说为什么OOM。请求进来的时候我先为每个请求从头算一遍完整流程这就意味着每个请求都独立占一份输入token的张量空间每个请求都要申请完整的显存上下文没有缓存第二个请求进来第一份显存还没释放就得再要一份如果用户的请求长度是2000个token每个token的隐层表示是4096维float162字节光保存所有输入token的中间激活就已经占了2000 × 4096 × 2 16MB这还只是一层的。32层就要512MB。再来3个并发请求一个引擎瞬间吃掉好几个GB显存OOM完全在预期之内。这里面最大的显存黑洞其实是KV Cache。前面提到自注意力计算每个token都需要跟历史token做Key/Value匹配。如果不做任何缓存每次生成新token都得把之前所有token重新算一遍Key和Value——这叫重计算算得慢但省显存。而现代推理引擎普遍选择缓存策略第一次生成时把每一步算出的Key和Value存下来后续每次只算新token的Key和Value再拼接进缓存。这样省掉了重复计算但代价是显存消耗随序列长度线性增长。拿一个7B模型粗略算一笔账32层、40个注意力头、每头维度128、每个token的Key和Value矩阵两层加起来大概是40 × 128 × 2 × 2K和V × 2字节 ≈ 40KB一个token。如果用户的对话长度是4096个token光KV Cache就占4096 × 40KB ≈ 160MB。这还没算模型权重和中间激活。所以一个7B模型在4096上下文下最低显存需求就已经接近10GB模型权重fp16大约14GB——这里用的参数比较粗实际模型大小取决于具体结构我做的是迷你版参数小一些但比例关系是一样的。这也是为什么很多人在自己电脑上跑大模型总卡顿、总报显存不足——不是模型放不下而是KV Cache在背后偷偷吃显存只是你看不见它。KV Cache带来的第二个问题是序列越长计算越慢。因为每生一个新token注意力计算都要跟历史上所有token的K做点积复杂度是O(序列长度)。也就是说让模型写一篇长文越到后面越慢。这跟人类写文章一样写到第2000字上下文已经积累了太多信息。那批处理是怎么回事推理引擎的常见做法是把多个请求拼成一个batch一起过GPU。因为GPU的特点是你让它算一个矩阵乘法跟让它同时算4个独立的矩阵乘法耗时差不多。所以最好的方案是让GPU一直满负荷干活。具体到自回归生成场景多个请求虽然长度不同但可以padding对齐在短序列后面补pad token到同一长度然后拼成一个大batch。这个技巧在学术上叫动态批处理dynamic batching在工业引擎里还有更进阶的连续批处理continuous batching允许一个序列生成完了就立刻把它换出去、让新序列插队进来。我在自己做第一版的时候先从最简单的静态批处理入手——固定batch大小固定序列长度上限做Padding。逻辑简单、容易排查先跑通再优化。这里必须提一个我自己踩过的大坑。批处理性能提升是有上限的不是batch越大越好。当你把4个不同长度的请求拼成一个batch的时候所有序列都被padding到最长那一个的长度。如果其中一个请求特别长比如生成了500个token另外三个只有50个token那三个短序列在每次迭代里都在为并不存在的token做计算GPU的算力就被浪费了。所以在做批处理的时候要设置合理的最大长度或者按长度分桶把长度相近的请求放一批。这个优化技巧看着不起眼但对吞吐量的影响非常大。自己写引擎的过程中我建议你优先跑通以下这条路线单请求、单序列跑通完整生成流程这就是完成了能推理单请求、加入KV Cache验证长文本生成是否不再爆炸式重计算性能第一台阶多请求、静态批处理验证显存占用和吞吐量变化性能第二台阶多请求、动态长度管理避免padding浪费性能第三台阶走到第4步你已经比大量只会调接口的人更了解大模型推理了。再往后才是量化、投机采样、paged attention这些高阶玩法。5. 从第一个朴素版本到现代引擎推理优化地图写到这里你可能会想我的引擎大概能跑、能批处理了然后呢离真正现代化还有多远坦白说还挺远的。但如果把现代推理引擎看成一张优化地图你目前在哪个位置、下一站去哪儿这张地图还是值得提前看一眼因为它决定了你后面往哪个方向使劲。地图上有几条主要的优化路线第一条模型量化。这条路线解决的是显存瓶颈。模型权重默认是fp1616位浮点数2字节一个7B模型光权重就要14GB。如果量化成int81字节就变成7GBint4更进一步压到3.5GB。你的引擎对显存的需求直接减半甚至减到四分之一意味着很多原本跑不动的场景比如个人电脑就能跑起来了。但量化不是白给的——精度会损失。尤其int4模型生成的语感会变差容易出现逻辑不连贯。一般建议int8精度损失很小int4要看任务。自己写引擎时可以先做最简单的那种权重只量化到int8版本验证显存-精度-速度三者之间的关系不用一上来就追求4bit极致压缩。第二条投机采样Speculative Decoding。这条路线解决的是生成速度瓶颈。大模型生成慢是因为一个token一个token地想——每一步都要跑一遍整个模型。投机采样的思路很反直觉用一个小模型draft model先快速猜出接下来要生成的几个token然后大模型target model并行核验。核验通过就一次性接受好几个token失败就校正。它的前提是小模型猜得对的地方多大模型核验比生成便宜。很多现代推理引擎都内置了这套机制但我建议你至少在朴素引擎跑通之后再考虑它。因为投机采样的debug难度较高它引入了两个模型协同的复杂度新手容易在batch对齐上翻车。第三条并行与分布。模型太大一块GPU放不下怎么办答案是模型切分。水平切分张量并行把同一个矩阵拆成几块每块放一块GPU算完再拼起来垂直切分流水线并行把Transformer层分几段第一块GPU跑1-11层第二块跑12-22层第三块跑23-32层数据一层层往下传。这两条路线解决的是单卡放不下大模型的问题。但你自己写引擎的话我建议先别碰——它在单卡上根本不会暴露出问题一上多卡通信瓶颈、负载不均、死锁、梯度同步各种诡异问题都来了。在没有绝对必要比如你想在本地跑70B模型的情况下先享受单卡带来的简单幸福。第四条注意力机制的变体paged attention等。这个解决的是KV Cache的显存碎片化问题。现代引擎比如vLLM的核心卖点之一就是paged attention——把连续的KV Cache物理空间打散成页式存储按需分配、按需释放避免一个长序列把显存占了一大块但实际只用了30%的浪费。这个思路其实是从操作系统内存的分页管理借鉴过来的它在显存利用率上的提升非常显著。但它也对应了相当复杂的内存映射逻辑实现起来极其考验心智。我打算把它放在第二期去搞第一期先跑通朴素版本。这里我想给一张优化路线的性价比表方便你规划优化项提升效果实现难度适合时机KV Cache长文本生成提速几倍低第一版引擎就该做静态批处理吞吐量提升数倍中跑通基础流程之后动态批处理吞吐量、显存利用率再提中高有服务化需求时int8量化显存减半推理稳中有部署需求时int4量化显存减至1/4高追求极致部署时投机采样生成速度提升1.5-3倍高基础吞吐优化完Paged Attention显存碎片化大降很高大并发大上下文时多卡并行单卡限制解除很高需要跑超大模型时我自己目前的规划是第一篇先讲清楚引擎是什么就是你正在看的这篇概览第二篇讲Tokenizer和词表构建第三篇讲Embedding和位置编码第四篇讲Transformer层和自注意力这是最硬核的一篇第五篇讲KV Cache和显存管理第六篇讲批处理和调度第七篇讲采样策略和生成质量第八篇讲量化和部署。这是一条从能算到算得快再到部署得好的完整路线每一篇我都希望能用最朴素的话把原理和代码串起来。最后分享一个我在动手写作过程中最深的一个感受大模型的智能感其实是从一层层普通计算中涌现出来的而不是藏在某个神奇的公式里。当你亲手把那些矩阵乘法写出来、看着一个token一个token地生成一个通顺的句子你会觉得哇不过是一堆浮点数在打架嘛。但正是这个不过如此的过程让你对这堆浮点数背后的原理有了真正的信任和掌控。下一章我们就从Tokenizer开始把人话变成数字。