从零搭建AI工程能力:手写KV Cache与动态批处理实战
发布时间:2026/10/2 16:26:51来源:尧图网络
1. 从零搭建AI工程能力为什么我劝你别一上来就调包这两年AI应用开发的门槛肉眼可见地降低了随便拉个框架、调个API就能跑出一个能对话的Demo。但我见过太多团队Demo阶段惊艳四座一上生产就原形毕露推理延迟飙到几秒、显存说爆就爆、换个模型整个pipeline推倒重来、线上效果和离线评测对不上。这些问题的根子几乎都指向同一件事——工程能力没有跟上模型能力。ai-engineering-from-scratch这个标题说的就是这件事不依赖现成的高级封装从最底层的张量运算、数据管道、推理调度、服务部署一步步搭起来把AI系统里每一层的“为什么”搞清楚。它不是一个具体的开源项目名而是一类学习路径和实践方法的统称。适合谁看适合那些已经会用框架跑模型、但一遇到性能瓶颈和线上故障就抓瞎的工程师也适合想从算法岗转向AI工程岗、需要补齐系统能力的人。我自己走过这条路也带过几个从零起步的团队。最大的体会是你调包时省下的每一分钟都会在生产环境里以十倍的调试时间还回来。这篇文章就把我从零搭建AI工程能力的完整思路拆开讲包括整体设计、核心细节、实操过程和踩坑记录尽量做到你照着就能复现。2. 整体设计与思路拆解先画地图再走路2.1 为什么选择“自底向上”而不是“自顶向下”大多数人学AI工程是从上往下走的先学LangChain这类编排框架再学推理引擎最后才碰底层。这条路入门快但有个致命问题——你永远不知道框架帮你做了什么也就永远不知道它什么时候会坑你。比如一个简单的流式输出框架帮你处理了token拼接、缓冲区管理、连接保活看起来很美可一旦出现输出截断或者乱序你连从哪查起都不知道。自底向上的路径反过来先搞明白一个矩阵乘法在CPU和GPU上分别怎么执行、内存怎么分配再理解KV Cache为什么能加速自回归生成然后才去看推理引擎怎么调度请求。这样搭起来的认知是可迁移的——换个框架、换个模型、换个硬件底层原理不变你只需要重新映射一遍就行。我一般把整个学习路径分成四层从下到上依次是层级核心内容关键产出计算层张量运算、自动微分、内存管理手写一个能跑通反向传播的迷你框架模型层Transformer结构、注意力机制、KV Cache从零实现一个可生成文本的小型语言模型推理层批处理、量化、算子融合、调度策略一个支持动态批处理的推理服务服务层API设计、并发控制、监控告警、灰度发布一个能扛住压测的线上服务这四层不是严格串行的但每一层都得有能跑通的代码不能只停留在看论文。我的经验是看十篇讲解KV Cache的文章不如自己实现一遍带Cache的注意力计算哪怕性能很差那个“为什么能省计算量”的直觉是看文章永远给不了的。2.2 技术选型的几个关键取舍从零搭建不等于什么都自己写。有些轮子该用还得用关键是知道边界在哪。我的选型原则是计算密集的部分用成熟库调度和业务逻辑自己写。具体来说张量运算和基础算子直接用PyTorch或者NumPy没必要自己写CUDA核函数除非你的目标就是学CUDA。但推理调度、请求队列、批处理策略这些一定要自己实现一遍。因为这部分是AI工程和普通后端工程最大的差异点也是最容易出问题的地方。举个例子动态批处理continuous batching是现在推理服务的标配但它的实现逻辑其实不复杂——维护一个请求队列每次有新请求进来或者有请求完成时重新组一个batch送进模型。自己写一遍你就能理解为什么batch size不是越大越好、为什么prefill和decode要分开调度、为什么显存碎片会导致OOM。这些认知直接决定了你线上调参的方向。另一个取舍是评测体系。很多人从零搭AI系统时只顾着跑通忽略了评测。我的建议是在写第一行推理代码之前先把评测集和指标定下来。语言模型至少要有困惑度perplexity和任务准确率两个维度生成任务还要加上人工评估或者BLEU/ROUGE这类自动指标。没有评测你后面所有的优化都是盲人摸象。3. 核心细节解析与实操要点每一层的关键决策3.1 计算层手写自动微分到底值不值计算层最核心的是自动微分。PyTorch的autograd用起来太方便了方便到很多人从来没想过它是怎么实现的。我强烈建议你手写一个迷你版的自动微分引擎不用支持太多算子能跑通一个两层全连接网络就行。实现思路很简单每个张量除了存数据再存一个grad字段和一个_backward函数。前向计算时构建计算图反向传播时按拓扑逆序调用每个节点的_backward。核心代码大概长这样class Tensor: def __init__(self, data, requires_gradFalse): self.data data self.grad None self._backward lambda: None self._prev set() self.requires_grad requires_grad def __add__(self, other): out Tensor(self.data other.data, self.requires_grad or other.requires_grad) out._prev {self, other} def _backward(): if self.requires_grad: self.grad out.grad if self.grad is None else self.grad out.grad if other.requires_grad: other.grad out.grad if other.grad is None else other.grad out.grad out._backward _backward return out这段代码看起来简单但写一遍你就能理解几个关键问题为什么梯度要累加而不是覆盖、为什么需要拓扑排序、为什么原地操作会破坏计算图。这些在调包时都是隐式处理的但线上出问题时往往就是这些地方在作祟。注意手写自动微分只适合学习生产环境绝对不要自己实现。但理解它的原理能让你在看PyTorch报错时快速定位是计算图的问题还是数据的问题。3.2 模型层KV Cache是第一个性能分水岭模型层最值得深挖的是Transformer的推理优化而KV Cache是第一个必须拿下的山头。自回归生成时每生成一个token都要重新计算所有历史token的Key和Value这是巨大的浪费。KV Cache的思路是把历史token的K和V缓存下来每步只计算新token的K和V然后拼接。实现上有个细节很容易踩坑Cache的存储布局。如果你用[batch, seq_len, num_heads, head_dim]的布局每次拼接都要重新分配内存效率很低。更好的做法是预分配一个最大长度的缓冲区用索引来标记当前写到哪了。这个思路和后面要讲的PagedAttention是一脉相承的。另一个细节是Cache的精度。很多人在做量化时会忽略KV Cache也可以量化。把KV Cache从FP16降到INT8显存占用直接减半对长序列场景效果非常明显。但要注意KV Cache的量化误差会累积因为每一步的输出都依赖上一步的Cache。我的经验是KV Cache用INT8权重量化用INT4这个组合在大多数场景下精度损失可控。3.3 推理层动态批处理的实现要点推理层的核心是调度。静态批处理要求所有请求长度一致实际场景根本不可行。动态批处理允许不同请求在不同时间加入和退出实现上需要维护一个请求状态机。每个请求有这几个状态等待中、prefill中、decode中、已完成。调度器每次循环做三件事把新请求加入等待队列、把等待队列里能塞进当前batch的请求做prefill、把decode中的请求继续生成。关键参数是最大batch size和最大等待时间——batch size越大吞吐越高但延迟越大等待时间越长batch越容易填满但首token延迟越高。我实测下来的一组经验值在线对话场景最大batch size设32到64最大等待时间设10到20毫秒离线批量推理场景batch size可以拉到128以上等待时间设100毫秒以上。这些数字不是绝对的但可以作为你调参的起点。还有一个容易忽略的点prefill和decode要分开调度。prefill是计算密集型的decode是内存带宽密集型的混在一起会导致GPU利用率上不去。现在主流的做法是chunked prefill把长prompt切成小块和decode请求混批这样既能填满计算单元又不会让decode请求等太久。3.4 服务层并发控制和优雅降级服务层最核心的是并发控制。AI推理服务和普通Web服务最大的区别是每个请求消耗的资源差异巨大。一个短prompt可能只占几百MB显存一个长prompt可能直接把显存打满。所以不能用简单的请求数来做限流要基于token数或者预估显存来做。我的做法是在入口处加一个准入控制器根据当前显存水位和请求的预估token数来决定是否接收。如果显存水位超过80%新请求要么排队要么直接返回降级结果。降级策略可以是用更小的模型、更短的输出长度、或者直接返回缓存结果。监控指标也要专门设计。除了常规的QPS、延迟、错误率AI服务必须监控这几个首token延迟TTFT、每token延迟TPOT、显存利用率、batch填充率、Cache命中率。TTFT影响用户体验TPOT影响生成速度batch填充率反映调度效率Cache命中率反映前缀复用情况。这些指标在普通Web监控里是没有的需要自己埋点。4. 实操过程与核心环节实现从零到一跑通全流程4.1 环境准备与依赖选择环境准备这块我踩过不少坑说几个关键决策。Python版本建议3.10或3.113.12有些库的wheel还没跟上。PyTorch版本要和CUDA版本匹配这个去官网查兼容性表就行别自己猜。显存监控用pynvml比nvidia-smi解析文本靠谱得多。依赖管理我强烈建议用uv或者conda别用裸pip。AI项目的依赖冲突太常见了特别是transformers、accelerate、bitsandbytes这几个库版本不匹配直接报错。我的做法是先用uv pip compile生成锁定文件确保每次环境重建都是一致的。# 用uv创建环境并锁定依赖 uv venv ai-eng --python 3.11 source ai-eng/bin/activate uv pip install torch2.3.0 transformers4.42.0 accelerate0.31.0 uv pip freeze requirements.lock提示如果你的机器有多张卡记得设置CUDA_VISIBLE_DEVICES来指定用哪张不然PyTorch默认会用第一张很容易和别人冲突。4.2 从零实现一个带KV Cache的生成循环这是整个实操里最核心的一段。我把它拆成四步模型加载、Cache初始化、prefill阶段、decode循环。模型加载没什么好说的用transformers的AutoModelForCausalLM就行。关键是Cache初始化我不用transformers自带的Cache类而是自己用张量管理这样能更清楚地控制内存布局。import torch class SimpleKVCache: def __init__(self, batch_size, max_seq_len, num_layers, num_heads, head_dim, dtype, device): self.max_seq_len max_seq_len self.num_layers num_layers # 预分配缓冲区避免动态拼接 self.k_cache torch.zeros( num_layers, batch_size, num_heads, max_seq_len, head_dim, dtypedtype, devicedevice ) self.v_cache torch.zeros_like(self.k_cache) self.seq_len 0 def append(self, layer_idx, k_new, v_new): # k_new: [batch, num_heads, new_len, head_dim] new_len k_new.shape[2] end self.seq_len new_len self.k_cache[layer_idx, :, :, self.seq_len:end, :] k_new self.v_cache[layer_idx, :, :, self.seq_len:end, :] v_new if layer_idx self.num_layers - 1: self.seq_len end def get(self, layer_idx): return ( self.k_cache[layer_idx, :, :, :self.seq_len, :], self.v_cache[layer_idx, :, :, :self.seq_len, :] )prefill阶段把整个prompt送进去一次性计算所有位置的K和V并写入Cache。decode阶段每次只送一个token从Cache里取出历史的K和V和新计算的拼接后做注意力。这里有个性能陷阱每次decode都调用get返回切片会创建新的张量视图虽然不复制数据但会增加Python层的开销。优化方法是在注意力计算里直接传Cache对象和索引避免频繁创建视图。这个优化在长序列场景下能带来10%到20%的吞吐提升。4.3 动态批处理调度器的实现调度器我用一个简单的循环来实现核心是维护三个队列等待队列、运行队列、完成队列。class Scheduler: def __init__(self, max_batch_size, max_wait_ms): self.waiting [] self.running [] self.max_batch_size max_batch_size self.max_wait_ms max_wait_ms def step(self): # 1. 把等待超时的请求移入运行队列 now time.time() * 1000 ready [r for r in self.waiting if now - r.arrive_time self.max_wait_ms] # 2. 如果运行队列没满从等待队列补充 while len(self.running) self.max_batch_size and self.waiting: self.running.append(self.waiting.pop(0)) # 3. 对运行队列做一次前向 if self.running: self._forward_batch() # 4. 把完成的请求移出 self.running [r for r in self.running if not r.finished]这个调度器很粗糙但能跑通。实际生产里还要考虑prefill和decode分开调度、显存不足时的抢占、请求优先级、超时取消等。但先把最简版本跑通再逐步加功能比一上来就设计一个完美架构要靠谱得多。我实测下来这个简单调度器在batch size为32时吞吐比静态批处理高3到5倍首token延迟增加不到20毫秒。对于大多数中小规模场景这个性价比已经足够了。4.4 服务化与压测服务化我用FastAPI因为它的异步支持好适合IO密集型的推理服务。关键是要把推理放在单独的线程或者进程里避免阻塞事件循环。from fastapi import FastAPI from concurrent.futures import ThreadPoolExecutor app FastAPI() executor ThreadPoolExecutor(max_workers1) app.post(/generate) async def generate(request: GenerateRequest): loop asyncio.get_event_loop() result await loop.run_in_executor(executor, inference_engine.generate, request) return {text: result}压测用locust或者wrk都行但要注意压测的请求分布要接近真实场景。我一般会准备三组请求短prompt短输出、短prompt长输出、长prompt短输出按6:3:1的比例混合。这样压出来的数据才有参考价值。压测时重点看三个指标P99首token延迟、P99每token延迟、显存峰值。如果P99延迟远高于P50说明调度有问题大概率是长请求阻塞了短请求。这时候就要考虑优先级调度或者chunked prefill了。5. 常见问题与排查技巧实录5.1 显存OOM的排查路径显存OOM是AI工程里最高频的问题没有之一。排查思路要按这个顺序来第一确认是模型权重还是激活值还是KV Cache占的显存。用torch.cuda.memory_summary()能看到详细分布。如果是权重占大头考虑量化如果是激活值考虑梯度检查点或者减小batch如果是KV Cache考虑量化Cache或者限制最大序列长度。第二检查显存碎片。PyTorch的缓存分配器会产生碎片导致明明有足够空闲显存却分配失败。解决办法是设置PYTORCH_CUDA_ALLOC_CONFexpandable_segments:True让分配器能合并碎片。第三检查内存泄漏。最常见的是在循环里不断创建新的张量而没有释放或者把中间结果存到了全局变量里。用torch.cuda.memory_allocated()监控每一步的显存变化如果持续增长不下降基本就是泄漏。注意nvidia-smi显示的显存占用包含CUDA上下文和缓存比实际张量占用大不少。排查时以torch.cuda.memory_allocated()为准。5.2 生成结果异常的定位方法生成结果异常分几种重复、截断、乱码、不遵循指令。排查方法各不相同。重复问题通常是重复惩罚参数没设好或者采样策略有问题。如果用的是贪心解码重复是正常的换成top-p采样一般能解决。如果top-p也重复检查一下是不是模型本身的问题换个小模型对比一下。截断问题先看max_new_tokens是不是设小了再看eos_token是不是被错误地屏蔽了。有些模型的tokenizer会把eos加在特殊位置如果生成时没正确处理模型永远生成不出eos就会一直生成到max长度。不遵循指令的问题最复杂可能是prompt格式不对也可能是模型能力不够。先用官方推荐的prompt模板试一下如果还不行考虑换更大的模型或者做微调。5.3 性能不达预期的优化清单性能优化要按投入产出比排序我一般按这个顺序来优化项预期收益实现难度优先级KV Cache2-5倍吞吐低最高动态批处理3-10倍吞吐中高权重量化显存减半低高算子融合10-30%提速中中投机解码1.5-2倍提速高低张量并行线性扩展高按需先把KV Cache和动态批处理做了这两个是性价比最高的。量化看场景如果显存够用可以先不做。算子融合和投机解码属于锦上添花等前面都做完了再考虑。5.4 常见问题速查表现象可能原因排查方法解决方案首token延迟高prefill计算量大看prompt长度分布chunked prefill每token延迟高batch太大或Cache未命中看batch填充率和Cache命中率调小batch或优化Cache吞吐上不去GPU利用率低看GPU util和显存带宽增大batch或算子融合输出乱码tokenizer不匹配检查tokenizer和模型是否配套换正确的tokenizer显存缓慢增长内存泄漏监控memory_allocated检查全局变量和循环引用请求超时长请求阻塞看请求长度分布优先级调度或超时取消6. 我踩过的坑和给你的建议第一个坑是过早优化。我一开始就想着要做PagedAttention、要做投机解码结果基础版本都没跑通浪费了两周。后来老老实实从最简单的实现开始先跑通再优化反而更快。能跑通的烂代码比跑不通的好代码有价值一百倍。第二个坑是忽略评测。有段时间我一直在调性能吞吐确实上去了但生成质量下降了都不知道。后来补上评测才发现量化到INT4之后某些任务的准确率掉了5个点。性能和质量的权衡必须建立在量化评测的基础上凭感觉调参迟早翻车。第三个坑是不重视日志。线上出问题时如果没有详细的请求日志和性能日志根本无从查起。我现在每个请求都会记录请求ID、prompt长度、生成长度、TTFT、TPOT、使用的batch size、Cache命中情况。这些日志在排查问题时能省下大量时间。最后一个建议从零搭建的目的是理解不是替代。你不需要真的在生产环境用自己写的推理引擎但你需要知道vLLM、TensorRT-LLM这些引擎在做什么、为什么这么做。这样当它们出问题时你能快速定位当它们不满足需求时你能自己改。这个能力才是AI工程师真正的护城河。
网站建设高端定制企业官网