从零手搓AI工程框架：推理引擎、批处理与显存管理实战

1. 为什么我要从零手搓一套AI工程框架第一次看到ai-engineering-from-scratch这个标题我脑子里蹦出来的不是“又一个教程仓库”而是过去两年带团队踩过的那些坑。市面上讲AI工程的文章十篇里有八篇在教你调transformers的pipeline剩下两篇在讲怎么用现成的编排框架搭个RAG。但真到了生产环境模型加载慢、显存炸了、推理延迟抖动、批处理吞吐上不去、服务一重启就雪崩——这些问题没有一个能靠调库解决。ai-engineering-from-scratch这个项目标题核心诉求非常明确不依赖高层封装从底层把AI工程的关键环节自己实现一遍。它面向的不是刚学完Python语法的新手而是那些已经会用PyTorch或TensorFlow跑通demo、但一到部署和优化就抓瞎的工程师。这个项目要解决的核心问题是当框架帮你隐藏了所有细节你反而失去了对性能、内存、并发的掌控力。通过从零实现推理引擎、批处理调度、显存管理、服务化封装这些模块你能真正理解一个AI系统从训练完成到对外提供服务之间到底发生了什么。我自己的经历很典型。2022年做一个实时推荐系统模型本身不大但QPS一上来GPU利用率死活上不去延迟从20ms飙到200ms。排查了两周最后发现是数据预处理和模型推理串行执行GPU大部分时间在等CPU。这种问题你翻遍框架文档都找不到答案因为框架默认假设你懂这些。ai-engineering-from-scratch的价值就在这——它逼你把每一层都拆开看清楚数据怎么流动、内存怎么分配、请求怎么排队。适合谁来参考三类人一是想从算法工程师转型AI系统工程师的二是需要自己搭建推理服务但不想被框架绑架的三是面试时被问到“如何优化模型推理性能”却只能背八股文的。这篇文章我会按照这个项目的核心思路把AI工程从零实现的关键环节拆解一遍包括整体架构设计、核心模块实现、实操步骤、以及我踩过的那些坑。2. 整体架构设计与技术选型思路2.1 为什么选择“从零实现”而不是“基于框架二次开发”很多人第一反应是有现成的TensorRT、ONNX Runtime、Triton为什么要自己写这个问题我认真想过。直接说结论从零实现的目的不是替代生产级框架而是建立对底层机制的直觉。你用Triton部署模型遇到性能瓶颈时如果不知道它内部怎么做批处理、怎么做内存池、怎么做CUDA流调度你连调优的方向都找不到。这个项目的架构设计遵循一个原则每一层都只依赖下一层的抽象不跳级。最底层是张量计算和内存管理往上是模型加载和算子调度再往上是批处理策略和请求队列最顶层才是HTTP/gRPC服务接口。这样设计的好处是你可以在任何一层做替换或优化而不影响其他层。比如你想把批处理策略从静态批处理换成连续批处理只需要改中间那一层服务接口完全不用动。技术选型上我建议用Python做胶水层用C或CUDA做计算密集型部分。但如果你只是想理解原理纯Python加NumPy也够用只是性能会差很多。这个项目的重点不是性能数字而是让你亲手实现一遍那些框架帮你做的事。比如显存池你可以先用一个简单的字典模拟记录每个张量的引用计数等引用归零再回收。虽然简陋但你能直观感受到为什么显存碎片会导致OOM。2.2 核心模块划分与依赖关系整个系统我把它拆成五个核心模块每个模块的职责和依赖关系如下模块名称核心职责依赖模块关键输出张量抽象层管理内存分配、设备转移、生命周期无统一的Tensor对象模型加载器解析模型文件、初始化权重、构建计算图张量抽象层可执行的模型实例推理调度器请求排队、动态批处理、超时控制模型加载器批处理后的推理任务执行引擎算子调度、CUDA流管理、异步执行推理调度器推理结果服务接口HTTP/gRPC端点、序列化、限流执行引擎对外API这个划分方式的好处是关注点分离。张量抽象层不关心模型是什么模型加载器不关心请求怎么来调度器不关心算子怎么执行。每一层都可以独立测试和优化。我见过太多项目把所有这些逻辑揉在一个文件里最后改一行代码牵一发而动全身。注意模块划分不是越细越好。我一开始把算子调度也拆成独立模块结果发现调度器和执行引擎之间的接口调用开销比算子本身还大。后来合并了性能反而提升。拆分的粒度取决于你的性能瓶颈在哪。2.3 数据流与控制流的设计数据流的设计直接决定了系统的吞吐和延迟。这个项目采用生产者-消费者模型请求从服务接口进来被放入一个线程安全的队列调度器从队列里取请求攒成批次交给执行引擎。执行引擎异步执行结果通过回调或Future返回。控制流的关键在于背压机制。如果请求生产速度超过消费速度队列会无限增长最后OOM。所以必须设置队列上限满了就拒绝新请求或阻塞。我试过用有界队列加阻塞策略结果发现高并发下线程切换开销很大。后来改成无锁队列加自旋等待性能好了不少但CPU占用上去了。最终方案是混合策略低负载用阻塞高负载用自旋阈值根据队列长度动态调整。这个设计思路背后的逻辑是没有银弹只有权衡。阻塞队列实现简单但延迟高无锁队列性能好但实现复杂。你需要根据实际场景选择。如果你的QPS稳定在几百阻塞队列完全够用。如果QPS上万那就得考虑无锁方案。3. 核心模块的从零实现细节3.1 张量抽象层内存池与引用计数张量抽象层是整个系统的地基。框架里的Tensor看起来很简单但背后涉及内存分配、设备管理、生命周期跟踪。从零实现时我建议先做一个简单的内存池预分配一大块显存然后用一个空闲链表管理。每次需要显存时从链表取用完还回去。这样避免了频繁调用CUDA的cudaMalloc和cudaFree这两个操作非常慢。引用计数是另一个关键。Python有GC但GPU显存不归GC管。你需要自己跟踪每个张量被多少地方引用引用归零才能回收。实现方式很简单每个张量对象维护一个计数器创建时加一赋值时加一离开作用域时减一。减到零就还给内存池。class Tensor: def __init__(self, shape, dtype, devicecuda): self.shape shape self.dtype dtype self.device device self._ref_count 1 self._data memory_pool.allocate(shape, dtype, device) def retain(self): self._ref_count 1 return self def release(self): self._ref_count - 1 if self._ref_count 0: memory_pool.free(self._data)这段代码看起来简单但坑很多。比如循环引用怎么处理多线程并发修改引用计数怎么保证原子性我的经验是先实现单线程版本跑通了再加锁。一上来就搞无锁原子操作调试成本太高。实操心得内存池的块大小要设计好。如果块太小大张量分配不了如果块太大小张量浪费严重。我一般按2的幂次划分从1MB到1GB分十几档每档一个空闲链表。分配时找最接近的档位不够就往上取。3.2 模型加载器权重解析与计算图构建模型加载器的核心任务是把训练好的权重文件读进来初始化到张量里然后构建计算图。计算图可以用邻接表表示每个节点是一个算子边是张量依赖。从零实现时不需要支持所有算子先支持矩阵乘、卷积、激活函数这几个最常用的。权重解析的坑在于格式兼容性。PyTorch的.pt文件本质是pickle但直接torch.load会依赖PyTorch。如果你想完全从零可以用safetensors格式它就是一个简单的二进制文件头部是JSON描述后面是连续的张量数据。解析起来非常直接。import json import struct def load_safetensors(path): with open(path, rb) as f: header_len struct.unpack(Q, f.read(8))[0] header json.loads(f.read(header_len)) data f.read() tensors {} for name, info in header.items(): if name __metadata__: continue start, end info[data_offsets] shape info[shape] dtype info[dtype] tensors[name] parse_tensor(data[start:end], shape, dtype) return tensors计算图构建的关键是拓扑排序。你需要确定算子的执行顺序保证每个算子的输入都已经计算完成。从零实现时可以用Kahn算法先统计每个节点的入度入度为0的入队然后依次出队把下游节点入度减一减到0就入队。这样得到的顺序就是合法的执行顺序。3.3 推理调度器动态批处理与超时控制动态批处理是提升GPU利用率的核心手段。原理很简单把多个请求攒成一个批次一起送进GPU计算。但实现起来要考虑几个问题批次多大合适等多久攒批超时了怎么办我的实现方案是双阈值控制设置最大批次大小和最大等待时间。调度器从队列取请求如果当前批次没满且等待时间没超就继续等如果满了或超时了就立即执行。这样在低负载时延迟低高负载时吞吐高。class BatchScheduler: def __init__(self, max_batch_size32, max_wait_ms10): self.max_batch_size max_batch_size self.max_wait_ms max_wait_ms self.queue deque() def schedule(self): batch [] start_time time.time() while len(batch) self.max_batch_size: if self.queue: batch.append(self.queue.popleft()) elif time.time() - start_time self.max_wait_ms / 1000: break else: time.sleep(0.0001) return batch这个逻辑看起来简单但实际跑起来问题不少。比如time.sleep的精度不够Windows上最小15msLinux上1ms左右。高QPS下这个精度会导致批次大小波动很大。我的解决办法是用select或epoll做超时等待精度能到微秒级。注意动态批处理不是万能的。如果请求的输入长度差异很大攒批后需要paddingpadding越多浪费越大。这时候可以考虑按长度分桶同桶内的请求攒批。我试过按长度分10个桶吞吐提升了30%以上。3.4 执行引擎CUDA流与异步执行执行引擎负责把计算图变成实际的GPU操作。从零实现时你需要管理CUDA流。默认流是同步的一个操作做完才做下一个。用多个流可以实现异步执行比如数据拷贝和计算重叠。# 创建两个流 stream_compute cuda.Stream() stream_copy cuda.Stream() # 在拷贝流上做数据传输 with cuda.stream(stream_copy): input_data.copy_to_device() # 在计算流上做推理 with cuda.stream(stream_compute): output model(input_data) # 同步两个流 stream_compute.synchronize() stream_copy.synchronize()这个模式的关键是依赖管理。计算流必须等拷贝流完成才能开始否则会读到脏数据。CUDA提供了cudaStreamWaitEvent来做流间同步。从零实现时你可以用一个事件对象来标记拷贝完成然后让计算流等待这个事件。我踩过的一个坑是流太多反而慢。每个流都有调度开销流之间同步也有开销。我一开始给每个请求创建一个流结果GPU上下文切换频繁性能还不如单流。后来改成固定几个流轮流使用性能才稳定下来。一般4到8个流就够了具体看GPU型号。4. 完整实操流程与关键步骤4.1 环境准备与依赖安装这个项目对环境的依赖其实很少核心就是Python、NumPy、CUDA驱动。如果你想用GPU加速需要装CUDA Toolkit和PyCUDA或Numba。我建议用conda建一个干净的环境避免和系统Python冲突。conda create -n ai-scratch python3.10 conda activate ai-scratch pip install numpy safetensors # 如果需要GPU支持 conda install -c nvidia cuda-toolkit pip install pycuda版本选择上CUDA 11.8和12.x都可以但要注意驱动版本兼容。我遇到过驱动太老导致CUDA初始化失败的情况报错信息很隐晦排查了半天。建议先用nvidia-smi确认驱动版本再对照CUDA文档选版本。提示如果你没有GPU纯CPU也能跑通整个流程只是速度慢。NumPy的矩阵乘在CPU上也有优化小模型推理完全够用。先跑通逻辑再考虑加速。4.2 从零实现一个简单的推理服务我以一个简单的全连接网络为例展示从模型加载到服务化的完整流程。假设你已经训练好一个模型保存为safetensors格式。第一步加载权重并构建计算图from tensor import Tensor from model_loader import load_safetensors, build_graph weights load_safetensors(model.safetensors) graph build_graph(weights, input_shape(1, 784))第二步初始化调度器和执行引擎from scheduler import BatchScheduler from engine import ExecutionEngine scheduler BatchScheduler(max_batch_size16, max_wait_ms5) engine ExecutionEngine(graph, devicecuda)第三步启动服务循环from http_server import HTTPServer def handle_request(request): input_tensor preprocess(request.data) scheduler.enqueue(input_tensor) result engine.wait_result(input_tensor.id) return postprocess(result) server HTTPServer(port8080, handlerhandle_request) server.start()这个流程看起来简单但每一步都有细节。比如preprocess要做归一化、resize、通道转换这些操作如果在Python里做会成为瓶颈。我的做法是用NumPy向量化或者用CUDA kernel在GPU上做。实测下来预处理用GPU做比CPU快5到10倍。4.3 性能测试与调优记录跑通之后我用wrk和ab做了压测。初始版本QPS只有50左右延迟P99在200ms。分析后发现几个瓶颈瓶颈点现象优化方案优化后QPS预处理在CPUCPU占用100%GPU空闲预处理移到GPU120批处理等待过长批次大小平均只有4降低等待时间到2ms180内存分配频繁每次推理都malloc启用内存池220单流执行GPU利用率60%改用4个CUDA流280这个调优过程让我深刻理解了一件事AI工程的性能瓶颈往往不在模型本身而在数据流动和资源调度。模型计算只占整个链路的一小部分大部分时间花在数据搬运、内存分配、线程同步上。实操心得压测时一定要用真实数据分布。我用随机数据压测时QPS很高换成真实数据后掉了30%因为真实数据的长度分布不均匀padding浪费严重。后来按长度分桶才解决。5. 常见问题与排查技巧实录5.1 显存泄漏与OOM排查显存泄漏是AI工程里最头疼的问题之一。现象是服务跑一段时间后OOM但重启就好。排查思路是记录每次分配和释放的调用栈找出没有配对的分配。我的做法是在内存池里加一个调试模式每次分配时记录调用栈释放时删除记录。服务跑一段时间后打印剩余的分配记录就能看到哪些张量没被释放。import traceback class MemoryPool: def __init__(self, debugFalse): self.debug debug self.allocations {} def allocate(self, shape, dtype, device): ptr self._do_allocate(shape, dtype, device) if self.debug: self.allocations[ptr] traceback.format_stack() return ptr def free(self, ptr): if self.debug and ptr in self.allocations: del self.allocations[ptr] self._do_free(ptr)常见泄漏原因有几个一是异常路径没有释放比如推理中途报错张量没回收二是循环引用两个张量互相引用引用计数永远不归零三是缓存没清理比如KV cache只增不减。我的经验是所有分配都要用try-finally包起来确保异常时也能释放。5.2 推理结果不一致的排查有时候你会发现同样的输入两次推理结果不一样。如果模型是确定性的那问题一定在工程实现上。常见原因有浮点累加顺序不同多线程或不同批次大小会导致浮点运算顺序变化结果有微小差异。这是正常的但如果差异很大说明有bug。内存未初始化分配的内存没清零残留了上次的数据。解决办法是分配时强制清零或者用cudaMemset。竞态条件多个线程同时读写同一块内存。用cuda-memcheck或compute-sanitizer可以检测。我遇到过一次结果不一致排查了一整天最后发现是批处理时padding的位置不对。不同批次的padding值不同导致注意力机制算出了不同的结果。解决办法是padding用固定的mask值并且确保mask在计算时生效。5.3 高并发下的服务稳定性问题高并发下最常见的问题是雪崩请求堆积导致队列满队列满导致拒绝服务拒绝服务导致客户端重试重试加剧堆积。打破这个循环的关键是限流和降级。我的方案是三层防护第一层在入口做限流超过阈值的请求直接返回429第二层在调度器做超时控制等待超过阈值的请求直接丢弃第三层在执行引擎做熔断如果连续失败超过阈值暂停服务一段时间。class CircuitBreaker: def __init__(self, failure_threshold10, recovery_timeout30): self.failure_count 0 self.failure_threshold failure_threshold self.recovery_timeout recovery_timeout self.last_failure_time 0 self.state CLOSED def call(self, func, *args, **kwargs): if self.state OPEN: if time.time() - self.last_failure_time self.recovery_timeout: self.state HALF_OPEN else: raise Exception(Circuit breaker is open) try: result func(*args, **kwargs) if self.state HALF_OPEN: self.state CLOSED self.failure_count 0 return result except Exception as e: self.failure_count 1 self.last_failure_time time.time() if self.failure_count self.failure_threshold: self.state OPEN raise e这个熔断器实现很简单但效果很好。我实测下来加了熔断之后服务在过载时不会完全不可用而是快速失败给系统恢复的时间。5.4 常见问题速查表问题现象可能原因排查方法解决方案服务启动报CUDA错误驱动版本不匹配nvidia-smi查看驱动版本升级驱动或降级CUDA推理速度突然变慢GPU降频或显存不足nvidia-smi -q -d PERFORMANCE检查散热和显存占用批次大小上不去请求到达率低或等待时间短打印批次大小分布调整等待时间或增加并发结果随机波动浮点累加顺序或竞态固定随机种子单线程测试统一计算顺序加锁内存持续增长泄漏或缓存未清理内存池调试模式修复泄漏设置缓存上限高并发下超时队列满或线程池耗尽监控队列长度和线程数限流、扩容、异步化这张表是我过去两年遇到问题的总结基本上覆盖了80%的常见故障。剩下的20%需要具体问题具体分析但排查思路是一样的先定位瓶颈在哪一层再深入那一层的细节。6. 从零实现后的经验总结与扩展方向做完这个项目我最大的感受是AI工程的门槛不在算法而在系统。你会调model.generate()不代表你会部署模型你会写训练脚本不代表你会做推理服务。从零实现一遍那些曾经模糊的概念——批处理、内存池、CUDA流、背压——都会变得具体。这个项目后续可以往几个方向扩展。一是支持更多算子比如注意力机制、LayerNorm、卷积让它能跑Transformer类模型。二是加入量化支持把FP16或INT8的推理实现一遍理解量化的精度损失和加速比。三是做分布式推理把模型切分到多张GPU上实现张量并行或流水线并行。每一个方向都能让你对AI系统的理解更深一层。我个人在实际操作中的体会是不要一开始就追求完美。先写一个能跑通的版本哪怕性能很差然后再逐步优化。我第一版推理服务QPS只有10但正是这个简陋的版本让我看清了所有瓶颈在哪。如果一上来就照着Triton的架构抄你永远不知道那些设计决策背后的原因。从零开始慢慢迭代这才是ai-engineering-from-scratch的真正意义。