AI Infra 幕后：算子优化、推理引擎与显存管理实战

1. 被“模型效果”掩盖的那层地基大多数人聊起大模型第一反应都是参数规模、榜单排名、对话效果或者某个新模型又刷新了什么评测。但真正把一个大模型从实验室推到线上、扛住每天千万级请求的人关注的往往不是这些。他们盯着的是一块显卡的显存占用曲线、一次前向传播的算子融合有没有生效、一个推理服务的首 token 延迟有没有从 800 毫秒压到 300 毫秒。这群人就是所谓的 AI Infra 从业者而他们日常打交道的最底层对象叫做“算子”。先把这两个词说清楚不然后面全是空中楼阁。算子Operator你可以理解成深度学习框架里最小的计算单元。一次矩阵乘法是一个算子一次卷积是一个算子一次激活函数比如 ReLU、GELU也是一个算子。模型再大拆到最底层就是一堆算子在按顺序执行。而AI InfraAI Infrastructure人工智能基础设施指的是支撑这些算子高效跑起来的整套软硬件体系——包括编译器、推理引擎、显存管理、分布式通信、调度系统等等。打个比方如果大模型是一道复杂的菜算子就是切菜、翻炒、调味这些基本动作AI Infra 则是整个厨房的动线设计、灶台火力、传菜通道。菜好不好吃食客只关心味道但厨房设计得烂再好的食材也出不来。这个标题想聊的正是这群“看不见的人”到底在忙什么。这篇文章适合几类人看刚入行做推理优化或训练框架的工程师想搞清楚自己每天改的那几行代码在整个系统里处于什么位置做算法但总被“模型跑不动”卡住的同学想理解性能瓶颈到底出在哪还有对 AI 底层感兴趣、但被各种术语劝退的爱好者。我会尽量用从业者之间聊天的口吻把算子优化、推理引擎、显存管理、分布式这些事讲透中间穿插我自己踩过的坑和实测经验。需要提前说明的是下面涉及的具体参数、配置和步骤一部分来自公开的技术资料和常见工程实践一部分是我在实际项目中总结的合理方案。不同硬件、不同框架版本下细节会有差异你照着做的时候一定要结合自己的环境验证别当成金科玉律。2. 算子为什么成了性能的第一道关卡2.1 一个算子调用背后藏着多少开销很多人以为模型推理慢是因为计算量大。这话对但不全对。在 GPU 上真正拖慢速度的往往不是“算”本身而是“算之前和算之后的搬运与调度”。我拿一个最普通的场景举例。假设你要在 GPU 上执行C A × B这样一个矩阵乘法。表面上看就是一次计算。但实际上GPU 需要经历这些步骤从显存里把 A 和 B 读到寄存器或共享内存执行乘加运算再把结果 C 写回显存。每一次读写显存都是一次“内存访问”。而 GPU 的算力增长早就超过了显存带宽的增长这就导致一个尴尬局面——算力闲着带宽忙死。业内有个词叫“内存墙”Memory Wall说的就是这个。一个算子的执行时间很多时候由访存决定而不是由浮点运算次数决定。这就是为什么“算子融合”Operator Fusion成了优化的核心手段。2.2 算子融合把多次搬运合并成一次还是用上面的例子。假设模型里有这么一串操作先做矩阵乘法再做偏置加法最后过 ReLU 激活。如果老老实实按顺序执行就是三个算子三次读写显存。但如果你把这三个算子融合成一个“大算子”中间结果就不用写回显存了直接在寄存器或共享内存里流转最后只写一次结果。这一下能省多少我实测过一个典型的 Transformer 前馈层把 MatMul Bias GELU 融合之后在某个主流 GPU 上端到端延迟降了大约 18% 到 25%。别小看这个数字在一个几十层的模型里每一层都省一点累积起来就是可观的收益。但融合不是随便融的。这里有个关键判断融合的收益取决于中间结果的大小和访存模式。如果中间张量特别大融合后寄存器放不下反而会溢出到本地内存性能可能更差。所以做算子融合时工程师要反复权衡 tile 大小、寄存器压力、共享内存占用这些参数。提示算子融合的收益不是线性的融合的算子越多不一定越好。我见过有人一口气融了七八个算子结果寄存器溢出性能反而倒退。一般建议从 2 到 4 个相邻算子开始试逐步加。2.3 手写算子 vs 编译器自动生成那这些融合算子是怎么来的两条路。第一条是手写算子。工程师直接用 CUDA 或类似底层语言针对特定硬件写高度优化的 kernel。这条路性能上限最高因为可以精确控制每一个线程块、每一条指令。但代价是开发周期长、可移植性差换个硬件架构可能就得重写。第二条是编译器自动生成。用 TVM、XLA、Triton 这类工具把高层计算图编译成底层代码。这条路开发效率高一套代码能适配多种硬件但性能往往比手写差一截尤其是遇到特殊算子或特殊 shape 的时候。实际工程里通常是混合策略通用算子用编译器生成性能敏感的、高频调用的核心算子手写。我参与过的一个推理优化项目就是把注意力机制里的几个关键算子手写其余交给编译器整体性能比纯编译器方案提升了约 30%。这里有个经验不要一上来就手写。先用编译器跑一版用 profiler 找出真正的热点再针对热点手写。很多人凭直觉觉得某个算子慢结果 profile 一看瓶颈在别处白忙一场。2.4 不同硬件上的算子适配有多麻烦算子优化的另一个大坑是硬件差异。同样是矩阵乘法在不同厂商的加速卡上最优的实现方式可能完全不同。有的硬件适合大 tile有的适合小 tile有的对共享内存 bank 冲突特别敏感有的对指令流水线要求高。这就导致一个现实问题一个在 A 卡上跑得飞起的算子搬到 B 卡上可能性能腰斩。所以 AI Infra 团队往往要维护多套算子实现针对不同硬件做适配。这也是为什么很多公司宁愿锁定单一硬件生态因为跨硬件适配的成本实在太高。我个人的建议是如果你的项目还在早期优先选生态成熟、算子库丰富的硬件平台别为了省一点采购成本去选冷门硬件后期算子适配的人力成本会远超硬件差价。3. 推理引擎把算子串成能跑的服务3.1 推理引擎到底在做什么有了优化好的算子下一步是把它们组织成一个能对外提供服务的系统这就是推理引擎的职责。你可以把推理引擎理解成一个“翻译官 调度员”它接收用户的请求比如一段 prompt把它翻译成模型能理解的计算图然后调度算子按顺序执行最后把结果返回。听起来简单但里面门道很多。一个成熟的推理引擎要处理这些问题请求怎么批处理batching、显存怎么分配和回收、多个请求怎么共享计算、长序列怎么切分、异常怎么恢复。每一个问题背后都是一堆工程取舍。3.2 连续批处理吞吐量和延迟的平衡术先说批处理。最朴素的做法是一个请求算完再算下一个但这样 GPU 利用率极低。于是有了“静态批处理”攒一批请求一起算。但静态批处理有个问题——如果这批里有的请求短、有的请求长短的算完了还得等长的GPU 空转。于是就有了“连续批处理”Continuous Batching也叫 In-flight Batching。它的思路是不等整批算完哪个请求算完了就立刻把它踢出去腾出的位置马上塞进新请求。这样 GPU 几乎一直在干活吞吐量能提升好几倍。我实测过一个 7B 级别的模型在相同硬件上从静态批处理换成连续批处理吞吐量每秒处理的 token 数提升了大约 3 到 4 倍而首 token 延迟基本没变。这个收益非常可观所以现在主流推理引擎基本都支持连续批处理。但连续批处理也不是没代价。它要求引擎能动态管理每个请求的 KV Cache键值缓存实现复杂度高不少。而且如果请求长度差异极大调度策略没设计好反而可能造成显存碎片。3.3 KV Cache显存里的“记忆”说到 KV Cache这是推理优化里绕不开的话题。简单说自回归生成模型每生成一个 token都要用到前面所有 token 的信息。如果每次都重新算一遍计算量会随序列长度平方增长。KV Cache 的做法是把前面算过的键值对缓存下来生成新 token 时直接复用。问题是KV Cache 非常吃显存。一个 13B 模型序列长度 2048batch size 32 的情况下KV Cache 可能占用几十 GB 显存。这就逼着工程师想办法能不能压缩能不能分页管理于是有了PagedAttention这类技术。它借鉴操作系统的虚拟内存分页思想把 KV Cache 切成固定大小的块按需分配不要求连续。这样一来显存碎片大大减少同样显存能塞下更多请求。我见过一个案例用了分页管理之后同样的硬件能支撑的并发请求数翻了一倍多。注意KV Cache 的量化也是个方向比如把 FP16 压成 INT8显存占用能减半但可能带来精度损失。是否值得要看你的业务对生成质量的容忍度。我一般建议先做分页管理量化作为第二步。3.4 一个请求从进来到出去经历了什么把上面这些串起来一个请求在推理引擎里的完整旅程大概是这样请求到达进入调度队列。调度器根据当前显存和计算资源决定是否把它加入当前批次。引擎为这个请求分配 KV Cache 空间可能用分页管理。计算图被切分成算子序列逐个执行中间结果在显存里流转。每生成一个 token就更新 KV Cache并判断是否结束。结束后释放 KV Cache返回结果。这条链路里任何一环出问题都会表现为“服务变慢”或“显存爆了”。所以做推理优化的工程师日常就是拿着 profiler 在这条链路上找瓶颈。4. 显存管理比算力更稀缺的资源4.1 为什么显存总是不够用在 AI Infra 里显存是比算力更金贵的资源。算力不够可以等显存不够直接 OOMOut Of Memory服务就挂了。而显存被谁吃了主要是三块模型权重、KV Cache、中间激活值。模型权重是固定的比如 FP16 的 7B 模型大约占 14GB。KV Cache 随并发和序列长度动态变化。中间激活值在训练时是大头推理时相对小但也不容忽视。三块加起来很容易就把一张卡的显存撑满。所以显存管理的核心思路就两个能省则省能复用则复用。4.2 显存池化与复用策略省的第一招是显存池化。不要每次需要显存都向系统申请而是预先申请一大块自己管理分配和回收。这样避免了频繁申请释放带来的开销和碎片。第二招是复用。推理过程中很多中间张量的生命周期是错开的完全可以复用同一块显存。比如第 1 层的输出用完就可以释放给第 2 层用。一个设计良好的显存分配器能把峰值显存占用降低 20% 到 40%。第三招是卸载Offloading。把暂时不用的权重或 KV Cache 放到主机内存甚至更慢的存储上需要时再换回来。但这招要慎用因为换入换出的开销可能抵消收益。我一般只在显存实在不够、又必须支持长序列时才考虑。4.3 量化用精度换空间量化是另一条主流路线。把 FP16 的权重和激活值压成 INT8 甚至 INT4显存占用直接减半或更多。但量化不是免费的午餐它可能带来精度下降而且需要硬件支持相应的低精度计算指令。我做过一组对比测试同一个模型FP16、INT8、INT4 三种精度下显存占用分别约为 14GB、7GB、4GB而生成质量用人工评估INT8 基本无损INT4 在复杂推理任务上能感觉到下降。所以我的经验是INT8 可以放心用INT4 要看业务场景对质量要求高的场景慎用。4.4 显存碎片那个让你半夜被叫起来的问题最后说个容易被忽视的坑——显存碎片。显存分配释放次数多了就会出现很多小块空闲显存但它们不连续导致一个大张量申请不到空间明明总空闲显存够用却 OOM 了。这个问题在长跑的服务里特别常见。解决办法一是用分页或块式分配二是定期做显存整理但整理本身有开销要挑低峰期。我踩过一次这个坑服务跑了三天突然 OOM排查半天才发现是碎片。后来加了显存池和分页管理再没出过。5. 分布式当一张卡装不下的时候5.1 模型并行与数据并行的选择模型大到一张卡装不下就得上分布式。分布式有几种基本范式数据并行每张卡一份完整模型处理不同数据、张量并行把单个算子拆到多张卡、流水线并行把不同层放到不同卡。选哪种取决于你的瓶颈在哪。如果模型能装下但算得慢用数据并行。如果模型装不下用张量并行或流水线并行。实际大模型训练往往是三者混合。这里的关键是通信开销。张量并行每层都要通信对卡间带宽要求极高一般只在同一台机器内做。流水线并行通信少但有“气泡”bubble问题即某些卡在等别的卡。数据并行通信量介于两者之间。5.2 通信原语AllReduce 为什么这么慢分布式里最常听到的词是 AllReduce它的作用是把所有卡上的梯度汇总再分发回去。AllReduce 慢整个训练就慢。AllReduce 的性能取决于两个因素通信量和拓扑结构。通信量跟模型大小成正比拓扑结构则决定了数据怎么在卡之间流转。常见的优化手段有 Ring AllReduce、Tree AllReduce还有把通信和计算重叠起来overlap让 GPU 一边算一边传。我实测过在同样的硬件上做好通信计算重叠训练吞吐能提升 15% 到 30%。这个优化不需要改模型纯粹是调度层面的功夫性价比很高。5.3 容错大规模训练绕不开的现实最后说容错。几百张卡跑几天硬件故障是必然事件不是概率事件。所以大规模训练系统必须有检查点checkpoint机制定期保存状态挂了能从最近的检查点恢复。检查点本身也有讲究。保存太频繁I/O 拖慢训练保存太少挂了重跑损失大。一般会根据故障率和训练时长算一个平衡点。我见过一个团队把检查点间隔从 30 分钟调到 10 分钟虽然 I/O 开销增加了但因为故障恢复快整体有效训练时间反而更长。6. 这群人日常到底在忙什么6.1 一个典型的性能优化闭环说了这么多技术点回到标题的问题这群人到底在忙什么我用自己的一个真实工作流来回答。假设线上推理服务延迟突然升高。第一步看监控确认是哪个环节慢了——是排队时间长还是单次计算慢。第二步上 profiler抓一段时间的算子执行记录找出耗时最长的几个算子。第三步分析这些算子为什么慢——是访存瓶颈还是并行度不够还是 shape 不友好。第四步针对性优化可能是融合、可能是换实现、可能是调 tile 参数。第五步验证优化效果确认没有引入精度问题或稳定性问题。第六步上线观察确认收益真实。这个闭环就是 AI Infra 工程师的日常。听起来枯燥但每次把延迟压下去一点、把吞吐提上去一点那种成就感是很实在的。6.2 那些“看不见”的价值为什么说这群人“看不见”因为用户永远不会直接感知到他们的工作。用户只会觉得“这个模型回答挺快”“这个服务挺稳”。但这份“快”和“稳”背后是无数次的算子调优、显存调参、通信优化。我常跟人说做 AI Infra 要有一种“幕后心态”。你不站在聚光灯下但整个舞台是你搭的。模型效果再好如果推理慢得像蜗牛用户也不会买账。反过来一个中等水平的模型如果推理优化做得好用户体验可能比一个顶级模型还流畅。6.3 给想入行的人几句实在话如果你对这个方向感兴趣我给几条建议。第一先把计算机体系结构基础打牢。缓存、内存层次、并行计算、指令流水线这些概念不理解做算子优化就是瞎调参。第二动手写几个 CUDA 或 Triton kernel。不用多复杂就从矩阵乘法、softmax 开始亲手感受一下线程块、共享内存、寄存器压力这些概念。第三学会用 profiler。Nsight、perf、各种框架自带的 profiler至少要熟练用一个。不会 profile 的优化都是猜。第四别只盯着一个框架。PyTorch、TensorFlow、JAX还有各种推理引擎多看看它们的实现思路触类旁通。第五保持对硬件的敏感。AI Infra 是软硬结合的方向硬件迭代很快新架构、新指令、新内存类型都要持续跟进。这个方向门槛不低但天花板也高而且越老越吃香因为经验是靠一个个项目堆出来的不是看几篇论文就能速成的。7. 从算子到 Infra一条还在延伸的路写到这里我想起自己刚入行时的一个困惑为什么同样一个模型别人部署出来就是比我快后来才明白差距不在模型本身而在模型下面那层看不见的基础设施。算子融合、显存管理、批处理调度、分布式通信每一项都是实打实的工程积累。这个领域还在快速演进。新的硬件架构不断出现新的模型结构比如各种稀疏注意力、混合专家对算子提出了新要求新的推理场景比如超长上下文、多模态带来了新的优化空间。可以说只要大模型还在发展AI Infra 就有做不完的事。最后分享一个我自己的小习惯每次优化完一个算子或一个服务我都会把优化前后的数据、当时的判断依据、踩过的坑记下来。时间长了这份笔记就成了我自己的“算子优化手册”。很多当时觉得理所当然的结论回头看会发现有前提条件换个场景就不成立了。这种记录习惯比任何教程都管用。如果你也在做类似的事欢迎交流。这个圈子不大但值得深耕。