Toonflow 这个项目做了半年多一直有一条主线贯穿始终生成质量往上提显存却撑不住批处理堆上去延迟又开始飘。中间试过不少常规手段像减小 batch、砍序列长度、换更小的模型都属于花钱买平安训练指标是保住了但 GPU 的算力利用率一直上不去资源在那里空转。进入进阶阶段后我把重心从调参转向推理链路和显存生命周期本身从算子层面、CUDA Graph、KV Cache、激活显存这几个方向逐个拆才算真正把性能和显存管理纳入同一个优化框架里。这篇博文我会把在 Toonflow 上的完整优化实践展开来讲不绕弯子直接进入技术细节包括怎么用性能分析工具定位瓶颈、为什么最终放弃 TensorRT 而选择算子融合加 CUDA Graph 的组合、显存管理里有哪几个隐藏的大坑以及每一步优化前后实测出来的数据变化。内容会包含可以抄作业的命令、配置和代码片段也有踩过坑之后的判断逻辑。1. 瓶颈扫描延迟和显存到底消耗在哪里很多优化项目一开始容易凭感觉下手觉得模型大就换小模型延迟高就降 batch这种做法往往治标不治本。Toonflow 早期的困境很典型单卡单 batch 推理一次要 4.2 秒显存峰值稳定在 17GB 左右按理说资源占用也不算特别夸张但 P99 延迟能到 6.9 秒波动非常明显。这时候最需要做的不是改代码而是先搞清楚延迟和显存到底烧在了什么地方。1.1 先跑通 Profiling 工具链再做改动我建议所有优化工作开始前先花时间把 profiling 工具链跑通。在 Toonflow 这个项目里我主要用了三样东西PyTorch Profiler 看算子级耗时和显存分配、NVIDIA Nsight Systems 看 kernel 启动和 Host 侧的同步等待、以及 roofline 模型分析理论计算量和实际吞吐的差距。# 基于 PyTorch Profiler 的轻量级 profiling 入口 from torch.profiler import profile, ProfilerActivity, record_function with profile( activities[ProfilerActivity.CPU, ProfilerActivity.CUDA], record_shapesTrue, profile_memoryTrue, with_stackTrue, ) as prof: with record_function(toonflow_generation): output model.generate(input_ids, max_new_tokens512) print(prof.key_averages().table( sort_bycuda_time_total, row_limit30 ))这一步跑出来的结果非常关键。Top 20 耗时算子里面排序靠前的竟然不是卷积或者注意力计算而是一堆 reshape、transpose、cat 之类的内存搬运操作加起来占了总 GPU 时间的 37%。这非常不合理明显是指计算图设计得碎片化小算子太多kernel launch 开销被无限放大。另一个让我意外来自 Nsight Systems 的数据是整个推理过程里GPU 的空闲时间接近 28%。换句话说GPU 很多时候在等着 CPU 下发指令而不是真的在计算。Host 和 Device 之间的同步等待严重拉低了实际吞吐。1.2 瓶颈定性判断的三个关键指标看完一轮 profiling 数据之后别急着动手改先定一下瓶颈属于哪一类。我习惯通过三个指标来判断第一个是 GPU 利用率曲线。如果平均利用率低于 40%大概率是 launch-bound 或者 memory-bound而不是 compute-bound。第二个是 kernel 单次执行时长。如果大量 kernel 运行时间在几十微秒以内说明计算密度太低图融合和 kernel 合并会非常有效。第三个是 Host 侧 stall 时间。利用 Nsight Systems 看 CPU 和 GPU 之间的 time gap如果 gap 占比超过 15%就要考虑 CUDA Graph 或者减少 python 侧调度。Toonflow 当时的数据是这样的GPU 平均利用率只有 46%kernel 数量达到 3400 多个其中超过 70% 的 kernel 执行时间小于 30 微秒Host-side stall 占比 28%。结论已经非常清晰了这是一个典型的 launch-bound 加 memory-bound 混合问题点对点的算子优化空间有限最好是做全局层面的计算图和显存规划。2. GPU 推理优化的三板斧算子融合、CUDA Graph 和自适应批处理定位到 launch-bound 和 memory-bound 之后我把优化工作分成了三个互相配合的层面算子级融合减少显存读写和 kernel 数量CUDA Graph 消除 Host 调度瓶颈再加上自适应批处理策略解决吞吐和延迟之间的矛盾。这三个层面不是孤立生效的而是叠加在一起的。2.1 为什么 TensorRT 不适合这个项目而手写算子融合可以很多做 GPU 优化的朋友会第一时间想到 TensorRT把整个模型转成 engine效果通常也不差。但 Toonflow 这个项目的特殊性在于模型结构迭代非常快基本每周都有新模块加进来比如新的注意力变体或者自定义的 normalization 方法。每次结构改动都要重新导出 ONNX、过一遍解析器、处理各种 op fallback维护成本实在太高。实测下来TensorRT 能带来的加速大概在 15% 左右但踩坑时间远远超过预期碰到不支持的算子还得退回原生 PyTorch 或者用 plugin 补齐投入产出比太低。相比之下手写关键路径上的算子融合只针对耗时最集中的几个模式下手收益更容易控制。Toonflow 里我重点融合了两个模式第一个是 QKV 投影融合。原始实现里存在三个独立的 LinearQ、K、V 分开计算每层都有三次大矩阵乘法、三次显存分配、三次 kernel launch。融合方案是把这三个 Linear 的权重横向拼接成一个大的 Linear一次计算就能输出所有 QKV 结果。这个改动带来大约 11% 的端到端加速代码改动只有几行。# 融合前Q/K/V 分别做线性投影 query self.q_proj(hidden_states) key self.k_proj(hidden_states) value self.v_proj(hidden_states) # 融合后合并为一个 Linear权重横向拼接 qkv self.qkv_proj(hidden_states) query, key, value qkv.chunk(3, dim-1)第二个是 LayerNorm 激活函数 残差连接这三者融合。尽量避免在 HBM 里面反复读写中间张量比如先把 hidden_states 加上残差再做 LayerNorm最后紧跟 GELU用一个自定义 CUDA kernel 一口气算完而不是每个环节生成一个临时中间张量写回显存。注意别一上来就全网络做算子融合收益最大、风险最低的做法是先用 profiler 找到前 10 个耗时算子只对这些热点模式动手。融合一个非热点算子花大力气优化对整体延迟影响几乎为零。2.2 CUDA Graph 的接入和手写 catch 的细节Host-side stall 问题用 CUDA Graph 解决是最直接的。整个模型的 forward 过程先执行一次捕获之后每次推理只是 replay 一遍 graph省掉了 python 侧大量 kernel 启动和调度开销。但 CUDA Graph 不是白白就能用的有几个非常容易踩坑的地方必须处理好。第一个是静态输入输出 buffer。graph 捕获时的输入输出 tensor 地址会被固化所以每次推理要用相同 shape 的输入在同一个预分配的 buffer 上写数据输出也从同一个 buffer 读。第二个是动态 shape 问题如果你的模型内部存在比如长度依赖的分支、python 层的 if 条件、随输入长度变化的循环graph 捕获会直接失败。Toonflow 里我做了两步很关键的处理把生成阶段的序列长度固定为 max_new_tokens配合一个提前终止机制空余位置的 logits 直接 mask 掉对于 batch 变化设计了一个按 bucket 缓存 graph的策略只预创建几种常见 batch 配置1、4、8、16对应的 graph来一个新 batch 尺寸就查找最近的 bucket 复用。# CUDA Graph 复用示例按 batch bucket 缓存 graph graph_bucket {} def get_cuda_graph(batch_size, model, sample_input_shape): key batch_size # 简单策略按 batch 维度分组 if key not in graph_bucket: # 预热 捕获阶段 s torch.cuda.Stream() s.wait_stream(torch.cuda.current_stream()) with torch.cuda.stream(s): for _ in range(3): model(*sample_input_shape) torch.cuda.current_stream().wait_stream(s) g torch.cuda.CUDAGraph() with torch.cuda.graph(g): static_output model(*sample_input_shape) graph_bucket[key] (g, static_output) return graph_bucket[key]需要特别提醒的是capture 之前一定要做一次 warmup把可能触发的 lazy initialization、cuDNN autotuning、显存分配都提前跑一遍。否则 capture 过程中如果碰到显存申请或者算法选择整个 graph 捕获会失败。实测下来接入 CUDA Graph 之后P99 延迟从 6.9 秒降到了 3.2 秒Host-side stall 从 28% 降到了 9% 左右这一项拿到了整轮优化里最大的一块收益。2.3 动态批处理与抢占式调度的实践思路Toonflow 线上推理服务的请求模式波动比较大高峰期并发请求多闲时几乎没有流量。固定 batch 的做法会出现在高峰期排队过长在闲时又浪费算力。我引入了一套轻量级的动态批处理模块请求先进入队列调度器每 50ms 采集一次队列里的请求数量如果超过预设阈值就触发一次批量推理并把 batch 内请求的最大生成长度作为本次推理的序列长度。这里有一个延迟和吞吐的平衡问题。等待更多请求凑成一个大 batch能提升吞吐但每个请求在队列里的等待时间会变长。我的策略是设置一个等待窗口上限距离第一个请求进入队列超过 200ms无论 batch 是否凑满都立刻开始推理。这个策略在 Toonflow 里把整体吞吐提升了约 1.9 倍同时 P95 延迟依旧控制在 2.8 秒以内。另外可以留意调度器的优先级问题。我针对实时性要求较高的请求设置了一个独立的高优先级队列但高优先级队列不能无限抢占资源否则批处理策略会被打散。我的做法是当高优先级队列存在且等待时间超过 100ms最多让当前 batch 拆掉一半资源去服务这些请求其余仍然保持批处理。3. 显存管理从被动爆显存到主动规划每一块空间显存管理这个主题表面上看起来只是显存不够就减小 batch或者换成更大显存的卡但实际上大模型推理的显存消耗远不是一张模型权重表那么简单。Toonflow 在优化过程中遇到过好几次看起来显存占用不高却莫名其妙 OOM的场景最后排查下来问题都出在显存碎片化和临时 buffer 反复申请上。3.1 推理阶段的显存占用拆解推理阶段显存主要花在四个地方模型权重、优化器状态训练中才有、KV Cache自回归生成才有、以及中间激活值。Toonflow 用的主干模型是一个 7B 级别的生成式架构半精度模型权重大概是 14GB 左右。如果第一次推理的时候就一块块地加载权重再为每个中间结果申请临时显存前几个 batch 会频繁出现显存碎片化。我建议设计一个显存规划表在服务启动时就明确分出几块区域权重区固定不重复申请、KV Cache 区预分配按最大 batch 和序列长度算好、激活区动态但复用 buffer。这样做的意义在于减少显存管理器的动态分配压力让每次推理的空间地址尽量稳定。3.2 激活显存的重计算与 buffer 复用激活值占用的显存很容易被忽视。在 transformer 层数很深的时候每一层的中间激活张量都会驻留在显存中直到反向传播训练阶段或某一层输出被后续层完全消费完推理阶段才能释放。推理阶段一个直接有效的策略是即时释放在代码里减少对大中间结果的长期引用确保临时张量超出作用域后能被 python 的垃圾回收机制及时处理。另外可以把多个中间 buffer 放到同一个 pre-allocated 池里反复使用。比如 attention 打分矩阵、FFN 中间特征、LayerNorm 输入输出这些张量大小相对规整完全可以用显存池来复用避免每次推理重新触发 allocator。# 使用显存池复用中间 buffer 的简化示意 buffer_pool torch.cuda.mem_pool() def forward_with_pool(x, layer): # 从池中获取一块显存作为中间激活 hidden buffer_pool.get_buffer(x.shape) hidden.copy_(x) # layer 计算 out layer(hidden) # 归还 buffer供下一层复用 buffer_pool.release_buffer(hidden) return out显存池的思路跟线程池很像都是把反复创建销毁的重开销变成一次初始化加多次复用。实测中的效果非常明显连续跑 200 次推理显存峰值从 17GB 降到 13.5GB并且没有出现 OOM而优化前跑到大约 70 次左右显存就开始出现锯齿状爬升。3.3 KV Cache 的预分配策略与序列长度关系自回归生成场景中KV Cache 的显存消耗是随序列长度线性增长的。Toonflow 的部分模块需要生成长文本最多 4096 tokenKV Cache 占用的显存超过了权重的一半这是个非常大的隐性开销。我的优化思路是按最大生成长度预分配 KV Cache 空间而不是按已生成长度动态扩容。动态扩容的过程中会不断触发显存复制和重新分配既慢又容易碎。预分配方案虽然看起来更浪费但换来的是运行时完全不需要管 KV Cache 的扩容问题。# 预分配 KV Cache 的伪代码示例 max_seq_len 4096 max_batch_size 16 num_layers 32 num_heads 32 head_dim 128 k_cache torch.empty( (num_layers, max_batch_size, num_heads, max_seq_len, head_dim), dtypetorch.float16, devicecuda ) v_cache torch.empty_like(k_cache) # 推理时只需要往 pre-allocated 区域写入当前 step 的 K 和 V另外我用了 PagedAttention 类似的策略在 Toonflow 里做了一个简化版把 KV Cache 划分成固定大小的 page每个请求按需分配 page而不是每个请求独占一整块连续空间。这样做的好处是显存利用率更高尤其在不同请求生成长度差异很大的场景下不会因为一个长文本请求占掉空间导致其他请求无缓存可用。4. 实测结果优化前后延迟、吞吐、显存占用对比与分析技术方案讲完之后贴上真实的对比数据。我习惯把优化前后记录成一张表格便于后续回归和复盘。Toonflow 的基准环境如下单张 24GB 显存 GPUPyTorch 版本 2.xCUDA 12.x模型为 7B 级别输入长度 128 token、输出长度 512 token。4.1 核心性能指标对比指标优化前优化后提升幅度P50 延迟4.2s1.4s66.7%P99 延迟6.9s3.2s53.6%吞吐量请求/秒1.23.5191.7%GPU 平均利用率46%81%35 个百分点显存峰值17GB12.6GB25.9%单请求 Token 生成速度82 tok/s227 tok/s176.8%这份数据是几种优化叠加之后的效果单纯靠某一个单项优化拿不到这么明显的变化。举个例子只做算子融合时端到端延迟大约缩短到 3.3 秒但结合 CUDA Graph 后跳到 1.8 秒再加上批处理和显存规划最终落到 1.4 秒左右。4.2 显存曲线和稳定性观察显存峰值下降之外另一个更重要的变化曲线显存占用变得非常平稳。优化前每跑完一次推理显存占用会有一点残留长跑之后逐渐累积最终逼近 OOM。优化后由于采用了固定 buffer 池和 KV Cache 预分配的方式显存曲线始终在 12.6GB 附近微微波动没有爬升趋势。这点对线上服务很关键。很多服务上线初期看着没事跑一周之后突然 OOM其实都是显存碎片加缓存残留慢慢积累出来的。主动规划显存区域能消除这类稳定性风险。4.3 什么场景下这套方案收益有限实践过程中也注意到这套优化不是万能的。当模型本身是稠密计算型、batch 很小且序列很短时launch-bound 和 memory-bound 问题就不突出算子融合和 CUDA Graph 的收益会明显缩小。另外如果显存已经非常紧张比如权重就占了可用显存的 90% 以上那显存规划策略的效果也会打折扣这时候最先考虑的应该是模型量化或者权重卸载。5. 踩坑记录与排查经验这些坑你可能也会遇到优化过程中踩过不少坑记录下来希望对大家有帮助。每一条都是真实遇到并通过日志和性能工具定位解决的。5.1 CUDA Graph 捕获时报operation not permitted这个问题在引入 CUDA Graph 之后立刻遇到。捕获阶段如果触发任何 cudaMalloc、cuDNN benchmark、或者 lazy 初始化都会直接报错或者导致 graph 失效。这里的操作顺序非常重要先做 3 到 5 次 warmup 推理再用一个单独的 stream 做捕获捕获过程中不要插入任何 python 侧的数据处理逻辑。如果使用了动态 shape尽量把形状统一到固定值。排查工具推荐在捕获代码前后各加一次显存快照对比是否出现显存增长。如果捕获过程触发了显存分配快照对比会直接暴露出来。5.2 显存碎片化导致的显存足够但 OOM这是最让人头疼的一类问题。从nvidia-smi看显存还有 6GB 空闲但只要申请一块 4GB 连续空间就直接 OOM。原因在于显存被大量小对象切割成碎片没有足够的连续大块空间。解决办法有两个层面一是减少动态分配次数用显存池或者 buffer 复用二是引入显存整理机制在每次推理结束后调用torch.cuda.empty_cache()但这只是辅助手段不能依赖它解决设计问题。5.3 精度与性能之间的权衡算子融合和 CUDA Graph 在多数情况下不会改变数值结果但有一些融合实现如果使用了近似计算或者降低了中间精度可能会出现生成质量轻微下降。Toonflow 里做注意力和 FFN 融合时我特意保留了 FP16 计算没有使用 FP8因为在长文本生成中精度损失会被逐步放大最终影响可读性。6. 优化框架沉淀一套可以复用的 GPU 推理与显存治理路线Toonflow 的优化过程走到这里其实已经不是一次性改几个函数的事了。我最后把方法和经验收敛成一个可复用的优化框架后续做其他模型推理服务时可以直接套用。这个框架分四个阶段采集、定位、优化、验证。6.1 采集与定位阶段用数据说话不要猜首先把 profiling 工具统一起来。PyTorch Profiler 负责算子级别耗时和显存分配Nsight Systems 负责了解 GPU 活动范围和 Host 侧调度等待显存快照对比用来定位碎片和泄漏。每个优化动作之前和之后都要各跑一组基准固定输入、固定 batch、固定随机种子避免结果不可比。6.2 优化优先级排序先显存后延迟先全局后局部踩过一轮坑之后我现在做优化的优先级排序固定为先看显存规划预分配、buffer 池再看全局执行结构CUDA Graph、动态批处理最后才做局部算子层面的融合。原因很简单显存是地基如果频繁 OOM 或者碎片严重后面做 CUDA Graph、批处理都会受到牵连而算子融合是局部收益应该在全局结构稳定之后再去收尾。6.3 可观测性把优化成果固化到监控体系里优化完成之后建议把关键指标接入监控显存峰值、显存曲线走势、P50/P99 延迟、GPU 利用率、kernel launch 数量。这些指标能帮助你在模型更新或数据分布变化时快速发现是否退化。Toonflow 后续每次做模型结构调整都会重新跑一遍 profiler权重加载到新模型后先做一次小流量验证确认延迟和显存曲线没有异常再全量放开。如果大家正在做的项目规模更大可以进一步考虑多卡推理、张量并行、流水线并行或者更细粒度的显存调度算法思路是类似的先测量、再分析、再做针对性优化。不要一上来就砸钱换更大的卡先把现有资源的利用率提上去后面很多事情顺了之后问题会变得清晰很多。
