基于Agent的SGLang Diffusion推理Kernel自动化调优实践

我今年大概有三分之一的工作时间都花在模型推理优化上说实话普通LLM的serving已经卷得差不多了SGLang和vLLM两家把KV cache、调度、CUDA Graph这些能榨的油水都榨了。但让我真正头疼的是Diffusion模型的推理性能调优尤其是Qwen-Image、FLUX.2这类生成模型还有视频生成模型它们的Kernel特征和纯文本模型完全不是一个路数。最开始我是一边用nsight compute跑profile一边手动去改SGLang的kernel配置改一版测一版一天下来可能只试了三五个组合而且很多组合之间还有隐藏的依赖关系——比如attention backend切换之后FP8的scale策略又得跟着变。后来我索性做了一个基于Agent的调优闭环让语言模型Agent来接管“读profile、提假设、改配置、跑benchmark、复盘结论”这个折腾人的循环。这篇文章就把我这段时间踩过的坑、实际跑出来的优化路径、以及Agent系统怎么设计的一次性讲清楚内容围绕Qwen-Image、FLUX.2和视频模型三个场景展开适合正在做Diffusion推理加速、或者想把自己的SGLang后端调得更顺手的朋友参考。1. 项目初衷为什么我要用Agent去调SGLang的Diffusion Kernel1.1 SGLang的Diffusion后端到底解决了什么问题先说一个容易误解的点SGLang并不是只能跑LLM。它的核心价值在于一套能够感知前缀复用、支持连续批处理和CUDA Graph的调度器这套东西放在Diffusion生成任务里同样是加分项。常规SD WebUI那种一次生成一张图的玩法压根用不到调度器的威力但当你需要同时处理多个用户的生成请求或者是视频模型这种多条候选序列同时去噪的场景批处理和显存复用就变得非常重要。SGLang这套东西套在Diffusion模型上核心逻辑是把Diffusion模型的迭代去噪过程当作一条特殊的前向计算图每一轮denoising step产生的中间张量都可以通过调度器来复用显存、合并算子、共享KV Cache文本编码器部分。尤其对Qwen-Image这类基于MMDiTMulti-Modal DiT的模型文本token和图像token要在attention层做joint attention中间会反复使用text encoder的KV cacheSGLang的缓存机制能把这部分省下来。我在实际部署中碰到的真正痛点是Diffusion模型里Kernel的可配置面比LLM大很多。LLM的推理主结构就是decoder-only transformerattention和MLP相对固定而Diffusion模型里既有卷积、又有线性层、还有各种跨模态attention再加上不同的precision选项、不同的attention backend、不同的内存layout这些维度组合起来几乎不可能靠手工一盘一盘点出最优解。1.2 手工调Kernel的“体力活”陷阱如果只是偶尔调一个模型手工来并不是不行。我记得第一次调Qwen-Image流程大概是这样的先跑一遍torch profiler拿到每个算子的耗时分布发现attention占了40%以上接着把SDPA换成flash-attn显存和延迟立刻好看了一截。但问题很快出现了换了flash-attn之后某个自定义的LayerNorm Kernel反而变成了瓶颈因为设备带宽被flash-attn压得更满全局的性能瓶颈发生了转移。这就是手工调优最恶心的地方局部的最优解会不断改变全局的热点分布。你以为自己在优化实际上是在玩一个动态平衡游戏。我在SGLang的Diffusion后端里切换过attention kernel、调整过FP8量化粒度、改过CUDA Graph捕获策略每个改动单独看都有收益合在一起却可能带来微妙的性能回退。这种情况反复发生之后我就意识到要想高效地找到全局近似最优配置必须有一个能够快速做多轮“假设-验证-归因”循环的系统。1.3 Agent方案的整体选型思路之所以选Agent而不是传统AutoML或者贝叶斯优化是因为Diffusion Kernel优化的搜索空间里很多决策不是简单的连续参数而是需要“理解上下文”的离散选择。比如某个kernel实现明明在GPU A上很快在GPU B上却慢得离谱你需要去判断这是不是与SM数量、共享内存大小、或者LLC容量有关再比如FP8量化对FLUX.2这类flow-matching模型几乎没有损失但对某些视频模型的CFG guidance部分就会崩画质这个判断需要模型层面的知识。我把这些“领域知识”交给Agent本质上不是让Agent凭空创新而是让它基于已有的知识库和规则做推理再利用我提供的一组工具去快速执行实验、获取反馈。整个系统采取“人在环上”的架构Agent负责大量重复性的探测和建议生成我负责最终的代码合并和关键决策审批。做下来之后效果明显胜过我一个人手动折腾。2. Diffusion推理优化到底在调什么核心原理拆解2.1 Diffusion推理全链路中的Kernel热点从宏观链路来看一个完整的Diffusion生成任务包括三个大阶段Text Encoder编码、Denoising循环、VAE Decode。Qwen-Image、FLUX.2、视频模型虽然结构差异很大但这条链路基本一致。Text Encoder阶段Qwen-Image和FLUX.2通常使用双向编码器比如Qwen2.5-VL的text encoder或者T5系列这部分是典型的prefill任务算力消耗大头在长文本的attention上。因为文本token通常就几十到几百个所以这个阶段优化空间不算大但SGLang的RadixAttention在这里能发挥作用——同一批prompt的prefix被多个请求共享时KV cache可以直接复用。Denoising循环才是真正的热点。以Qwen-Image为例20步左右去噪迭代中每一步要处理图像token、文本token、以及它们之间的joint attention。图像分辨率越大token数越多attention的计算量呈平方级上涨。FLUX.2引入了flow matching和MoE结构的DiT去噪过程中的kernel调用更加动态化某些expert分支可能只在特定step或者特定denoising区间被激活。视频模型更狠多了一维时间轴attention变成空间加时间类似3D attention计算量和显存占用完全不是一个量级。VAE Decode阶段相对快主要是卷积层但在视频模型里因为要解码多帧图像VAE的计算也会变成一个不忽略的瓶颈尤其是时序VAE对帧间一致性的处理。2.2 Attention Kernel与精度策略的选择逻辑在SGLang里做Diffusion Kernel优化最核心的两个旋钮就是attention backend和计算精度。Attention backend方面我们实际比较过PyTorch原生SDPA、flash-attn包括2.x版本、以及SGLang内部基于triton实现的专属attention kernel。结论是Qwen-Image和FLUX.2的QKV维度、head数都不一样没有一个kernel是全局最优的。flash-attn 2在大多数情况下是首选但遇到某些batch size小、序列长度又特别长的场景SGLang里针对Diffusion动态shape做的triton kernel反而能通过自动调整block size拿到更优的显存带宽利用率。精度策略上FP8是绕不开的话题。H系列GPU对FP8支持得很好Qwen-Image的MMDiT主干做FP8量化之后显存占用降到原来的60%左右生成质量肉眼几乎不可见损失。但FLUX.2情况略不一样尤其它的guidance机制会对中间特征做额外的高精度累加FP8量化如果scale选得不好画面上会出现轻微的水波纹我们后面是通过per-tensor改成per-channel量化解决的。2.3 Qwen-Image、FLUX.2和视频模型的结构差异这三个模型放在一起调优是因为它们的Kernel特征差异足够大值得作为Agent学习的三组样本。Qwen-Image本质上是2B参数的MMDiT模型图像和文本在注意力层直接融合Kernel热点集中在joint attention和feed-forward层。优化它时我优先关注的是如何让QKV融合kernel用上减少中间张量的写回带宽。FLUX.2是更大体系下的flow-matching DiT模型还引入了MoEMixture of Experts结构。去噪时只激活部分expert导致kernel dispatch动态性极强。SGLang对这种动态性的支持主要体现在CUDA Graph策略上如果盲目把整个模型捕获进一张静态CUDA Graph里一旦激活的expert集合发生变化graph就必须重新捕获那开销小不了。所以Agent很快就学到了一个经验FLUX.2适合用分段CUDA Graph按expert activation pattern分组捕获。视频模型则是另一个极端以时序建模为特征的模型会把视频潜变量拆成T, H, W, C的四维形式attention要在空间维和时间维分别做。这种条件下flash-attn的标准接口不太好用而SGLang的自定义kernel需要额外处理temporal attention的layout变换。显存管理方面视频模型多帧共享同一个文本KV cache这个特性让我们在调度器里做了专门的帧间缓存复用。3. Agent系统设计与实现工具、记忆与迭代闭环3.1 Agent的整体架构与各组件职责我们的Agent优化系统分三层。最上层是决策层我用了Qwen系列的指令模型作为大脑它负责理解优化目标、分析profiling报告、提出下一轮实验假设。为了模拟一个“资深优化工程师”的思考过程我在System Prompt里塞了大量Diffusion Kernel优化的先验知识包括各类kernel的适用条件、FP8量化的常见坑、SGLang调度器的行为特性等等。中间层是工具层Agent不直接读写文件、不直接跑命令而是通过调用我封装好的工具来做事情。每个工具都有严格的输入输出schemaAgent只需要学会“用工具”就行不需要真的去拼命令行。最底层是执行层就是SGLang的推理服务、profiler、benchmark脚本以及显存监控。这一层负责产生真实数据所有数据再流回Agent的上下文形成闭环。3.2 给Agent的四大工具设计工具设计是整个Agent能否落地的关键。我一开始犯过错误给了Agent一个“万能Shell工具”让它什么命令都能跑结果它有一次递归调用了profiler导致服务OOM教训很惨痛。后面我把工具收敛成四个第一个是ProfileTool封装了torch profiler和NVIDIA nsight compute的调用。它输出一份结构化JSON报告包含每个算子的调用次数、平均耗时、GPU kernel利用率、显存使用峰值以及卡在tensor core上的比例。这份报告是Agent做归因分析的最核心输入。第二个是ConfigTool它只允许Agent读取和修改特定配置文件里的白名单字段。这些字段包括attention_backend、precision、quantization_scheme、cuda_graph_mode、max_batch_size等。每次修改之前ConfigTool都会做一次非法值拦截防止Agent填出一些明显不存在的kernel名字。第三个是BenchmarkTool固定跑一组我提前定义好的真实负载包括不同batch size、不同分辨率、不同去噪步数的组合。它输出的指标是SGLang里大家最关心的单图生成延迟、每步平均耗时、吞吐、以及显存峰值。第四个是HintTool这个工具很“作弊”它是我把维护多年的经验知识做成的检索库。Agent在犹豫不决的时候可以调用它比如“这个模型用FP8之前先检查一下guidance支不支持”HintTool会返回对应的提示词。后来我又顺手接了一个基于qwen3-reranker的模块对候选优化方案做相关性排序帮Agent更快过滤掉明显的坏点子。3.3 Agent的推理、记忆与收敛策略有了工具还不够Agent要真正干好活还差两个关键点短期记忆和收敛判断。短期记忆我直接采用了最朴素的方式把每一轮实验的假设、执行结果、Agent自己的复盘结论都追加到上下文中。因为Qwen模型的上下文窗口很长跑40轮实验完全够用。每轮实验结束后我还让Agent把“当前认为最优的配置”单独存下来方便最后回滚或者导出。收敛判断也做得很工程化。我给Agent设置了两个硬性条件连续5轮实验没有性能提升或者总耗时超过预算就自动终止。此外还有一个GuardrailTool负责实时监测benchmark结果如果发现本轮配置比历史最优值差了超过10%立即给出警告Agent必须决定是回滚还是继续微调。3.4 Agent主循环代码骨架实际写出来的主循环并不复杂用伪代码表达大概是下面这样。这也方便你在自己的项目里快速搭一个原型出来。# agent_diffusion_tuner.py # 去掉工程细节之后的骨架理解为主 class DiffusionKernelTuner: def __init__(self, agent_model, tools, max_iters40, budget_sec1800): self.agent agent_model self.tools tools self.max_iters max_iters self.budget_sec budget_sec self.memory [] self.best_config None self.best_score float(inf) def run(self, task_desc): start time.time() for i in range(self.max_iters): elapsed time.time() - start if elapsed self.budget_sec: break # Agent基于已有记忆提出下一步实验假设 hypothesis self.agent.propose(task_desc, self.memory) # 按hypothesis修改SGLang配置 config self.tools[config].apply(hypothesis) if config is None: continue # 跑benchmark拿到结构化结果 metrics self.tools[benchmark].run(config) self.memory.append({hypothesis: hypothesis, config: config, metrics: metrics}) # 更新历史最优 if self.is_better(metrics) and self.is_safe(metrics): self.best_config config self.best_score self.score(metrics) # 收敛判断连续N轮无提升 if self.no_improvement_streak() 5: self.agent.summarize(self.memory) break # Guardrail回滚危险配置 if self.regression_exceeds(metrics, threshold0.10): self.tools[config].rollback() return self.best_config这里特别说明一下Agent的propose方法不是简单的只生成一段文本它会生成一个结构化的“实验提案”包括对当前瓶颈的判断、想改的配置项、预期效果和风险。工具层解析这个提案时会做类型校验和取值枚举不合法就直接拒绝并把原因反馈给Agent让它重新提。4. 实操全记录从Qwen-Image到FLUX.2再到视频模型4.1 环境准备与基线测量先交代实验环境。我们用的是8卡H800节点CUDA 12.4PyTorch 2.5.1flash-attn 2.6.3SGLang是当时的main分支代码。模型方面Qwen-Image用的官方2B版本FLUX.2用safetensors从diffusers转换到SGLang的Diffusion runtime里视频模型用的是我们内部基于时空DiT练的模型支持16帧768分辨率输入。基线测量阶段我没有让Agent直接介入而是先手工确认服务能正常跑通。Qwen-Image当batch size等于1、分辨率512x512、20步去噪时单次完整生成的延迟大概在4.2秒左右其中denoising loop占3.5秒VAE decode占0.5秒。FLUX.2同样条件下要重一些因为模型参数量更大单张图延迟接近7.8秒。视频模型16帧768分辨率完整生成一次要超过48秒明显有大量优化空间。基线数据很重要因为Agent后续的所有判断都是基于“相比基线有没有提升”来做决策的所以基线必须稳定。我建议每次优化前都至少跑三轮benchmark取中位数避免系统噪声干扰Agent的判断。4.2 第一轮优化Attention Kernel切换与FP8精度Agent接管之后做的第一轮决策跟我之前的经验高度重合先把Qwen-Image的attention backend从SDPA切到flash-attn结果显示单步平均耗时从175ms降到了132ms降幅约25%。这个收益主要来自flash-attn对tiling和shared memory的优化减少了全局内存读写。接下来Agent开始尝试FP8。它对Qwen-Image做了per-tensor的FP8量化但第一次benchmark结果让我有点意外显存降下来了延迟却几乎没变。Agent后来归因发现问题出在量化转换kernel上每次前向都要把weight做一次dequant再计算等于多了一轮额外的内存搬运。正确的做法是提前把权重在加载阶段转成FP8存储前向计算时直接走FP8 GEMM。修改之后单步平均耗时降到108ms显存峰值从16.2GB降到10.8GB。这一轮还有一个非常值得说的细节Agent尝试把整个MMDiT的主干都塞进CUDA Graph结果第一次捕获就失败了因为attention计算内部有动态shape分支。HintTool及时给出了提示让Agent改用“按shape分组捕获”的策略问题迎刃而解。这算是我最满意的一个Agent自救案例。4.3 FLUX.2的独特优化点MoE DiT与Joint Attention在Qwen-Image上吃到甜头之后我把同一套Agent流程推到FLUX.2上。FLUX.2和Qwen-Image的结构差异很快暴露了出来。首先最明显的是MoE结构。Agent通过ProfileTool发现FLUX.2在denoising的不同阶段激活的expert数量差异巨大早期step只激活2个expert后期可能激活6个甚至更多。这导致模型的计算强度是动态的静态CUDA Graph完全不适配。Agent提出的方案是把denoising过程按step分成几个区间每个区间分别捕获一张CUDA Graph这样虽然增加了首次构建的耗时但实际推理时kernel launch开销大幅下降单步延迟最终降了18%。其次是Joint Attention的处理。FLUX.2的joint attention把文本和图像当作一个序列来处理序列长度扩展到原来的1.5倍。这一个改动直接让attention计算量暴增。Agent的应对是引入split-k策略把长序列的attention在sequence维度上切分成多段用不同的stream并行计算然后再做segment级合并。这个优化在batch size大于1时效果尤其明显因为并行度填充得比较满。我额外手动介入了一次因为Agent一开始想在FLUX.2上做与Qwen-Image一样的FP8量化但我根据线上效果踩坑经验让它先检查CFG guidance的数值范围。果然Agent检测到CFG分支的中间特征标准差偏大如果沿用固定scale量化误差会被guidance放大。最终方案是CFG计算路径保持BF16主生成路径用FP8代价是显存优化幅度比Qwen-Image小一点但画质完全无损。4.4 视频模型时空注意力与KV Cache复用视频模型的优化是整个项目里最惨烈也最有收获的部分。我让Agent先对基线做了完整profile发现单步去噪耗时里有将近一半花在了空间attention上另外有20%的时间消耗在3D张量到2D张量的layout变换上。视频模型用的是3D DiT但底层flash-attn只支持2D矩阵乘所以每一帧的token都要重新reshape带来大量L2 cache miss和同步开销。Agent针对这个问题的解法是把时空attention解耦成两个阶段先沿空间维度做完整的cluster attention再沿时间维度做一个轻量级的temporal attention。这样虽然理论上多了一次attention pass但每一层的矩阵维度更规整反而能打出更高的throughput。实测下来空间注意力那些layout变换几乎全部消失整体单步耗时反而降了22%。KV Cache复用机制也是视频模型优化的重头戏。视频模型一次生成16帧这16帧共享同一个文本编码结果。Agent通过在SGLang调度器里增加了帧级KV Cache池让16帧生成的文本KV只计算一次后续所有去噪step都能复用。最终视频模型从48秒降到了33秒左右其中注意力优化贡献三分之一KV缓存复用贡献三分之一剩下来自kernel精度调整和显存分配策略的优化。5. 踩坑实录与问题排查技巧5.1 Agent调参最容易翻车的情况Agent在自动调优过程中翻车场景比我原本预想的多。最典型的坑是“过拟合benchmark集合”。有几次Agent为了降低benchmark里的平均延迟疯狂针对batch size等于1的负载优化结果真实业务里batch size等于4的请求进来后性能全线崩溃。后来我给BenchmarkTool增加了一个权重机制线上流量分布决定不同负载的加权占比Agent优化的是加权后的综合收益才算是治好了这个问题。另一个常见翻车是Agent对“精度无损”过于乐观。它在量化FLUX.2时只看了benchmark指标没看生成效果导致某轮实验的FID图像质量指标差了3个点。妥协的方案是工具层加入了图像质量自检模块每次benchmark结束之后用一组固定的prompt生成样例图计算与基线输出之间的LPIPS距离。超过阈值整轮实验就不予采纳。5.2 常见错误速查表下面的表格是我在实际优化过程中遇到的错误、原因和对应解决办法按出现频率排了个序。异常现象根因分析解决办法CUDA error: no kernel image available编译出一版精度格式不兼容的kernelSGLang调用时的dispatch key与实际注册kernel不匹配重建flash-attn和SGLang扩展确保与CUDA版本匹配在ConfigTool中增加kernel可用性探测torch.accelerator error: cuda error OOM视频模型帧数增加后KV Cache池没有自动扩容或者CUDA Graph预留显存过大让Agent读取实际显存峰值动态调节graph预留缓冲区帧间缓存复用开启Agent连续几轮实验都是相同配置工具输出格式不明确Agent没理解到“当前配置已经是历史最优”在ProfileTool输出中加入与基线对比的差值字段引导Agent重新归因Benchmark结果波动超过5%节点上有其他任务抢占GPU增加预热次数benchmark前后查询GPU利用率FP8量化后生成图像出现色块scale选择粒度过粗某些通道动态范围过大切换per-token或per-channel量化CFG分支保留高精度Flash-Attn在视频模型时空注意力上报错3D attention需要输入维度是3D直接用2D kernel处理4D张量导致崩溃显式做空间/时间注意力分离分别在降维后的2D张量上调用kernel5.3 我保留的几个“反直觉”小技巧最后分享几个不太常见、但在我实际项目中让Agent少走弯路的小技巧。第一让Agent定期清理自己的“记忆”。上下文虽然很长但几十轮实验下来早期无效实验会对决策产生噪声。我在每次Agent做新的proposal之前会调用一个compress方法把历史实验用几句话概括成“教训”压缩进上下文的最前面形成一个动态的经验卡片。效果非常明显后20轮的决策质量明显高于前20轮。第二不要期待Agent产出完美代码。我把Agent定位成“实验设计器和数据分析师”而不是“代码提交者”。它负责告诉我“应该试什么”真正的代码修改还是由我来写。这样既保留了Agent的高效探索能力又规避了它生成低质量代码带来的风险。第三裸的Agent跑不满GPU别让它“想好了再动手”。我一开始给Agent设置了很激进的节能策略每步实验间隔3秒方便它思考后来发现完全多此一举。LLM推理agent本身是在CPU上跑的和GPU实验完全可以并行我改成“异步推进”之后整个优化循环的吞吐提升了将近一倍。这个优化系统的意义对我来说不只是把Qwen-Image做到了单张图3秒出头、FLUX.2做到了5.6秒、视频模型压到了26秒更在于它把“调kernel”这件事从一门手艺变成了一个可复制、可复盘、不给人类工程师添堵的流程。你完全可以把这套Agent思路推广到文本生成、语音合成等其他推理场景核心就是构建好工具闭环、设定清晰的收敛条件剩下的脏活累活让Agent去干就好。