从零手搓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 核心模块划分与依赖关系

整个系统我把它拆成五个核心模块，每个模块的职责和依赖关系如下：

这个划分方式的好处是关注点分离。张量抽象层不关心模型是什么，模型加载器不关心请求怎么来，调度器不关心算子怎么执行。每一层都可以独立测试和优化。我见过太多项目把所有这些逻辑揉在一个文件里，最后改一行代码牵一发而动全身。

注意：模块划分不是越细越好。我一开始把算子调度也拆成独立模块，结果发现调度器和执行引擎之间的接口调用开销比算子本身还大。后来合并了，性能反而提升。拆分的粒度取决于你的性能瓶颈在哪。

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, device='cuda'): 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_size=32, max_wait_ms=10): 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上最小15ms，Linux上1ms左右。高QPS下这个精度会导致批次大小波动很大。我的解决办法是用select或epoll做超时等待，精度能到微秒级。

注意：动态批处理不是万能的。如果请求的输入长度差异很大，攒批后需要padding，padding越多浪费越大。这时候可以考虑按长度分桶，同桶内的请求攒批。我试过按长度分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 python=3.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_size=16, max_wait_ms=5) engine = ExecutionEngine(graph, device='cuda')

第三步，启动服务循环：

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(port=8080, handler=handle_request) server.start()

这个流程看起来简单，但每一步都有细节。比如preprocess要做归一化、resize、通道转换，这些操作如果在Python里做，会成为瓶颈。我的做法是用NumPy向量化，或者用CUDA kernel在GPU上做。实测下来，预处理用GPU做比CPU快5到10倍。

4.3 性能测试与调优记录

跑通之后，我用wrk和ab做了压测。初始版本QPS只有50左右，延迟P99在200ms。分析后发现几个瓶颈：

这个调优过程让我深刻理解了一件事：AI工程的性能瓶颈往往不在模型本身，而在数据流动和资源调度。模型计算只占整个链路的一小部分，大部分时间花在数据搬运、内存分配、线程同步上。

实操心得：压测时一定要用真实数据分布。我用随机数据压测时QPS很高，换成真实数据后掉了30%，因为真实数据的长度分布不均匀，padding浪费严重。后来按长度分桶才解决。

5. 常见问题与排查技巧实录

5.1 显存泄漏与OOM排查

显存泄漏是AI工程里最头疼的问题之一。现象是服务跑一段时间后OOM，但重启就好。排查思路是：记录每次分配和释放的调用栈，找出没有配对的分配。

我的做法是在内存池里加一个调试模式，每次分配时记录调用栈，释放时删除记录。服务跑一段时间后，打印剩余的分配记录，就能看到哪些张量没被释放。

import traceback class MemoryPool: def __init__(self, debug=False): 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_threshold=10, recovery_timeout=30): 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 常见问题速查表

这张表是我过去两年遇到问题的总结，基本上覆盖了80%的常见故障。剩下的20%需要具体问题具体分析，但排查思路是一样的：先定位瓶颈在哪一层，再深入那一层的细节。

6. 从零实现后的经验总结与扩展方向

做完这个项目，我最大的感受是：AI工程的门槛不在算法，而在系统。你会调model.generate()不代表你会部署模型，你会写训练脚本不代表你会做推理服务。从零实现一遍，那些曾经模糊的概念——批处理、内存池、CUDA流、背压——都会变得具体。

这个项目后续可以往几个方向扩展。一是支持更多算子，比如注意力机制、LayerNorm、卷积，让它能跑Transformer类模型。二是加入量化支持，把FP16或INT8的推理实现一遍，理解量化的精度损失和加速比。三是做分布式推理，把模型切分到多张GPU上，实现张量并行或流水线并行。每一个方向都能让你对AI系统的理解更深一层。

我个人在实际操作中的体会是：不要一开始就追求完美。先写一个能跑通的版本，哪怕性能很差，然后再逐步优化。我第一版推理服务QPS只有10，但正是这个简陋的版本让我看清了所有瓶颈在哪。如果一上来就照着Triton的架构抄，你永远不知道那些设计决策背后的原因。从零开始，慢慢迭代，这才是ai-engineering-from-scratch的真正意义。