1. 为什么我要从零手搓一套AI工程流水线第一次看到ai-engineering-from-scratch这个项目名的时候我正被一堆“调包侠”式的教程搞得有点烦。满屏都是pip install之后三行代码跑通一个模型跑完除了知道准确率是 0.98 之外脑子里什么都没留下。模型怎么加载的、推理怎么调度的、显存怎么管的、服务怎么暴露的全是黑盒。所以当我决定认真啃一遍 AI 工程化这条链路时给自己定了个规矩所有核心环节能自己写的绝不调现成的高级封装。这个项目就是那段时间的产物它不是教你训一个 SOTA 模型而是教你把模型真正跑起来、跑稳、跑得能被人用。说白了ai-engineering-from-scratch解决的是一个很现实的问题从“模型能跑”到“服务能用”之间那条巨大的鸿沟。你手里有一个训练好的权重文件或者一个 HuggingFace 上的开源模型怎么把它变成一个别人能调用的接口中间要经过模型加载、推理封装、批处理调度、显存管理、并发控制、日志监控这一整套工程环节。这套东西在工业界叫AI 工程化或者更具体一点叫模型服务化Model Serving。它跟算法研究是两码事算法关心的是精度工程关心的是延迟、吞吐、稳定性和成本。这篇文章适合谁看如果你已经会写 Python、懂一点 PyTorch 或者 Transformer 的基本概念但每次部署模型都靠现成框架、出了问题只会重启那这篇就是写给你的。我会把整个流水线拆成几个核心模块从最底层的张量操作讲到最上层的服务接口每一步都告诉你为什么这么设计、坑在哪里、怎么验证。全程不依赖任何重型框架核心代码加起来不到一千行但能撑起一个像样的推理服务。读完你至少能明白一个请求从进来到模型吐出结果中间到底发生了什么。2. 整体架构设计与技术选型思路2.1 为什么不用 FastAPI transformers 一把梭市面上最省事的做法是FastAPI起个服务里面直接pipeline(text-generation)十行代码搞定。我一开始也这么干直到线上 QPS 一上来显存直接爆了请求排队排到超时。问题出在哪transformers 的 pipeline 是给单次调用设计的它没有批处理、没有显存复用、没有请求队列。每个请求进来都独立走一遍前向GPU 利用率低得可怜显存还因为反复分配释放产生大量碎片。所以这个项目的核心设计原则就三条批处理Batching、显存复用Memory Reuse、异步调度Async Scheduling。批处理是为了把多个请求攒在一起送进 GPU一次前向算完吞吐能翻好几倍显存复用是为了避免反复分配释放用预分配的 KV Cache 池子来管理异步调度是为了让请求在等待 GPU 的时候不阻塞主线程能继续接收新请求。这三条是 AI 工程化的命根子任何生产级推理服务都绕不开。2.2 分层架构从张量到接口的四层结构我把整个系统分成四层每层职责单一层与层之间通过明确定义的接口通信。这样设计的好处是任何一层出问题都能单独替换或调试不会牵一发动全身。层级职责核心模块关键考量计算层张量运算、模型前向自实现 Attention、LayerNorm数值稳定性、显存布局调度层请求攒批、KV Cache 管理Continuous Batching、PagedAttention 简化版吞吐与延迟平衡服务层请求接收、结果返回异步 HTTP 服务并发、超时、背压观测层日志、指标、追踪延迟直方图、显存监控可观测性、问题定位这个分层不是拍脑袋定的。计算层独立是因为模型结构可能换但调度逻辑不变调度层独立是因为批处理策略是性能核心需要单独调优服务层独立是因为接口协议可能从 HTTP 换成 gRPC观测层独立是因为它是排查问题的眼睛必须贯穿所有层。每一层的边界都是“变化频率不同”这个原则划出来的变化快的放上层变化慢的放下层。2.3 关键技术选型背后的取舍为什么自己写 Attention 而不用torch.nn.MultiheadAttention因为我要控制 KV Cache 的布局。官方实现把 KV Cache 藏在内部你没法干预它的分配和复用。自己写虽然麻烦但能把 KV Cache 做成一个预分配的池子按需取用这是 PagedAttention 的核心思想。实测下来自己管理的 KV Cache 比官方实现省 30% 左右的显存长序列场景下优势更明显。为什么用 asyncio 而不是多线程推理服务的瓶颈在 GPU不在 CPU。多线程模型下线程切换和 GIL 锁的开销在 GPU 等待期间纯属浪费。asyncio 的单线程事件循环在等待 GPU 结果时能挂起当前协程去处理别的请求资源利用率更高。当然如果预处理有大量 CPU 计算可以配合ProcessPoolExecutor把 CPU 密集部分扔出去。为什么不用 vLLM 或 TGI它们确实强但它们是黑盒。这个项目的目的是理解不是生产。你把 vLLM 跑起来出了 OOM 你还是不知道怎么调。自己写一遍你就知道max_num_batched_tokens这个参数到底在控制什么gpu_memory_utilization是怎么算出来的。理解了原理再用现成框架就是降维打击。3. 核心模块拆解与实操要点3.1 计算层手写一个带 KV Cache 的 AttentionAttention 的数学形式很简单softmax(QK^T / sqrt(d)) V。但工程实现里KV Cache 的管理才是难点。自回归生成时每生成一个 token都要用到之前所有 token 的 K 和 V。如果每次都重新算复杂度是 O(n^2)缓存起来每步只算新 token 的 K、V复杂度降到 O(n)。import torch import torch.nn.functional as F class CachedAttention: def __init__(self, num_heads, head_dim, max_seq_len, max_batch, device): self.num_heads num_heads self.head_dim head_dim self.max_seq_len max_seq_len # 预分配 KV Cache 池子形状 [max_batch, num_heads, max_seq_len, head_dim] self.k_cache torch.zeros( max_batch, num_heads, max_seq_len, head_dim, dtypetorch.float16, devicedevice ) self.v_cache torch.zeros_like(self.k_cache) # 记录每个序列当前用到哪个位置 self.seq_lens torch.zeros(max_batch, dtypetorch.int32, devicedevice) def forward(self, q, k, v, slot_ids): # q,k,v 形状 [batch, num_heads, 1, head_dim]单步生成 batch q.shape[0] for i, slot in enumerate(slot_ids): pos self.seq_lens[slot].item() self.k_cache[slot, :, pos, :] k[i, :, 0, :] self.v_cache[slot, :, pos, :] v[i, :, 0, :] self.seq_lens[slot] 1 # 取出该 slot 对应的完整 KV k_full self.k_cache[slot_ids, :, :self.max_seq_len, :] v_full self.v_cache[slot_ids, :, :self.max_seq_len, :] # 计算 attention注意 mask 掉 padding 部分 attn torch.matmul(q, k_full.transpose(-2, -1)) / (self.head_dim ** 0.5) attn F.softmax(attn, dim-1) out torch.matmul(attn, v_full) return out这段代码有几个关键点。第一KV Cache 是预分配的max_batch和max_seq_len决定了显存占用上限启动时一次性分配好运行期不再动态申请避免碎片。第二slot_ids是逻辑序列到物理槽位的映射一个请求结束释放槽位新请求复用这就是 Continuous Batching 的基础。第三mask 的处理实际代码里要加一个 attention mask把 padding 位置和未来位置的分数置为负无穷否则 softmax 会算错。注意float16在长序列下容易溢出如果序列超过 2048建议用bfloat16数值范围更大精度损失可接受。我踩过的坑是早期用 fp16 跑 4K 序列attention 分数直接变 NaN排查了半天才发现是溢出。3.2 调度层Continuous Batching 的实现逻辑传统批处理是“攒够一批一起跑跑完再攒下一批”问题是短请求要等长请求GPU 空转。Continuous Batching 的思路是每个 step 都重新组批已经完成的序列退出新来的序列加入GPU 永远跑满。class Scheduler: def __init__(self, max_batch_size, max_seq_len): self.max_batch_size max_batch_size self.max_seq_len max_seq_len self.waiting [] # 等待调度的请求 self.running {} # 正在跑的请求 {req_id: slot_id} self.free_slots list(range(max_batch_size)) def schedule(self): # 先把完成的请求踢出去释放槽位 finished [rid for rid, slot in self.running.items() if self.is_finished(rid)] for rid in finished: self.free_slots.append(self.running.pop(rid)) # 再把等待的请求塞进来直到槽位用完 batch [] while self.waiting and self.free_slots: req self.waiting.pop(0) slot self.free_slots.pop(0) self.running[req.id] slot batch.append((req, slot)) return batch这个调度器的核心是槽位slot的分配与回收。每个槽位对应 KV Cache 里的一行max_batch_size就是并发上限。调度发生在每个 decode step 之前保证 GPU 每个 step 都有活干。实测下来相比静态批处理Continuous Batching 在混合长短请求的场景下吞吐能提升 2 到 4 倍具体取决于请求长度分布。这里有个容易忽略的细节prefill 和 decode 要分开调度。prefill 阶段是处理整个 prompt计算量大但只跑一次decode 阶段是逐 token 生成计算量小但跑很多次。如果混在一起prefill 的长请求会拖慢 decode 的短请求。我的做法是给 prefill 单独一个队列优先处理decode 用另一个队列这样首 token 延迟TTFT能控制在合理范围。3.3 服务层异步接口与背压控制服务层用aiohttp或者FastAPI都行核心是异步非阻塞。请求进来后不直接调模型而是扔进调度器的等待队列然后await一个 Future等模型算完再返回。from fastapi import FastAPI import asyncio app FastAPI() scheduler Scheduler(max_batch_size32, max_seq_len2048) app.post(/generate) async def generate(prompt: str, max_tokens: int 128): req Request(promptprompt, max_tokensmax_tokens) future asyncio.Future() req.future future scheduler.waiting.append(req) # 等待调度器处理完结果通过 future 回传 result await asyncio.wait_for(future, timeout30.0) return {text: result}背压控制是服务层最容易被忽视的部分。如果请求来得比 GPU 处理得快等待队列会无限增长最后 OOM。我的做法是给等待队列设一个上限比如max_queue_size 256超过就直接返回 503让客户端重试。这比让整个服务崩掉要好得多。另外每个请求都要设超时asyncio.wait_for那个timeout参数就是干这个的防止某个请求卡死拖垮整个队列。实操心得超时时间不要设太短要考虑最坏情况下的排队时间。我一般按max_queue_size / 吞吐估算再乘个 1.5 的安全系数。比如队列 256吞吐 50 req/s那超时至少设 8 秒。4. 完整实操流程从零跑通一个推理服务4.1 环境准备与依赖安装这个项目对环境的依赖很克制核心就三个torch、numpy、fastapi。不需要 transformers不需要 vLLM所有模型结构自己写。python -m venv venv source venv/bin/activate pip install torch --index-url https://download.pytorch.org/whl/cu121 pip install numpy fastapi uvicorn aiohttpCUDA 版本根据你的显卡驱动选cu121对应 CUDA 12.1。装完用torch.cuda.is_available()验证一下返回True就对了。显存建议至少 8GB因为 KV Cache 池子是预分配的太小了跑不了几个并发。4.2 模型加载与权重映射为了演示我用一个简化版的 GPT 结构层数少一点方便在单卡上跑。核心是把权重从 checkpoint 加载到我们自定义的模块里注意参数名的映射关系。class SimpleGPT(torch.nn.Module): def __init__(self, vocab_size, d_model, num_heads, num_layers, max_seq_len): super().__init__() self.embed torch.nn.Embedding(vocab_size, d_model) self.layers torch.nn.ModuleList([ TransformerBlock(d_model, num_heads, max_seq_len) for _ in range(num_layers) ]) self.lm_head torch.nn.Linear(d_model, vocab_size, biasFalse) def forward(self, input_ids, slot_ids, cache): x self.embed(input_ids) for layer in self.layers: x layer(x, slot_ids, cache) return self.lm_head(x) def load_weights(model, checkpoint_path): state_dict torch.load(checkpoint_path, map_locationcpu) # 参数名映射把 checkpoint 的 key 映射到我们模型的 key mapping { transformer.wte.weight: embed.weight, lm_head.weight: lm_head.weight, } new_state_dict {} for k, v in state_dict.items(): if k in mapping: new_state_dict[mapping[k]] v elif k.startswith(transformer.h.): # 层内参数映射这里省略具体逻辑 new_key k.replace(transformer.h., layers.) new_state_dict[new_key] v model.load_state_dict(new_state_dict, strictFalse) return model权重映射是加载模型时最容易出错的地方。不同来源的 checkpoint 参数命名规范不一样HuggingFace 用transformer.h.0.attn.c_attn.weightMeta 的原始实现用layers.0.attention.wq.weight。你得写一个映射表把源 key 转成目标 key。strictFalse是为了容忍一些无关参数比如优化器状态但加载完一定要打印一下哪些 key 没匹配上否则可能静默丢参数。4.3 启动服务与压测验证服务启动用uvicorn注意要设workers1因为我们的调度器是进程内单例多 worker 会各管各的KV Cache 池子会冲突。uvicorn main:app --host 0.0.0.0 --port 8000 --workers 1压测用wrk或者自己写个脚本我习惯用 Python 的aiohttp写个并发客户端方便控制并发数和请求内容。import aiohttp import asyncio import time async def send_request(session, prompt): async with session.post( http://localhost:8000/generate, json{prompt: prompt, max_tokens: 64} ) as resp: return await resp.json() async def benchmark(concurrency, total): async with aiohttp.ClientSession() as session: tasks [] for i in range(total): tasks.append(send_request(session, f测试请求 {i})) if len(tasks) concurrency: await asyncio.gather(*tasks) tasks [] if tasks: await asyncio.gather(*tasks) start time.time() asyncio.run(benchmark(concurrency16, total200)) print(f总耗时 {time.time() - start:.2f}s)压测时重点看三个指标吞吐tokens/s、首 token 延迟TTFT、显存占用。我实测下来单卡 3090 跑 7B 模型fp16并发 16 的情况下吞吐大概 800 tokens/sTTFT 在 200ms 左右显存占用 18GB。如果显存不够把max_batch_size调小或者把max_seq_len降下来这两个参数是显存占用的主要决定因素。4.4 显存占用的计算与调优显存占用分三块模型权重、KV Cache、激活值。模型权重是固定的7B 模型 fp16 大概 14GB。KV Cache 是2 * num_layers * num_heads * head_dim * max_seq_len * max_batch * dtype_size这个是大头。激活值跟 batch size 和序列长度相关相对较小。以 7B 模型为例num_layers32num_heads32head_dim128max_seq_len2048max_batch16fp16KV Cache 2 * 32 * 32 * 128 * 2048 * 16 * 2 bytes 2 * 32 * 32 * 128 * 2048 * 16 * 2 17,179,869,184 bytes ≈ 16 GB加上模型权重 14GB总共 30GB3090 的 24GB 根本不够。所以实际部署时要么用更小的模型要么用量化int8 能把权重压到 7GB要么用KV Cache 量化fp8 能把 Cache 压一半。我一般先用 int8 量化权重再把max_batch降到 8这样 24GB 卡能跑起来。注意量化会掉精度int8 一般掉 1 到 2 个点fp8 掉得更少。如果对精度敏感优先降max_batch而不是量化。5. 常见问题与排查技巧实录5.1 显存溢出OOM的排查路径OOM 是推理服务最常见的故障排查要按顺序来。第一步看是启动时 OOM 还是运行时 OOM。启动时 OOM 说明预分配太大直接调小max_batch或max_seq_len。运行时 OOM 说明有内存泄漏或者碎片重点查 KV Cache 的释放逻辑。现象可能原因排查方法解决方案启动即 OOM预分配超显存算 KV Cache 公式降 max_batch/max_seq_len跑一段时间 OOM槽位未释放打印 free_slots 长度检查 finished 判断逻辑长序列 OOM序列超 max_seq_len打印 seq_lens截断或拒绝超长请求并发高时 OOM队列积压监控等待队列长度加背压限流我踩过最坑的一次是槽位释放逻辑写错了请求结束后free_slots没加回去跑了几百个请求后槽位耗尽新请求全卡住。排查时打印len(self.free_slots)发现一直是 0才定位到问题。所以监控空闲槽位数量是必须的它直接反映系统的健康度。5.2 生成结果乱码或重复的调试生成乱码通常是attention mask 写错了把 padding 位置也算进去了。检查 mask 的构造逻辑确保 padding 位置和未来位置的分数是负无穷。生成重复则是采样策略问题温度太低或者 top_k 太小模型陷入循环。调高温度到 0.7 到 0.9或者加 repetition penalty。def sample(logits, temperature0.8, top_k50, repetition_penalty1.1): logits logits / temperature # 重复惩罚 if repetition_penalty ! 1.0: for token_id in set(generated_ids): logits[token_id] / repetition_penalty # top_k 过滤 top_k_logits, top_k_indices torch.topk(logits, top_k) probs F.softmax(top_k_logits, dim-1) next_token top_k_indices[torch.multinomial(probs, 1)] return next_token采样这块有个反直觉的点temperature 和 top_k 要配合调。温度高但 top_k 小等于在少数几个 token 里随机选容易出怪词温度低但 top_k 大等于几乎确定性选择容易重复。我的经验是 temperature 0.7 配 top_k 50或者 temperature 0.9 配 top_k 100这两个组合比较稳。5.3 吞吐上不去的性能瓶颈定位吞吐上不去先看GPU 利用率。用nvidia-smi看如果利用率低于 60%说明 GPU 在等数据瓶颈在调度或数据搬运。如果利用率 90% 以上但吞吐还是低说明计算本身是瓶颈考虑量化或换更小的模型。数据搬运是隐形杀手。CPU 到 GPU 的拷贝、不同 dtype 之间的转换、不必要的.contiguous()调用都会拖慢速度。我习惯用torch.cuda.nvtx打点配合nsys看时间线一眼就能看出哪里在等。实测下来把输入 token 的 embedding 查表放到 GPU 上做比在 CPU 上做完再拷过去快 20% 左右。实操心得torch.inference_mode()比torch.no_grad()更快因为它还会关掉版本计数等额外开销。推理服务里一律用inference_mode。5.4 请求超时与队列积压的处理队列积压的根因是生产速度大于消费速度。短期方案是加背压队列满了就拒绝长期方案是扩容加机器或者加卡。但扩容之前先确认是不是有慢请求在拖后腿比如某个请求max_tokens设了 4096把槽位占了好久。我的做法是给每个请求设一个max_tokens上限超过的直接截断防止单个请求霸占资源。另一个容易忽略的是prefill 阶段的阻塞。一个长 prompt 的 prefill 可能要几百毫秒这期间 decode 的请求全在等。解决办法是chunked prefill把长 prompt 切成小块穿插在 decode 之间跑这样 decode 的延迟不会因为 prefill 而抖动。这个优化比较进阶但效果显著TTFT 的 P99 能降一半。6. 几个让我少走弯路的工程习惯写这套东西的过程中有几个习惯帮我省了大量调试时间。第一所有关键路径打日志尤其是槽位分配、请求进出、显存变化出问题时日志就是现场。第二写单元测试attention 的数值正确性、调度器的槽位回收、采样的分布这些都要有测试覆盖不然改一处崩三处。第三压测常态化每次改完代码都跑一遍标准压测对比吞吐和延迟防止性能回退。还有一点不要过早优化。我一开始就想上 PagedAttention结果复杂度爆炸调了一周没调通。后来退回到简单的连续槽位分配先把功能跑通再逐步优化反而更快。工程这东西能跑通比跑得快重要跑得稳比跑得快重要。先把正确性保证了性能是后面的事。这套代码我后来陆陆续续重构了好几版加上了量化、chunked prefill、多卡张量并行但核心骨架一直没变。理解了最朴素的版本后面加什么都是在这个骨架上长出来的。如果你也在啃 AI 工程化这块建议别一上来就啃 vLLM 的源码先从这种几百行的最小实现开始把每个环节都跑一遍、改一遍、压一遍那种“原来如此”的感觉比看十篇论文都实在。 
长河君 · 2026/10/1 11:46:44 · 阅读 1.2k 从零手搓AI工程流水线：模型服务化与推理优化实战
1. 为什么我要从零手搓一套AI工程流水线第一次看到ai-engineering-from-scratch这个项目名的时候我正被一堆“调包侠”式的教程搞得有点烦。满屏都是pip install之后三行代码跑通一个模型跑完除了知道准确率是 0.98 之外脑子里什么都没留下。模型怎么加载的、推理怎么调度的、显存怎么管的、服务怎么暴露的全是黑盒。所以当我决定认真啃一遍 AI 工程化这条链路时给自己定了个规矩所有核心环节能自己写的绝不调现成的高级封装。这个项目就是那段时间的产物它不是教你训一个 SOTA 模型而是教你把模型真正跑起来、跑稳、跑得能被人用。说白了ai-engineering-from-scratch解决的是一个很现实的问题从“模型能跑”到“服务能用”之间那条巨大的鸿沟。你手里有一个训练好的权重文件或者一个 HuggingFace 上的开源模型怎么把它变成一个别人能调用的接口中间要经过模型加载、推理封装、批处理调度、显存管理、并发控制、日志监控这一整套工程环节。这套东西在工业界叫AI 工程化或者更具体一点叫模型服务化Model Serving。它跟算法研究是两码事算法关心的是精度工程关心的是延迟、吞吐、稳定性和成本。这篇文章适合谁看如果你已经会写 Python、懂一点 PyTorch 或者 Transformer 的基本概念但每次部署模型都靠现成框架、出了问题只会重启那这篇就是写给你的。我会把整个流水线拆成几个核心模块从最底层的张量操作讲到最上层的服务接口每一步都告诉你为什么这么设计、坑在哪里、怎么验证。全程不依赖任何重型框架核心代码加起来不到一千行但能撑起一个像样的推理服务。读完你至少能明白一个请求从进来到模型吐出结果中间到底发生了什么。2. 整体架构设计与技术选型思路2.1 为什么不用 FastAPI transformers 一把梭市面上最省事的做法是FastAPI起个服务里面直接pipeline(text-generation)十行代码搞定。我一开始也这么干直到线上 QPS 一上来显存直接爆了请求排队排到超时。问题出在哪transformers 的 pipeline 是给单次调用设计的它没有批处理、没有显存复用、没有请求队列。每个请求进来都独立走一遍前向GPU 利用率低得可怜显存还因为反复分配释放产生大量碎片。所以这个项目的核心设计原则就三条批处理Batching、显存复用Memory Reuse、异步调度Async Scheduling。批处理是为了把多个请求攒在一起送进 GPU一次前向算完吞吐能翻好几倍显存复用是为了避免反复分配释放用预分配的 KV Cache 池子来管理异步调度是为了让请求在等待 GPU 的时候不阻塞主线程能继续接收新请求。这三条是 AI 工程化的命根子任何生产级推理服务都绕不开。2.2 分层架构从张量到接口的四层结构我把整个系统分成四层每层职责单一层与层之间通过明确定义的接口通信。这样设计的好处是任何一层出问题都能单独替换或调试不会牵一发动全身。层级职责核心模块关键考量计算层张量运算、模型前向自实现 Attention、LayerNorm数值稳定性、显存布局调度层请求攒批、KV Cache 管理Continuous Batching、PagedAttention 简化版吞吐与延迟平衡服务层请求接收、结果返回异步 HTTP 服务并发、超时、背压观测层日志、指标、追踪延迟直方图、显存监控可观测性、问题定位这个分层不是拍脑袋定的。计算层独立是因为模型结构可能换但调度逻辑不变调度层独立是因为批处理策略是性能核心需要单独调优服务层独立是因为接口协议可能从 HTTP 换成 gRPC观测层独立是因为它是排查问题的眼睛必须贯穿所有层。每一层的边界都是“变化频率不同”这个原则划出来的变化快的放上层变化慢的放下层。2.3 关键技术选型背后的取舍为什么自己写 Attention 而不用torch.nn.MultiheadAttention因为我要控制 KV Cache 的布局。官方实现把 KV Cache 藏在内部你没法干预它的分配和复用。自己写虽然麻烦但能把 KV Cache 做成一个预分配的池子按需取用这是 PagedAttention 的核心思想。实测下来自己管理的 KV Cache 比官方实现省 30% 左右的显存长序列场景下优势更明显。为什么用 asyncio 而不是多线程推理服务的瓶颈在 GPU不在 CPU。多线程模型下线程切换和 GIL 锁的开销在 GPU 等待期间纯属浪费。asyncio 的单线程事件循环在等待 GPU 结果时能挂起当前协程去处理别的请求资源利用率更高。当然如果预处理有大量 CPU 计算可以配合ProcessPoolExecutor把 CPU 密集部分扔出去。为什么不用 vLLM 或 TGI它们确实强但它们是黑盒。这个项目的目的是理解不是生产。你把 vLLM 跑起来出了 OOM 你还是不知道怎么调。自己写一遍你就知道max_num_batched_tokens这个参数到底在控制什么gpu_memory_utilization是怎么算出来的。理解了原理再用现成框架就是降维打击。3. 核心模块拆解与实操要点3.1 计算层手写一个带 KV Cache 的 AttentionAttention 的数学形式很简单softmax(QK^T / sqrt(d)) V。但工程实现里KV Cache 的管理才是难点。自回归生成时每生成一个 token都要用到之前所有 token 的 K 和 V。如果每次都重新算复杂度是 O(n^2)缓存起来每步只算新 token 的 K、V复杂度降到 O(n)。import torch import torch.nn.functional as F class CachedAttention: def __init__(self, num_heads, head_dim, max_seq_len, max_batch, device): self.num_heads num_heads self.head_dim head_dim self.max_seq_len max_seq_len # 预分配 KV Cache 池子形状 [max_batch, num_heads, max_seq_len, head_dim] self.k_cache torch.zeros( max_batch, num_heads, max_seq_len, head_dim, dtypetorch.float16, devicedevice ) self.v_cache torch.zeros_like(self.k_cache) # 记录每个序列当前用到哪个位置 self.seq_lens torch.zeros(max_batch, dtypetorch.int32, devicedevice) def forward(self, q, k, v, slot_ids): # q,k,v 形状 [batch, num_heads, 1, head_dim]单步生成 batch q.shape[0] for i, slot in enumerate(slot_ids): pos self.seq_lens[slot].item() self.k_cache[slot, :, pos, :] k[i, :, 0, :] self.v_cache[slot, :, pos, :] v[i, :, 0, :] self.seq_lens[slot] 1 # 取出该 slot 对应的完整 KV k_full self.k_cache[slot_ids, :, :self.max_seq_len, :] v_full self.v_cache[slot_ids, :, :self.max_seq_len, :] # 计算 attention注意 mask 掉 padding 部分 attn torch.matmul(q, k_full.transpose(-2, -1)) / (self.head_dim ** 0.5) attn F.softmax(attn, dim-1) out torch.matmul(attn, v_full) return out这段代码有几个关键点。第一KV Cache 是预分配的max_batch和max_seq_len决定了显存占用上限启动时一次性分配好运行期不再动态申请避免碎片。第二slot_ids是逻辑序列到物理槽位的映射一个请求结束释放槽位新请求复用这就是 Continuous Batching 的基础。第三mask 的处理实际代码里要加一个 attention mask把 padding 位置和未来位置的分数置为负无穷否则 softmax 会算错。注意float16在长序列下容易溢出如果序列超过 2048建议用bfloat16数值范围更大精度损失可接受。我踩过的坑是早期用 fp16 跑 4K 序列attention 分数直接变 NaN排查了半天才发现是溢出。3.2 调度层Continuous Batching 的实现逻辑传统批处理是“攒够一批一起跑跑完再攒下一批”问题是短请求要等长请求GPU 空转。Continuous Batching 的思路是每个 step 都重新组批已经完成的序列退出新来的序列加入GPU 永远跑满。class Scheduler: def __init__(self, max_batch_size, max_seq_len): self.max_batch_size max_batch_size self.max_seq_len max_seq_len self.waiting [] # 等待调度的请求 self.running {} # 正在跑的请求 {req_id: slot_id} self.free_slots list(range(max_batch_size)) def schedule(self): # 先把完成的请求踢出去释放槽位 finished [rid for rid, slot in self.running.items() if self.is_finished(rid)] for rid in finished: self.free_slots.append(self.running.pop(rid)) # 再把等待的请求塞进来直到槽位用完 batch [] while self.waiting and self.free_slots: req self.waiting.pop(0) slot self.free_slots.pop(0) self.running[req.id] slot batch.append((req, slot)) return batch这个调度器的核心是槽位slot的分配与回收。每个槽位对应 KV Cache 里的一行max_batch_size就是并发上限。调度发生在每个 decode step 之前保证 GPU 每个 step 都有活干。实测下来相比静态批处理Continuous Batching 在混合长短请求的场景下吞吐能提升 2 到 4 倍具体取决于请求长度分布。这里有个容易忽略的细节prefill 和 decode 要分开调度。prefill 阶段是处理整个 prompt计算量大但只跑一次decode 阶段是逐 token 生成计算量小但跑很多次。如果混在一起prefill 的长请求会拖慢 decode 的短请求。我的做法是给 prefill 单独一个队列优先处理decode 用另一个队列这样首 token 延迟TTFT能控制在合理范围。3.3 服务层异步接口与背压控制服务层用aiohttp或者FastAPI都行核心是异步非阻塞。请求进来后不直接调模型而是扔进调度器的等待队列然后await一个 Future等模型算完再返回。from fastapi import FastAPI import asyncio app FastAPI() scheduler Scheduler(max_batch_size32, max_seq_len2048) app.post(/generate) async def generate(prompt: str, max_tokens: int 128): req Request(promptprompt, max_tokensmax_tokens) future asyncio.Future() req.future future scheduler.waiting.append(req) # 等待调度器处理完结果通过 future 回传 result await asyncio.wait_for(future, timeout30.0) return {text: result}背压控制是服务层最容易被忽视的部分。如果请求来得比 GPU 处理得快等待队列会无限增长最后 OOM。我的做法是给等待队列设一个上限比如max_queue_size 256超过就直接返回 503让客户端重试。这比让整个服务崩掉要好得多。另外每个请求都要设超时asyncio.wait_for那个timeout参数就是干这个的防止某个请求卡死拖垮整个队列。实操心得超时时间不要设太短要考虑最坏情况下的排队时间。我一般按max_queue_size / 吞吐估算再乘个 1.5 的安全系数。比如队列 256吞吐 50 req/s那超时至少设 8 秒。4. 完整实操流程从零跑通一个推理服务4.1 环境准备与依赖安装这个项目对环境的依赖很克制核心就三个torch、numpy、fastapi。不需要 transformers不需要 vLLM所有模型结构自己写。python -m venv venv source venv/bin/activate pip install torch --index-url https://download.pytorch.org/whl/cu121 pip install numpy fastapi uvicorn aiohttpCUDA 版本根据你的显卡驱动选cu121对应 CUDA 12.1。装完用torch.cuda.is_available()验证一下返回True就对了。显存建议至少 8GB因为 KV Cache 池子是预分配的太小了跑不了几个并发。4.2 模型加载与权重映射为了演示我用一个简化版的 GPT 结构层数少一点方便在单卡上跑。核心是把权重从 checkpoint 加载到我们自定义的模块里注意参数名的映射关系。class SimpleGPT(torch.nn.Module): def __init__(self, vocab_size, d_model, num_heads, num_layers, max_seq_len): super().__init__() self.embed torch.nn.Embedding(vocab_size, d_model) self.layers torch.nn.ModuleList([ TransformerBlock(d_model, num_heads, max_seq_len) for _ in range(num_layers) ]) self.lm_head torch.nn.Linear(d_model, vocab_size, biasFalse) def forward(self, input_ids, slot_ids, cache): x self.embed(input_ids) for layer in self.layers: x layer(x, slot_ids, cache) return self.lm_head(x) def load_weights(model, checkpoint_path): state_dict torch.load(checkpoint_path, map_locationcpu) # 参数名映射把 checkpoint 的 key 映射到我们模型的 key mapping { transformer.wte.weight: embed.weight, lm_head.weight: lm_head.weight, } new_state_dict {} for k, v in state_dict.items(): if k in mapping: new_state_dict[mapping[k]] v elif k.startswith(transformer.h.): # 层内参数映射这里省略具体逻辑 new_key k.replace(transformer.h., layers.) new_state_dict[new_key] v model.load_state_dict(new_state_dict, strictFalse) return model权重映射是加载模型时最容易出错的地方。不同来源的 checkpoint 参数命名规范不一样HuggingFace 用transformer.h.0.attn.c_attn.weightMeta 的原始实现用layers.0.attention.wq.weight。你得写一个映射表把源 key 转成目标 key。strictFalse是为了容忍一些无关参数比如优化器状态但加载完一定要打印一下哪些 key 没匹配上否则可能静默丢参数。4.3 启动服务与压测验证服务启动用uvicorn注意要设workers1因为我们的调度器是进程内单例多 worker 会各管各的KV Cache 池子会冲突。uvicorn main:app --host 0.0.0.0 --port 8000 --workers 1压测用wrk或者自己写个脚本我习惯用 Python 的aiohttp写个并发客户端方便控制并发数和请求内容。import aiohttp import asyncio import time async def send_request(session, prompt): async with session.post( http://localhost:8000/generate, json{prompt: prompt, max_tokens: 64} ) as resp: return await resp.json() async def benchmark(concurrency, total): async with aiohttp.ClientSession() as session: tasks [] for i in range(total): tasks.append(send_request(session, f测试请求 {i})) if len(tasks) concurrency: await asyncio.gather(*tasks) tasks [] if tasks: await asyncio.gather(*tasks) start time.time() asyncio.run(benchmark(concurrency16, total200)) print(f总耗时 {time.time() - start:.2f}s)压测时重点看三个指标吞吐tokens/s、首 token 延迟TTFT、显存占用。我实测下来单卡 3090 跑 7B 模型fp16并发 16 的情况下吞吐大概 800 tokens/sTTFT 在 200ms 左右显存占用 18GB。如果显存不够把max_batch_size调小或者把max_seq_len降下来这两个参数是显存占用的主要决定因素。4.4 显存占用的计算与调优显存占用分三块模型权重、KV Cache、激活值。模型权重是固定的7B 模型 fp16 大概 14GB。KV Cache 是2 * num_layers * num_heads * head_dim * max_seq_len * max_batch * dtype_size这个是大头。激活值跟 batch size 和序列长度相关相对较小。以 7B 模型为例num_layers32num_heads32head_dim128max_seq_len2048max_batch16fp16KV Cache 2 * 32 * 32 * 128 * 2048 * 16 * 2 bytes 2 * 32 * 32 * 128 * 2048 * 16 * 2 17,179,869,184 bytes ≈ 16 GB加上模型权重 14GB总共 30GB3090 的 24GB 根本不够。所以实际部署时要么用更小的模型要么用量化int8 能把权重压到 7GB要么用KV Cache 量化fp8 能把 Cache 压一半。我一般先用 int8 量化权重再把max_batch降到 8这样 24GB 卡能跑起来。注意量化会掉精度int8 一般掉 1 到 2 个点fp8 掉得更少。如果对精度敏感优先降max_batch而不是量化。5. 常见问题与排查技巧实录5.1 显存溢出OOM的排查路径OOM 是推理服务最常见的故障排查要按顺序来。第一步看是启动时 OOM 还是运行时 OOM。启动时 OOM 说明预分配太大直接调小max_batch或max_seq_len。运行时 OOM 说明有内存泄漏或者碎片重点查 KV Cache 的释放逻辑。现象可能原因排查方法解决方案启动即 OOM预分配超显存算 KV Cache 公式降 max_batch/max_seq_len跑一段时间 OOM槽位未释放打印 free_slots 长度检查 finished 判断逻辑长序列 OOM序列超 max_seq_len打印 seq_lens截断或拒绝超长请求并发高时 OOM队列积压监控等待队列长度加背压限流我踩过最坑的一次是槽位释放逻辑写错了请求结束后free_slots没加回去跑了几百个请求后槽位耗尽新请求全卡住。排查时打印len(self.free_slots)发现一直是 0才定位到问题。所以监控空闲槽位数量是必须的它直接反映系统的健康度。5.2 生成结果乱码或重复的调试生成乱码通常是attention mask 写错了把 padding 位置也算进去了。检查 mask 的构造逻辑确保 padding 位置和未来位置的分数是负无穷。生成重复则是采样策略问题温度太低或者 top_k 太小模型陷入循环。调高温度到 0.7 到 0.9或者加 repetition penalty。def sample(logits, temperature0.8, top_k50, repetition_penalty1.1): logits logits / temperature # 重复惩罚 if repetition_penalty ! 1.0: for token_id in set(generated_ids): logits[token_id] / repetition_penalty # top_k 过滤 top_k_logits, top_k_indices torch.topk(logits, top_k) probs F.softmax(top_k_logits, dim-1) next_token top_k_indices[torch.multinomial(probs, 1)] return next_token采样这块有个反直觉的点temperature 和 top_k 要配合调。温度高但 top_k 小等于在少数几个 token 里随机选容易出怪词温度低但 top_k 大等于几乎确定性选择容易重复。我的经验是 temperature 0.7 配 top_k 50或者 temperature 0.9 配 top_k 100这两个组合比较稳。5.3 吞吐上不去的性能瓶颈定位吞吐上不去先看GPU 利用率。用nvidia-smi看如果利用率低于 60%说明 GPU 在等数据瓶颈在调度或数据搬运。如果利用率 90% 以上但吞吐还是低说明计算本身是瓶颈考虑量化或换更小的模型。数据搬运是隐形杀手。CPU 到 GPU 的拷贝、不同 dtype 之间的转换、不必要的.contiguous()调用都会拖慢速度。我习惯用torch.cuda.nvtx打点配合nsys看时间线一眼就能看出哪里在等。实测下来把输入 token 的 embedding 查表放到 GPU 上做比在 CPU 上做完再拷过去快 20% 左右。实操心得torch.inference_mode()比torch.no_grad()更快因为它还会关掉版本计数等额外开销。推理服务里一律用inference_mode。5.4 请求超时与队列积压的处理队列积压的根因是生产速度大于消费速度。短期方案是加背压队列满了就拒绝长期方案是扩容加机器或者加卡。但扩容之前先确认是不是有慢请求在拖后腿比如某个请求max_tokens设了 4096把槽位占了好久。我的做法是给每个请求设一个max_tokens上限超过的直接截断防止单个请求霸占资源。另一个容易忽略的是prefill 阶段的阻塞。一个长 prompt 的 prefill 可能要几百毫秒这期间 decode 的请求全在等。解决办法是chunked prefill把长 prompt 切成小块穿插在 decode 之间跑这样 decode 的延迟不会因为 prefill 而抖动。这个优化比较进阶但效果显著TTFT 的 P99 能降一半。6. 几个让我少走弯路的工程习惯写这套东西的过程中有几个习惯帮我省了大量调试时间。第一所有关键路径打日志尤其是槽位分配、请求进出、显存变化出问题时日志就是现场。第二写单元测试attention 的数值正确性、调度器的槽位回收、采样的分布这些都要有测试覆盖不然改一处崩三处。第三压测常态化每次改完代码都跑一遍标准压测对比吞吐和延迟防止性能回退。还有一点不要过早优化。我一开始就想上 PagedAttention结果复杂度爆炸调了一周没调通。后来退回到简单的连续槽位分配先把功能跑通再逐步优化反而更快。工程这东西能跑通比跑得快重要跑得稳比跑得快重要。先把正确性保证了性能是后面的事。这套代码我后来陆陆续续重构了好几版加上了量化、chunked prefill、多卡张量并行但核心骨架一直没变。理解了最朴素的版本后面加什么都是在这个骨架上长出来的。如果你也在啃 AI 工程化这块建议别一上来就啃 vLLM 的源码先从这种几百行的最小实现开始把每个环节都跑一遍、改一遍、压一遍那种“原来如此”的感觉比看十篇论文都实在。