用AI Agent自动化优化SGLang Diffusion推理Kernel的实战记录

前一阵子我们组的推理服务在SGLang上跑diffusion模型结果有点尴尬同样一个Qwen-Image的1024分辨率请求我们内网一张A100要跑7秒多而参照服务能够压到5秒以内。一开始我以为是模型版本或者采样步数配置的问题后来用profiler抓了一遍kernel热点才发现瓶颈根本不在模型图本身而在kernel层——SGLang虽然把LLM推理的调度做到了极致但diffusion这条链路尤其是从Qwen-Image切到FLUX.2再切到视频模型时大量算子还是朴素的eager实现没有针对多模态DiT做融合CFG的batched执行路径也没完全打通。这篇文章就是我过去一段时间基于Agent来做SGLang Diffusion Kernel优化的完整记录涉及Qwen-Image、FLUX.2和一个16帧的视频扩散模型适合正在做diffusion模型serving、推理加速或者对AI Agent落地感兴趣的人参考。1. 为什么Diffusion推理在SGLang上还有那么多kernel可挖1.1 Diffusion推理和LLM推理根本不是一回事SGLang的价值主张一直是围绕LLM自回归解码设计的RadixAttention拿前缀树缓存重用KV CachePD分离把prefill和decode的负载拆开调度连续batch动态拼装请求。这些优化在token逐字生成的任务里收益巨大因为注意力计算在decode阶段是memory-boundKV Cache的命中率直接决定吞吐。但Diffusion模型完全是另一种玩法。Diffusion模型通过多步去噪生成图像或视频每一步的输入是一份带噪latent经过完整的Transformer或U-Net前向计算后输出一份新的latent。这里没有当前token和历史token的区别也没有KV Cache可复用的说法——每个step输入都在变网络结构从头到尾完整执行一遍。这意味着SGLang在LLM场景下的核心调度优势到了Diffusion场景并不能直接兑现。真正决定Diffusion推理快慢的是每一个kernel的实现质量QKV投影是否做了融合、attention是否用上了高效的flash变体、LayerNorm和linear能不能合并、FP8路径是否被正确启用。我在刚开始接手这个项目的时候同事说反正SGLang已经支持FLUX了直接用就行结果一跑profile就发现block里有大量零碎kernel一个linear拆成好几个小kernel执行attention里mask矩阵被显式物化出来文本token和图像token在同一个flash-attention分块逻辑里互相拖累。这些都是优化空间。1.2 Kernel层三大优化空间融合、精度、调度复用我把Diffusion在SGLang上的kernel优化空间归纳成三个层面第一个是算子融合。Diffusion模型一个block里通常是Norm、QKV投影、attention、FFN、残差连接这一串。eager模式下这串操作可能对应几十次kernel启动每次启动都要把中间结果写回显存下一阶段再读回来。融合的目标就是把这些算子揉成一个大kernel中间张量只活在寄存器或共享内存里。图像模型的latent序列已经是大几千到上万token级别每次kernel面对的矩阵都很大减少几十次显存读写收益非常可观。第二个是数值精度路径。当前Diffusion模型主流的推理精度是BF16但很多矩阵乘法是可以用FP8跑的。FP8 matmul的关键在于scaling策略——per-tensor还是per-channel什么时候做量化什么时候反量化不同shape下效率差异很大。如果只是简单地在模型外面包一个全局FP8开关很容易出现某些层变快、某些层反而变慢甚至数值崩坏的情况。按层决策是必须的这也是Agent很适合干的活。第三个是调度级复用。Diffusion推理普遍用Classifier-Free GuidanceCFG意味着每个采样步都跑两遍网络一遍带条件一遍不带条件。两遍网络在早期层的输入几乎一样区别主要来自文本条件参与注入的位置。如果在kernel层能做分支合并或者共享中间结果理论上能省掉接近一半的计算量。但这件事实现复杂度高而且容易出数值问题后面我在视频模型的章节会详细展开。1.3 SGLang的优势换成Diffusion后还剩下什么既然调度优势不能直接复用那SGLang在Diffusion场景还剩下什么这是整个优化工作的前提认知SGLang提供的是一套完整的模型运行时——显存管理、模型并行、后端自动调度的骨架以及一个活跃的社区不断往里面贡献针对Diffusion架构的kernel补丁。它比我基于原生PyTorch脚本写serving要扎实得多而且社区里对FLUX、Qwen-Image的适配进度很快很多坑已经被踩过。我要做的是在它的kernel层之上做文章而不是从零搭一套推理引擎。所以整个项目的切入点非常明确保留SGLang的运行时和调度框架把精力集中在替换/优化关键kernel并且用Agent来自动化发现瓶颈-提出优化-验证收益-沉淀经验这个循环。2. Kernel优化Agent的模块设计我把调优流程做成了闭环2.1 从人肉调优到智能体的动机Diffusion模型迭代太快了。上个月还在调Qwen-Image的MMDiT这个月FLUX.2出来了架构细节一变之前手工调优的经验有一半作废下个月又要上视频模型显存和时序维度又带来新的麻烦。如果每次都是人来盯profiler报告、人肉改kernel、人肉做A/B测试一个模型至少消耗两个人周。我在这个过程里明显感到手跟不上模型发布速度于是决定做一个Kernel优化Agent把调优流程自动化。这个Agent的定位不是自动写CUDA代码那种野心很大的东西而是像一个不知疲倦的调优实习生它会跑profile、看热点、翻知识库、提出改动方案、跑benchmark验证、把结果记下来。真正涉及深层kernel开发的活它负责生成候选实现和对应的验证报告人来拍板。是这样一个人机协同的定位。2.2 Agent的模块划分与工具集Agent本身没有上重型的Agent框架核心就是一个基于Python asyncio的编排层把下面几个工具封装成可被大模型调用的函数Profiler定时抓取torch profiler和Nsight Compute输出kernel级热点排名、每个算子的耗时占比、显存读写总量以及算术强度估算。Analyzer解析SGLang加载的模型配置文件把热点kernel映射到具体的Module和算子类型判断瓶颈是memory-bound还是compute-bound。这个映射是很多人容易忽略但极其重要的一步——不知道热点属于哪个结构就无法提出有效的优化假设。Planner基于瓶颈类型和历史实验记录从候选动作池里选优化项。候选动作池包括替换attention kernel、调整block size、融合normlinear、启用FP8、切换数据排布、调整双流并行方式等。Executor在隔离环境执行改动跑端到端benchmark。我们的benchmark协议固定为warmup 3次正式采样20次生成统计p50耗时、p95耗时、峰值显存。Verifier检查生成质量不只跑CLIP score还会做结构化相似度对比和简单的伪影检测。Memory每次实验的完整信息模型、环境、改动项、收益、副作用写入SQLite供Planner后续检索。这些模块之间的协作逻辑很简单Profiler产出事实Analyzer把事实翻译成结构Planner基于结构提假设Executor验证假设Verifier判断质量Memory沉淀经验。一个循环跑完要么合入一个优化项要么多一条此路不通的记录。2.3 一轮优化迭代是怎么跑通的拿Qwen-Image 1024分辨率、50步CFG的场景举例。Agent跑第一轮profile后发现attention内部的softmax与mask处理占了总耗时18%而且通过算术强度估算这一段是典型的memory-bound——问题在于当时的实现显式构造了一个大mask矩阵每次都要把它从显存里读一遍。Analyzer把热点定位到Attention模块的mask实现Planner从历史库里查到一条之前验证过的经验logits计算前先对mask做稀疏化只保留需要参与attend的token可以显著减少显存读取。Executor按照这条经验把mask改成了稀疏索引格式重新跑benchmarkattention耗时从18%降到9%端到端p50提升了差不多7%。Verifier确认生成图像质量没有任何下降。于是这个改动被合入候选补丁集Memory里多了一条Qwen-Image系列适用稀疏mask优化的记录。整个过程大概跑了40分钟如果换成人肉来做从profile到验证至少一个下午。而且Agent可以持续跑很多轮人只需要在晚上看一眼它提交的报告。这也是我在文中反复强调的一个观念Agent的价值不在于一次性找到一个惊天优化而在于把发现-验证-沉淀的循环从人肉执行变成自动化脚本让经验积累的速度追上模型迭代的速度。这里有必要补充一个工程纪律Agent默认运行在dry-run模式任何改动都必须先生成完整的基准测试报告和质量检查记录经过人确认之后才能进入SGLang的正式分支。没有人为干预Agent不能直接改线上服务的代码。这个约束避免了Agent在探索过程中把服务搞挂也让后续的优化项都可追溯、可回滚。3. Qwen-Image与FLUX.2架构差异下的kernel策略分叉3.1 两种模型架构的注意力组织方式很多做推理优化的同学有个习惯把模型当成一个黑盒只看kernel热点不看架构然后套统一的优化模板。这个习惯在LLM上问题不大因为各家LLM的架构相似度很高但在Diffusion模型上会踩大坑——Qwen-Image和FLUX.2的注意力组织方式完全不同优化策略几乎处处分叉。Qwen-Image是典型的MMDiT多模态DiT架构文本token和图像token在同一个Transformer block里做联合注意力。一句话prompt可能只有几十个到几百个token而1024分辨率的图像latent展平后是上万个token。在联合注意力里图像token占据绝对主导计算量和显存访问都被图像侧带走文本token是一个很小的尾巴。FLUX.2延续了FLUX.1的双流double stream设计图像流和文本流在模型前半段各自独立跑注意力到后半段的single-stream层才合并交互。图像流和文本流各自有独立的QKV投影和注意力计算序列长度差异巨大但计算流程是并行的。这个结构在SGLang的eager实现里默认被串行执行先跑完图像流的attention再跑文本流的attention白白浪费了GPU的并行能力。3.2 Qwen-Image上重点动的三个kernel针对Qwen-Image我们主要做了三个层面的优化这些优化都是Agent一轮轮跑出来的第一个是QKV投影与RoPE的合并。Qwen-Image的图像token数量大QKV投影对应的矩阵乘法是一个典型的大gemm。原来RoPE是单独一个kernel做的需要在投影之后把数据从显存读一遍再写一遍。合并之后RoPE直接在gemm的epilogue阶段完成省掉了整块中间张量的显存读写。这一个改动在1024分辨率下大约带来6%-8%的端到端提升。第二个是文本token的稀疏attention。由于文本token占比小跟图像token一起走标准的flash-attention分块逻辑时会因为分块末尾的padding产生大量空转。我们把文本token部分拆出来用一个小而专用的kernel单独处理注意力。由于文本token数量少这个小kernel的启动和计算开销都很低整体attention耗时大概省下来5%-10%不等。第三个是按层启用的FP8 matmul。Qwen-Image对FP8的容忍度比预期好但我们发现并不是所有层都能无脑量化attention里的Q/K投影如果转成FP8图像质量会有可感知的下降而FFN的第一层和第二层转FP8后没有肉眼可见的问题。所以最终方案是FFN层FP8 attention保持BF16配合per-tensor的scaling策略。这个组合让总耗时再降了约10%。3.3 FLUX.2的并行文本流从reduce到split的转变FLUX.2的优化路径跟Qwen-Image完全不同。前面说了双流架构里图像流和文本流各自独立但SGLang的默认实现是串行的。Agent在Analyzer阶段通过解析模型配置发现模型配置里存在两个独立的stream结构并且没有数据依赖。Planner随即提出了一个假设把两个流的attention kernel发到同一个CUDA stream上以不同的block group并行执行而不是串行等待。这个方案听上去简单但实现时涉及kernel的grid布局调整。两个流各自的序列长度、头数、head_dim都不同如果只是简单地把它们concat成一个batch会在数据排布和attention mask上互相污染。我们的做法是让融合kernel支持多组QKV的block切分每个流的attention块在各自的block group里执行再用原子计数或事件机制保证single-stream层开始前两个流都已完成。另一个FLUX.2的优化重点是single-stream层的合并交互。这个阶段两个序列要同时参与attention我们沿用了为多模态FlashAttention设计的分块策略把两个流的block按比例分配避免某一组block算完空等导致SM闲置。实际跑下来双流并行加上single-stream融合kernel让FLUX.2端到端p50降低了大约27%其中双流并行贡献居多。3.4 用同一套Agent策略结果差异为什么这么大我在第2章提到过Analyzer模块有一个关键能力把kernel热点映射到具体的模型结构。这个能力在Qwen-Image和FLUX.2的对比上体现得特别明显。Qwen-Image的瓶颈主要集中在大token量的线性层计算属于典型的compute-bound或者访存密集的gemmFLUX.2的瓶颈则更多出在两个流没有并行起来和single-stream层的block调度不平衡上属于执行布局问题。如果Agent没有架构感知能力只按热点排名去调整它会在Qwen-Image上试着调attention block size在FLUX.2上也会做同样的事——然后得到完全不同的收益曲线和副作用。架构感知让Agent能够识别出这里的问题是调度布局不是算子本身从而选择完全不同的动作。这也是我觉得未来Kernel优化Agent和传统AutoTVM最大的差别传统AutoTVM在算子级做搜索而基于Agent的方案在架构理解-瓶颈定位-动作选择这个更高级的层次上工作。4. 视频模型的三道坎显存墙、时序注意力与CFG合并4.1 显存墙视频序列让中间激活爆炸视频扩散模型本质上就是把多帧图像放到一个序列里生成但多了一个时间维度。一个16帧、每帧512x512的视频模型latent序列长度大约是单张图像的16倍。这个规模的变化直接导致中间激活爆炸eager实现下一个普通block的前向计算就能吃下大几GB显存整个模型跑起来峰值轻松突破70GB在80GB卡上只剩很少余量。这给kernel优化带来的约束是深刻的。算子融合通常是减少中间张量物化的利器但融合kernel往往需要更多寄存器或共享内存在序列极大时反而可能降低occupancy。所以视频模型上的融合策略必须更谨慎优先选择共享内存友好、block内不需要承载过大chunk的kernel设计。Agent在视频模型优化时的目标函数也要调整从单纯的最小化p50时延变成在峰值显存不超过阈值的情况下最小化端到端时延。我们对一个16帧的视频扩散模型做了chunked attention改造把长序列切成段每段内部完成attention计算段与段之间只传递必要的KV信息。这个过程类似LLM推理里的序列切分但更要小心时序维度的完整性。改完之后峰值显存从71.8GB降到了60.1GB单步时延也有约30%的下降因为每个kernel处理的数据量变小后L2 cache命中率上来了。4.2 时序注意力没有现成的fused kernel视频DiT和图像DiT最大区别就是多了temporal attention时序注意力。空间注意力在每一帧内部算时序注意力则要跨帧拉取同一位置的token做attention。这个操作在数据布局上意味着大量reshape和transpose如果按eager实现每一次跨帧采样都是一次显存搬运。市面上的flash-attention变体都是为标准的序列维度设计的没有现成kernel可以直接处理时间维跨帧采样这种访问模式。我们的做法是用Triton写了一个专门针对temporal attention的融合kernel把QKV投影、reshape、跨帧attention融合在一起用共享内存缓存跨帧数据块。这个kernel的block划分以时间维为基准每一组block负责一段帧区间区间内完成全部跨帧注意力。但是这里有个很现实的麻烦不同视频模型放temporal attention的位置、数量、配置差异非常大。为某个模型写死的kernel大概率无法复用到另一个模型。所以Agent在这条路径上需要根据模型配置动态生成Triton kernel的模板参数而不是维护一堆静态kernel。这让我意识到传统的手写CUDA优化在视频模型这个领域会越来越力不从心模板化和自动化是绕不开的方向。4.3 CFG合并激进方案被质量门禁拦下视频模型的CFG执行比图像模型更痛苦。图像模型可以直接在batch维度把conditioned和unconditioned两个样本拼在一起跑显存翻倍也就多十几GB在80GB卡上通常可以接受。视频模型如果照搬这个做法显存直接翻倍在16帧512x512的场景下会直接撞上显存上限根本跑不起来。SGLang默认的策略是对两个分支串行执行每次denoiser跑两遍显存可控但耗时翻倍。Agent一开始尝试了一个更激进的方案在early层让两个分支共享中间结果只从文本条件真正开始影响网络输出的层之后再做分支。这个方案理论上能省接近一半计算量但实际跑出来出现了数值偏差生成视频在若干帧上出现了肉眼可见的闪烁伪影。Verifier拦下了这次改动CLIP score虽然只掉了0.01但结构化相似度指标掉得非常明显。最后我们退回到保守路线两个分支保持串行但把conditioned和unconditioned的两次前向尽量放进同一个kernel里分块执行至少让权重的加载只做一次减少重复的显存读取和kernel启动开销。这个方案没有激进方案的理论上限高但胜在稳妥而且实测也有接近20%的端到端收益。这次经验被记进Memory之后Agent后续再遇到类似CFG分支共享的假设会先查一遍质量门禁历史不会轻易走激进路线。4.4 视频模型的评估标准不能只看单步时延视频模型优化有个特别容易犯的错误只看单步时延。单步从1.8秒优化到1.2秒听起来很漂亮但如果它是以显存峰值剧增换来的而显存爆掉之后SGLang开始做swap或者OOM重启端到端反而更慢。视频模型生成一条16帧视频通常要跑几十个去噪步端到端时延是去噪步数×单步时延叠出来的任何一步的稳定性问题都会被放大。所以Agent在视频模型上的benchmark协议跟图像模型不一样。除了统计p50单步时延我们还会要求它输出完整生成一条视频的总时长、峰值显存、以及连续生成10条视频是否出现OOM。只有三个指标全部达标优化项才会被标记为合入候选。这套评估标准后来也被证明是必要的——有几轮优化单步时延很漂亮但连续生成到第三条视频就触发显存碎片问题被直接否掉。5. 实测结果和四类踩坑记录5.1 基准测试方法与优化前后对比在展示数据之前先交代清楚测试环境方便你复现时对照GPU是A100 80G PCIeCUDA 12.4PyTorch 2.4SGLang基于社区最新release分支Agent运行在独立的一台CPU节点上通过SSH下发任务。测试集是10条中文prompt加10条英文prompt图像模型固定1024x1024分辨率视频模型固定16帧512x512。各项优化合入后的实测结果如下表所有数据都是warmup 3次之后正式采样20次取p50耗时和峰值显存。模型优化前p50优化前峰值显存优化后p50优化后峰值显存端到端提升Qwen-Image 1024x10247.2s26.0GB5.1s25.8GB约29%FLUX.2 1024x102411.4s30.0GB8.3s29.5GB约27%视频模型 16帧512x51295.2s71.8GB68.4s60.1GB约28%这里要强调一点这些数据是基于我们内部环境、我们选定的这版模型checkpoint、我们的测试prompt集合得到的不代表通用结论。不同的GPU型号、不同的SGLang版本、甚至不同的CUDA版本都可能让数字浮动但优化方向的参考价值是确定的。5.2 踩坑一CUDA error: no kernel image硬件代际与编译目标不匹配这个报错在diffusion模型部署里太典型了值得单独写一节。我们一开始在A100上跑没有问题但同样的代码挪到一台T4上SGLang加载模型后直接报torch.acceleratorerror: cuda error: no kernel image is available for execution on the device这个报错的意思是当前设备上找不到可执行的kernel镜像。根因是SGLang源码编译时指定或默认生成的CUDA架构列表里没有包含T4对应的SM_75架构导致GPU无法加载编译产物。排查路径是先用nvidia-smi确认GPU型号再检查编译时的TORCH_CUDA_ARCH_LIST环境变量然后用cuobjdump --list-elf查看编译产物里包含哪些架构的SASS。解决方法是重新编译把目标架构显式加全比如TORCH_CUDA_ARCH_LIST7.5;8.0;9.0并且把CUTLASS、FlashAttention等第三方依赖的架构参数一起改掉。这条坑给Agent的启示很直接部署环境的GPU架构信息必须是Agent输入的一部分。Agent生成或启用的kernel如果只针对单一架构做编译很可能在另一张卡上直接挂掉。所以我们在Agent的环境检查工具里加了一条GPU算力版本检测任何优化项合入前都要确认目标机器的架构列表覆盖。5.3 踩坑二dflash开启后的收益与副作用SGLang社区为Diffusion模型引入了一条高速执行路径社区里习惯叫dflash开启后FLUX系列模型的生成速度有明显改善。我们在标准1024x1024分辨率下实测开启dflash后FLUX.2的端到端耗时可以再降8%-10%这个收益很诱人。但副作用也随之而来。当输入分辨率不是标准正方形比如720x1280这种竖构图或者batch里混了多种分辨率时dflash的patch划分逻辑会失衡导致一部分SM空转。有一轮测试甚至出现了比关闭dflash还要慢12%的倒挂情况峰值显存也有小幅上涨。我们的处理方式是在Agent的决策规则里加一条分辨率约束只有当分辨率能够被dflash的patch size整除且宽高比落在合理区间时才启用dflash路径否则回退到通用融合kernel路径。这条规则不是Agent自己发现的是人review日志时发现异常后手动加进去的但它很快被验证有效后续也变成了Memory里的一条强力经验。5.4 踩坑三Agent的一次误优化被质量门禁拦下这是Agent体系上线以来最有价值的一次事故。Agent在优化Qwen-Image的attention时把flash-attention的block size从128调到了256benchmark显示p50提升了6%看起来是一个干净的胜利。但Verifier在质量检查时发现CLIP score从0.318下降到了0.302再一看生成图出现了明显的条纹伪影。排查后的根因是block size的变化改变了attention内部数值累加的顺序和精度这种细微的数值差异在LLM任务里可能无所谓但在图像/视频生成任务里会被迭代式去噪放大成可见伪影。生成模型对数值路径的敏感度远超判别式模型这是我们在优化LLM时不太容易形成的直觉。这次事故之后我们在Verifier里加入了三重质量检查结构化相似度指标、CLIP score、人工抽检。并且定了一条硬性规则任何会改变数值路径的优化项比如调block size、切换attention变体、改精度策略必须同时通过质量门禁才能合入。这条规则后来写进了MemoryAgent没有再犯过同样的错误。5.5 踩坑四Agent日志的可读性设计最后一个偏工程的坑。Agent跑久了之后日志量会变得非常大如果日志格式只适合人读检索效率极低如果只适合机器检索人review的时候又很痛苦。我们的解决方案是双轨制每个experiment生成一个JSONL格式的结构化记录方便SQLite导入和后续检索同时每个关键决策点会生成一段自然语言的summary方便人类在晚上快速浏览当天Agent干了什么。这个设计在一次排查线上问题时帮了大忙我们想查Agent之前是否在视频模型的CFG分支共享上做过尝试直接在SQLite里检索到两条记录一条是激进方案被质量门禁拦下的完整日志一条是保守方案的收益数据。如果没有结构化的Memory库这些经验大概率会被淹没在日志海里重蹈覆辙。这套Agent体系到现在为止已经在我们内部跑了两三个月累计沉淀了四十多条有效的kernel优化经验和二十多条此路不通的教训。我最大的感受是Agent并没有替代人的判断它把人的经验转化成可重复执行的实验循环让每一次试错都变成可检索的资产。真正值钱的不是某一轮跑出来的性能提升而是那个架构感知-瓶颈定位-动作选择-质量验证-经验沉淀的闭环本身。如果你也在做diffusion模型的serving优化我建议不要一开始就追求全自动先把profile、验证、记录这三个环节自动化再逐步把决策权交给Agent。这样一步步来踩坑的成本会低很多。