1. 为什么盯上SGLangDiffusion推理的调度痛点不是算力不够是机制不对先交代下背景。最近大半年我一直在折腾图像和视频生成模型的线上推理服务从基础的SD系列到最新的Qwen-Image、FLUX.2还有SVD、Wan系列这类视频模型跑了一圈下来有个非常直观的感受同样的模型在不同推理框架上的表现差距能大到让你怀疑是不是卡坏了。尤其是当请求并发上来、输入分辨率五花八门的时候vLLM和SGLang、以及一些自研的diffusers流水线之间差距不只在吞吐上还在延迟的稳定性上。vLLM在LLM场景下确实做得够好这点无可争议。但它对Diffusion模型的支持至少在很长一段时间里都停留在“能跑”的阶段——能加载模型、能生成图但调度策略基本就是FCFS先来先服务这一套对Diffusion这种“按步数迭代、中间状态有依赖”的推理模式没有做太多针对性的优化。更关键的是vLLM的continuous batching在LLM里是靠着KV Cache和prefill/decode解耦来实现的但Diffusion模型压根没有KV Cache这种天然可以做memory reuse的结构每一步都是完整的前向传播。这就导致vLLM在做Diffusion推理时所谓的continuous batching很大程度退化成了一种“伪并行”——请求按step对齐窗口分组组内并行组间串行。SGLang在这块是完全不同的思路。它的调度器本质上是把Diffusion推理当成一个有向无环图DAG任务来管理的每一个denoising step都显式建模成图中的一个节点节点之间用显式的依赖关系连接。这意味着SGLang可以做到真正的跨请求step级并行——请求A的第5步和请求B的第3步只要计算资源有空闲就可以同时执行不需要强制对齐step窗口。这个能力在实际场景里的收益是很可观的。为什么这么讲先说一个数据。我在A10080G上用Qwen-Image跑1024x1024、50步的推理单请求裸跑大约需要8到9秒。但如果同一时刻进来10个请求用SGLang的调度器和原生diffusers的Pipeline并行跑SGLang能把末尾请求的完成时间提前30%到40%。原因不复杂Qwen-Image前几步step 0到10做的是全局结构建模计算量大、每步耗时高后几步step 40到50基本是细节微调算力需求明显下降。原生Pipeline是傻乎乎地等所有请求都到了同一个step再往前推慢的请求就拖累快的请求。SGLang是哪个请求的依赖满足了解锁了就先跑哪个。这种差异在小并发下不明显并发一上来就完全是两种体验。所以我对SGLang的判断很明确它就是当前Diffusion模型推理里最值得深入研究的开源框架没有之一。本文后面要聊的所有东西都是基于SGLang的性能剖析框架特别是它的autotune机制和kernel优化接口来做的配合Agent来自动化整个调优闭环。2. Agent在这里面到底扮演什么角色不是写代码的是替你无聊的很多朋友一听“Agent优化kernel”第一反应是让AI自动写CUDA算子然后自动替换、自动验证。这个方向现在确实有团队在做但我个人的观点是在Diffusion推理这个场景下Agent最大的价值不是帮你发明新kernel而是帮你把“特征分析—瓶颈定位—候选kernel生成—benchmark验证—参数归档”这个循环自动化。说白了Agent是个不知疲倦的调参师和测试员它做的筛选和验证工作人工做不仅慢而且容易漏。我们团队搭的这套Agent系统本质上是一个多智能体协作的闭环。核心角色有四个Profling Agent剖析智能体负责加载模型和输入样本跑一次完整的推理流程然后采集性能数据。这里采集的不只是step耗时还包括每个op的GPU kernel耗时分布、显存占用峰值、kernel launch overhead、以及GPU util的波动曲线。Analysis Agent分析智能体读入剖析数据结合模型结构比如Transformer block的层数、FFN的隐藏维度、attention head数等来定位热点。它不是一个简单的“谁耗时长谁就是瓶颈”的判断器而是会做相关性分析——比如一个kernel耗时高但它的计算密度arithmetic intensity也很高那它可能是受制于带宽而不是受制于算力这时候盲上算子融合是没用的。Optimization Agent优化智能体根据Analysis Agent的输出去SGLang的kernel仓库、或者我们积累的自定义kernel库中挑选/生成候选项。这里有一个很关键的策略它会优先尝试改动最小、风险最低的方案比如调整cache策略、修改autotune配置、替换某一个qkv投影的kernel实现而不会一上来就动attention计算的核心结构。Validator Agent验证智能体这可能是整个链路里最容易被人忽略但最重要的一个。它不只是验证优化后的输出和优化前是否一致还会验证数值精度是否在可接受范围内Diffusion用的是fp16/bf16你替换fusion算法之后误差会积累得设置容差、显存是否真的降了、以及在同样并发条件下延迟和吞吐是否有真实提升。这套体系的运作方式是Optimization Agent生成的每个候选优化都会被Validator Agent自动评估评估通过的才写进配置归档评估不通过的会带着失败原因返回给Analysis Agent重新分析。整条链路无人值守运行我们只需要在最初设定好目标函数比如“P95延迟降20%显存占用不超过40G生成质量PSNR不低于原始实现的0.99倍”然后等它跑完输出报告。这里我必须强调一个容易踩的坑做Agent优化千万不要让它直接修改模型权重或替换模型结构。Agent的优化边界应该是“计算图变换”和“kernel选择”不要越界去改模型的语义。Diffusion模型的训练代价太高了一个经过大量蒸馏、量化、剪枝的模型你动它的结构哪怕只动一个residual connection生成质量都可能断崖式下跌。我们内部定过规矩任何优化只要让CLIP similarity或者PSNR下降超过1%直接一票否决不管推理速度提升了多少。后面第三、四、五章我会分别讲Qwen-Image、FLUX.2和视频模型这三类模型在Kernel优化上的具体实践以及Agent在这个过程中发挥的作用。这三类模型的架构差异极大优化策略也完全不同。3. Qwen-Image的Kernel优化MoE和Text Encoder才是真正的两只拦路虎先说Qwen-Image。它是阿里通义实验室出的MoE架构的DiTDiffusion Transformer模型在SGLang上跑起来之后我做的第一件事就是让Profling Agent跑一次完整的性能剖析把耗时分布拉出来。数据出来之后我人都愣了一下——按我的预想最大的瓶颈应该集中在UNet/DiT的主干网络特别是那几个大尺寸的Transformer block上。但实际数据告诉我在1024x1024输入下主干网络耗时占总耗时的比例不到60%。那另外40%去哪了两个地方。一个是文本编码器Text EncoderQwen-Image的文本编码器是Qwen2-VL-7B这种规模的模型对prompt做编码要跑几百个token的前向传播这个过程非常吃显存带宽和算子调度。另一个就是MoE的专家路由——Qwen-Image的MoE层不是所有专家都被激活的每层会通过router选择top-k个专家但这个路由决策本身是动态的、依赖当前step的hidden state的这导致SGLang的静态图优化很难对MoE部分做预编译kernel launch的开销被放得很大。先说文本编码器。Qwen2-VL这种规模的模型做文本编码计算量其实不大几千亿次的FLOPs对A100来说就是一瞬间的事真正拖后腿的是它那恐怖的kernel launch数量和内存访问模式。Text Encoder里有大量的小张量操作比如LayerNorm、GELU、reshape、attention score的mask处理这些op单个耗时都是微秒级别的但架不住数量多——跑一遍编码要触发几百次kernel launch。每次launch都有固定开销大概5到10微秒的空闲时间GPU从接到指令到真正执行之间的延迟几百次叠加起来就是几毫秒甚至十几毫秒的纯开销。这个问题的解法方向其实很明确算子融合operator fusion。把LayerNorm QKV投影 attention score计算 SoftMax这串操作融合成一个或者两个kernel把中间结果的写回显存、再读出来的过程省掉。SGLang本身对Transformer类的模型已经有不错的fusion支持但对Qwen2-VL这种新架构融合覆盖并不完全。我们的做法是用Agent去识别出文本编码器前向计算图里所有“小算子簇”然后自动匹配SGLang/torch.compile里已有的fusion模板。这里有个很讨巧的点文本编码器对生成质量的影响是间接的它只是把prompt编码成embedding所以我们对文本编码器做更激进的数值精度降低比如从fp16降到bf16甚至fp8对图像质量的影响微乎其微但速度提升却是实打实的。再来说MoE专家的路由。Qwen-Image的MoE层用的是类似Mixtral的机制每层有N个专家但每个token只激活top-2。这个机制引入的问题是专家权重矩阵无法被预先绑定到特定的计算流上因为每个step每个token选哪些专家是运行时才决定的。在SGLang的静态图优化框架下这就意味着专家计算的kernel无法像非MoE层那样被提前编译好只能在运行时动态dispatch。每个token都要做一次router前向计算、选专家、再从显存里加载对应的专家权重这个过程的开销在batch size比较小的时候甚至能占单层耗时的一半。对这个问题的优化思路我们做的是“专家权重的分组预加载”。原理很简单虽然每个token的动态路由是运行时的但在同一个batch内部被选中的专家集合往往是高度重叠的——也就是大家选的专家就集中在那么几个。我们通过分析1000个真实prompt在50个denoise step中的路由分布发现top-6专家被选中的概率超过90%。于是我们做了一个激进的cache策略在推理开始前把最热门的8个专家的权重预先加载进GPU的L2 cache和shared memory里而不是等运行时再慢慢拉。这个策略让MoE层的平均耗时降了大约22%。可能有人要问那直接用torch.compile不是也能做融合吗为什么不直接用我的回答是torch.compile的融合更多是在Python层和Triton层做的它对计算图的改写能力很强但对显存规划memory planning的控制很弱。对于Qwen-Image这种动辄几十GB的MoE模型显存布局一旦不合适L2 cache命中率会掉得很惨。SGLang的kernel优化接口允许你更精细地控制显存的分配和复用策略这才是真正的优势所在。总结一下Qwen-Image的优化要点我整理成一张表瓶颈模块根因优化手段实测收益文本编码器小算子数量多、kernel launch overhead大LayerNorm投影Attention算子融合单次full推理耗时降约15%MoE路由动态dispatch导致kernel无法预编译热门专家权重L2 Cache预加载MoE层平均耗时降约22%主干DiT大张量算子带宽利用率不足使用SGLang自带的autotune机制做GEMM阻塞参数搜索主干部分降约10%4. FLUX.2的Kernel实践T5双塔和Parallel Block是被低估的优化富矿如果说Qwen-Image的优化难点在于MoE的调度开销那FLUX.2的难点就在于架构本身的复杂性和对显存的贪婪。FLUX.2Black Forest Labs出品的FLUX系列我这边用的是FLUX.2-dev版本是典型的双塔文本编码器架构——同时使用T5-XXL和CLIP-L两个文本编码器而且主干网络是并行block设计每个block同时处理time embedding和文本embedding的条件注入。这种设计带来的问题是它对显存的占用比单塔模型高出一截而且T5-XXL这个大哥本身就是一个11B参数规模的巨无霸。先聊T5。在Qwen-Image里我说文本编码器是个小瓶颈在FLUX.2里T5就直接是个大瓶颈了。T5-XXL做一次完整的长文本编码512个token在A100上需要大约0.8秒而整个图像生成流程也才8到10秒。也就是说T5占了将近10%的总耗时这个占比高得吓人。而T5的问题和Qwen-Image的文本编码器还不完全一样——T5是一个完整的encoder-decoder架构它的计算图极大层数达到了24层encoder 24层decoder而且每一层中间还有很多残差连接和门控机制。这种结构对编译器来说非常不友好因为它做算子融合的时候需要跨越多层做依赖分析分析时间本身就超过了kernel的优化收益。我们试过几种方案。第一种是直接在编码完文本后把T5的fp16权重强制降到fp8来跑。效果很好耗时直接降了约40%。但代价是当你用FLUX.2跑特别长的、包含专业术语的prompt时生成的图像里会出现轻微的语义漂移比如把“red vintage car”里的“vintage”给丢掉了。后来我们做了个折中——前12层encoder保持fp16后12层和所有decoder层用fp8。前12层主要做的是基础的词法和句法分析对细节敏感不能降精度后12层和decoder做的是语义整合和条件嵌入生成对精度不那么敏感。这个方案的最终效果是T5耗时降了约35%但生成的图像和纯fp16跑出来的几乎一致。再来说并行Block。FLUX.2的并行block结构在整个Diffusion模型里算是比较特殊的。它的每个DiT block内部有两条并行的计算路径——一条是处理文本embedding的cross-attention路径另一条是处理图像tokens的self-attention路径。这两条路径在block的末端汇合做一次特征融合然后进入下一个block。这种结构的好处是信息流动更充分坏处是在kernel优化层面并行路径意味着显存带宽的竞争。因为两条路径要同时读不同的数据文本embedding和图像tokens如果缓存策略设计得不好L2 cache的命中率会直线下降。针对这个问题我们用了两个技巧。第一个是路径分离的kernel融合策略。SGLang默认会把block内所有kernel按依赖顺序做融合但这样往往会把并行路径的操作串行化。我们通过SGLang的compute graph重写接口把并行路径上的算子分别融合成两个独立的fusion group让它们在GPU上真正并发执行。这个改动对单batch的耗时影响不大但对并发场景非常有用——当多个请求同时通过同一个block时并行路径的并发执行能让GPU的utilization从70%左右升到90%以上。第二个技巧更有意思——**对time embedding的处理。**FLUX.2的time embedding层是一个非常小的MLP输入是一个标量timestep输出是一个高维向量通常是3072维或更高。就是这么一个小MLP在每个denoise step里都要被重新计算一次。我们统计过这个tiny MLP在单次推理50个step里触发了超过1000次的kernel launch。按照每一步平均3到5次launch来算光是time embedding就占了整个推理的10%到15%的kernel launch开销。这个优化很机械就是把这个MLP的输出做缓存——因为同一个请求的50个step的timestep是固定的那么time embedding的输出就只跟timestep有关可以预先算好保存下来后面每个step直接读取不用再计算。这个改动看起来小但在Agent的自动优化流程里被标记为“低成本高收益候选”实际部署后单请求延迟降了约8%。很香。FLUX.2这块另外还有一个值得注意的点就是它的显存占用模型。FLUX.2的DiT模型本身就有差不多12B参数加上T5-XXL和CLIP-L总共接近30B。这意味着在A100 80G上跑一个小batch都非常紧张。我们试过用Agent自动搜索最佳的前置KV Cache逐出策略cache eviction policy让在不同batch大小下自动决定哪些中间激活值可以提前释放以腾出显存给新的请求。这一步优化做完我们能够稳定地在A100上同时跑4个FLUX.2请求而不会OOM单卡吞吐提升了将近3倍。这个结果对线上服务来说意义是巨大的。5. 视频模型的时间维度为什么SVD和Wan的Kernel优化要完全换一套打法视频模型我这边主要跑的是Stable Video DiffusionSVD和Wan系列。说实话视频模型的推理优化比图像模型整整复杂一个数量级核心原因就四个字时间维度。图像模型是2D的——你有一张NxN的latent做50步denoise每一步计算的是NxN的分布。视频模型是3D的——你有一个NxNxF的latentF是帧数做同样50步的denoise每一步都要处理F个帧的联合分布。这就导致视频模型有两个天生的负担一是单步计算量是图像模型的F倍二是跨帧的attention计算temporal attention引入了额外的显存消耗和计算复杂度。先聊时间维度的attention问题。视频模型的DiT在每一层里除了常规的空间attention对每帧内部的token做注意力还要做时间attention对不同帧对应位置的token做注意力。时间attention的张量形状是[batch, frame, height*width, height*width]在大分辨率视频下这个矩阵的尺寸会膨胀到让人绝望——512x512、30帧的视频height*width4096时间attention矩阵就是4096x4096接近1700万个元素光这个矩阵就要占34MB的显存fp16在一个层里它还不止一个头而且时间attention要做N层。这个矩阵的显存占用是个巨大的坑因为它的生命周期太短——算完之后马上要被softmax和V矩阵乘法消耗掉但在这个过程中你必须把它完整地放在显存里。对于A100的40GB显存来说一个层吃掉几十MB内存不算什么但模型有大量层数而且并行处理4个视频请求的话情况就完全不一样了——直接在显存层面OOM。我们通过分析发现时间attention矩阵是视频模型显存OOM的头号贡献者占比超过总显存峰值的45%。对这个问题的优化做了一个非常经典的kernel级优化操作——online softmax重计算。传统实现是先把attention score矩阵完整算出来、存在显存里、然后做softmax、再乘V。online softmax的做法是不保存中间attention score矩阵而是用一个running max和running sum来增量式计算softmax每算出一行score就立刻乘V并把结果累加进输出。这样我们就把峰值显存占用从O(N²)降到了O(N)。对512x512、30帧这个配置光这一项优化就让显存峰值下降了差不多20GB。代价是实现复杂度比较高因为在SGLang的计算图上你需要把原来的attention子图完全拆开重写。这块Agent发挥的作用很大——它在搜索kernel时会自动去匹配这个“attention重计算”的模式并生成对应的Triton实现。再聊视频模型在SGLang上的调度问题。视频模型比图像模型的单请求计算时间长得多——一个5秒、30帧的视频生成一次可能要2分钟以上。这意味着经典的队列调度策略FCFS在视频场景下会引发极其严重的尾部延迟——如果队列前面排了一个大视频请求后面即使是一个只需要1秒的小图像请求也得排队等2分钟后才能跑。对线上服务来说这完全不可接受。Agent优化系统发现了一个更好的模式我把这个模式称为“时间步分片调度Timestep-Chunked Scheduling”。简单来说就是把视频模型50个denoise step划分为多个chunk比如每个chunk包含10步每个chunk作为一个独立的调度单元。当一个视频请求的chunk执行完调度器会暂停它检查队列里是否有更小更紧急的请求如果有就插队执行小请求然后回到视频请求继续跑下一个chunk。这个策略让视频请求的完成时间稍微长了一点因为有上下文切换开销大概增加2%但小请求的P99延迟直接降了70%以上。这个思路在LLM里很常见类似preemption但在Diffusion视频推理中很少有人提。对视频模型来说其实还有一个比较有前景的方向就是利用时间维度的冗余来做推理优化。视频相邻帧之间的内容高度相似这导致时间attention在很多位置上算出来的值是近似重复的。我们尝试过一个思路——在推理的前期step0到15步每两帧共享一部分时间attention的计算结果也就是把30帧看成15组每组内的2帧复用同一份attention矩阵。在实验中发现前期步骤下这种近似对最终视频质量的影响几乎无法用肉眼察觉但速度提升在低step数下很明显。6. 那些Agent踩过的坑和我事后复盘出来的教训前面讲的都是Agent帮我找到的优化点但搞Agent优化这件事并不是一帆风顺的。这一章我要专门说说这套系统在实践过程中踩过的坑以及我自己事后复盘出来的一些判断。这些经验比那些优化参数本身更值钱。第一个大坑是**“过度调优”导致的过拟合**。用Agent做autotune它会疯狂地在性能数据里找优化空间有些优化在固定的benchmark上效果好到惊人但一换输入就全废。最典型的例子是一个Qwen-Image的GEMM阻塞参数搜索——Agent在1024x1024输入下找到了一个最优的block size和num_stages组合单次推理快了18%。但当我们把输入分辨率改成2048x2048之后这个组合的性能直接倒退了15%。原因不复杂不同分辨率下矩阵乘法的维度差异决定了最适合的block size不同。所以Agent优化时必须跑多组不同输入形态的泛化验证不能只盯着一个benchmark跑。我们后来在Agent系统里加了一个“泛化测试集”包含至少5种不同的分辨率、3种不同的batch大小任何优化必须在这个测试集上全部通过才能归档。这个改动让很多看起来很美、实际不一定能用的优化项目被自动淘汰掉了。第二个坑是在线服务里显存优化不等于延迟优化。Agent系统在离线benchmark上发现某个kernel优化能让显存占用降25%觉得赚翻了——因为降显存意味着可以开更大的并发。但上了线之后P99延迟反而变高了。查了半天才发现这个优化是通过减少L2 cache的使用来降低显存占用的但这导致同一个数据被反复从HBM显存主存储器读取而HBM的带宽远低于L2 cache。结果就是显存占用的确降了但计算时间因为cache miss变多了。这是一个很典型的“A/B指标倒挂”案例。我给大家的建议是不要让Agent只看一个目标指标必须给它建立一个多指标联合约束。我们后来给Agent设了三个硬性门槛——显存峰值不超过某个阈值、P95延迟不能比优化前差、吞吐必须提升超过5%三个门槛同时满足才接受这个优化。第三个坑也最有价值是关于数值误差积累的判断。Diffusion模型是迭代结构每步的输出是下一步的输入这意味着哪怕每个step只引入了浮点噪声经过50步的迭代后噪声也会被指数级放大。我们最开始用Agent做bf16替换的时候每一层的输出误差对比fp16是几乎可以忽略的相对误差在0.1%以内但跑到第50步之后生成的图像和fp16版本之间出现了明显的色偏。后来分析原因是bf16的尾数精度太低导致某些残差连接处的信息被抹掉了。这是个教训对Diffusion模型做任何精度降低不能只看单step的误差必须看完整50步之后的最终输出误差。后来我把Agent的验证环节改了——所有优化必须在推理全部step完成后对最终生成结果做CLIP similarity和SSIM双指标对比任何一个指标低于阈值就自动打回。复盘到现在我对Agent做Diffusion Kernel优化这件事的整体判断是它现在更像是一个极其勤快的助手而不是一个敢拍板的决策者。它能一天跑几百组实验把你从枯燥的benchmark循环里解放出来但真正决定优化方向正确与否的还是你对模型架构的理解和对业务目标的把握。拿最后这句话当个结尾吧别把Agent当成能独立思考的工程师把它当成一个可以无限加班的实习生——它真的能替你打很多杂但关键的判断还是得你自己来做。
1. 为什么盯上SGLangDiffusion推理的调度痛点不是算力不够是机制不对先交代下背景。最近大半年我一直在折腾图像和视频生成模型的线上推理服务从基础的SD系列到最新的Qwen-Image、FLUX.2还有SVD、Wan系列这类视频模型跑了一圈下来有个非常直观的感受同样的模型在不同推理框架上的表现差距能大到让你怀疑是不是卡坏了。尤其是当请求并发上来、输入分辨率五花八门的时候vLLM和SGLang、以及一些自研的diffusers流水线之间差距不只在吞吐上还在延迟的稳定性上。vLLM在LLM场景下确实做得够好这点无可争议。但它对Diffusion模型的支持至少在很长一段时间里都停留在“能跑”的阶段——能加载模型、能生成图但调度策略基本就是FCFS先来先服务这一套对Diffusion这种“按步数迭代、中间状态有依赖”的推理模式没有做太多针对性的优化。更关键的是vLLM的continuous batching在LLM里是靠着KV Cache和prefill/decode解耦来实现的但Diffusion模型压根没有KV Cache这种天然可以做memory reuse的结构每一步都是完整的前向传播。这就导致vLLM在做Diffusion推理时所谓的continuous batching很大程度退化成了一种“伪并行”——请求按step对齐窗口分组组内并行组间串行。SGLang在这块是完全不同的思路。它的调度器本质上是把Diffusion推理当成一个有向无环图DAG任务来管理的每一个denoising step都显式建模成图中的一个节点节点之间用显式的依赖关系连接。这意味着SGLang可以做到真正的跨请求step级并行——请求A的第5步和请求B的第3步只要计算资源有空闲就可以同时执行不需要强制对齐step窗口。这个能力在实际场景里的收益是很可观的。为什么这么讲先说一个数据。我在A10080G上用Qwen-Image跑1024x1024、50步的推理单请求裸跑大约需要8到9秒。但如果同一时刻进来10个请求用SGLang的调度器和原生diffusers的Pipeline并行跑SGLang能把末尾请求的完成时间提前30%到40%。原因不复杂Qwen-Image前几步step 0到10做的是全局结构建模计算量大、每步耗时高后几步step 40到50基本是细节微调算力需求明显下降。原生Pipeline是傻乎乎地等所有请求都到了同一个step再往前推慢的请求就拖累快的请求。SGLang是哪个请求的依赖满足了解锁了就先跑哪个。这种差异在小并发下不明显并发一上来就完全是两种体验。所以我对SGLang的判断很明确它就是当前Diffusion模型推理里最值得深入研究的开源框架没有之一。本文后面要聊的所有东西都是基于SGLang的性能剖析框架特别是它的autotune机制和kernel优化接口来做的配合Agent来自动化整个调优闭环。2. Agent在这里面到底扮演什么角色不是写代码的是替你无聊的很多朋友一听“Agent优化kernel”第一反应是让AI自动写CUDA算子然后自动替换、自动验证。这个方向现在确实有团队在做但我个人的观点是在Diffusion推理这个场景下Agent最大的价值不是帮你发明新kernel而是帮你把“特征分析—瓶颈定位—候选kernel生成—benchmark验证—参数归档”这个循环自动化。说白了Agent是个不知疲倦的调参师和测试员它做的筛选和验证工作人工做不仅慢而且容易漏。我们团队搭的这套Agent系统本质上是一个多智能体协作的闭环。核心角色有四个Profling Agent剖析智能体负责加载模型和输入样本跑一次完整的推理流程然后采集性能数据。这里采集的不只是step耗时还包括每个op的GPU kernel耗时分布、显存占用峰值、kernel launch overhead、以及GPU util的波动曲线。Analysis Agent分析智能体读入剖析数据结合模型结构比如Transformer block的层数、FFN的隐藏维度、attention head数等来定位热点。它不是一个简单的“谁耗时长谁就是瓶颈”的判断器而是会做相关性分析——比如一个kernel耗时高但它的计算密度arithmetic intensity也很高那它可能是受制于带宽而不是受制于算力这时候盲上算子融合是没用的。Optimization Agent优化智能体根据Analysis Agent的输出去SGLang的kernel仓库、或者我们积累的自定义kernel库中挑选/生成候选项。这里有一个很关键的策略它会优先尝试改动最小、风险最低的方案比如调整cache策略、修改autotune配置、替换某一个qkv投影的kernel实现而不会一上来就动attention计算的核心结构。Validator Agent验证智能体这可能是整个链路里最容易被人忽略但最重要的一个。它不只是验证优化后的输出和优化前是否一致还会验证数值精度是否在可接受范围内Diffusion用的是fp16/bf16你替换fusion算法之后误差会积累得设置容差、显存是否真的降了、以及在同样并发条件下延迟和吞吐是否有真实提升。这套体系的运作方式是Optimization Agent生成的每个候选优化都会被Validator Agent自动评估评估通过的才写进配置归档评估不通过的会带着失败原因返回给Analysis Agent重新分析。整条链路无人值守运行我们只需要在最初设定好目标函数比如“P95延迟降20%显存占用不超过40G生成质量PSNR不低于原始实现的0.99倍”然后等它跑完输出报告。这里我必须强调一个容易踩的坑做Agent优化千万不要让它直接修改模型权重或替换模型结构。Agent的优化边界应该是“计算图变换”和“kernel选择”不要越界去改模型的语义。Diffusion模型的训练代价太高了一个经过大量蒸馏、量化、剪枝的模型你动它的结构哪怕只动一个residual connection生成质量都可能断崖式下跌。我们内部定过规矩任何优化只要让CLIP similarity或者PSNR下降超过1%直接一票否决不管推理速度提升了多少。后面第三、四、五章我会分别讲Qwen-Image、FLUX.2和视频模型这三类模型在Kernel优化上的具体实践以及Agent在这个过程中发挥的作用。这三类模型的架构差异极大优化策略也完全不同。3. Qwen-Image的Kernel优化MoE和Text Encoder才是真正的两只拦路虎先说Qwen-Image。它是阿里通义实验室出的MoE架构的DiTDiffusion Transformer模型在SGLang上跑起来之后我做的第一件事就是让Profling Agent跑一次完整的性能剖析把耗时分布拉出来。数据出来之后我人都愣了一下——按我的预想最大的瓶颈应该集中在UNet/DiT的主干网络特别是那几个大尺寸的Transformer block上。但实际数据告诉我在1024x1024输入下主干网络耗时占总耗时的比例不到60%。那另外40%去哪了两个地方。一个是文本编码器Text EncoderQwen-Image的文本编码器是Qwen2-VL-7B这种规模的模型对prompt做编码要跑几百个token的前向传播这个过程非常吃显存带宽和算子调度。另一个就是MoE的专家路由——Qwen-Image的MoE层不是所有专家都被激活的每层会通过router选择top-k个专家但这个路由决策本身是动态的、依赖当前step的hidden state的这导致SGLang的静态图优化很难对MoE部分做预编译kernel launch的开销被放得很大。先说文本编码器。Qwen2-VL这种规模的模型做文本编码计算量其实不大几千亿次的FLOPs对A100来说就是一瞬间的事真正拖后腿的是它那恐怖的kernel launch数量和内存访问模式。Text Encoder里有大量的小张量操作比如LayerNorm、GELU、reshape、attention score的mask处理这些op单个耗时都是微秒级别的但架不住数量多——跑一遍编码要触发几百次kernel launch。每次launch都有固定开销大概5到10微秒的空闲时间GPU从接到指令到真正执行之间的延迟几百次叠加起来就是几毫秒甚至十几毫秒的纯开销。这个问题的解法方向其实很明确算子融合operator fusion。把LayerNorm QKV投影 attention score计算 SoftMax这串操作融合成一个或者两个kernel把中间结果的写回显存、再读出来的过程省掉。SGLang本身对Transformer类的模型已经有不错的fusion支持但对Qwen2-VL这种新架构融合覆盖并不完全。我们的做法是用Agent去识别出文本编码器前向计算图里所有“小算子簇”然后自动匹配SGLang/torch.compile里已有的fusion模板。这里有个很讨巧的点文本编码器对生成质量的影响是间接的它只是把prompt编码成embedding所以我们对文本编码器做更激进的数值精度降低比如从fp16降到bf16甚至fp8对图像质量的影响微乎其微但速度提升却是实打实的。再来说MoE专家的路由。Qwen-Image的MoE层用的是类似Mixtral的机制每层有N个专家但每个token只激活top-2。这个机制引入的问题是专家权重矩阵无法被预先绑定到特定的计算流上因为每个step每个token选哪些专家是运行时才决定的。在SGLang的静态图优化框架下这就意味着专家计算的kernel无法像非MoE层那样被提前编译好只能在运行时动态dispatch。每个token都要做一次router前向计算、选专家、再从显存里加载对应的专家权重这个过程的开销在batch size比较小的时候甚至能占单层耗时的一半。对这个问题的优化思路我们做的是“专家权重的分组预加载”。原理很简单虽然每个token的动态路由是运行时的但在同一个batch内部被选中的专家集合往往是高度重叠的——也就是大家选的专家就集中在那么几个。我们通过分析1000个真实prompt在50个denoise step中的路由分布发现top-6专家被选中的概率超过90%。于是我们做了一个激进的cache策略在推理开始前把最热门的8个专家的权重预先加载进GPU的L2 cache和shared memory里而不是等运行时再慢慢拉。这个策略让MoE层的平均耗时降了大约22%。可能有人要问那直接用torch.compile不是也能做融合吗为什么不直接用我的回答是torch.compile的融合更多是在Python层和Triton层做的它对计算图的改写能力很强但对显存规划memory planning的控制很弱。对于Qwen-Image这种动辄几十GB的MoE模型显存布局一旦不合适L2 cache命中率会掉得很惨。SGLang的kernel优化接口允许你更精细地控制显存的分配和复用策略这才是真正的优势所在。总结一下Qwen-Image的优化要点我整理成一张表瓶颈模块根因优化手段实测收益文本编码器小算子数量多、kernel launch overhead大LayerNorm投影Attention算子融合单次full推理耗时降约15%MoE路由动态dispatch导致kernel无法预编译热门专家权重L2 Cache预加载MoE层平均耗时降约22%主干DiT大张量算子带宽利用率不足使用SGLang自带的autotune机制做GEMM阻塞参数搜索主干部分降约10%4. FLUX.2的Kernel实践T5双塔和Parallel Block是被低估的优化富矿如果说Qwen-Image的优化难点在于MoE的调度开销那FLUX.2的难点就在于架构本身的复杂性和对显存的贪婪。FLUX.2Black Forest Labs出品的FLUX系列我这边用的是FLUX.2-dev版本是典型的双塔文本编码器架构——同时使用T5-XXL和CLIP-L两个文本编码器而且主干网络是并行block设计每个block同时处理time embedding和文本embedding的条件注入。这种设计带来的问题是它对显存的占用比单塔模型高出一截而且T5-XXL这个大哥本身就是一个11B参数规模的巨无霸。先聊T5。在Qwen-Image里我说文本编码器是个小瓶颈在FLUX.2里T5就直接是个大瓶颈了。T5-XXL做一次完整的长文本编码512个token在A100上需要大约0.8秒而整个图像生成流程也才8到10秒。也就是说T5占了将近10%的总耗时这个占比高得吓人。而T5的问题和Qwen-Image的文本编码器还不完全一样——T5是一个完整的encoder-decoder架构它的计算图极大层数达到了24层encoder 24层decoder而且每一层中间还有很多残差连接和门控机制。这种结构对编译器来说非常不友好因为它做算子融合的时候需要跨越多层做依赖分析分析时间本身就超过了kernel的优化收益。我们试过几种方案。第一种是直接在编码完文本后把T5的fp16权重强制降到fp8来跑。效果很好耗时直接降了约40%。但代价是当你用FLUX.2跑特别长的、包含专业术语的prompt时生成的图像里会出现轻微的语义漂移比如把“red vintage car”里的“vintage”给丢掉了。后来我们做了个折中——前12层encoder保持fp16后12层和所有decoder层用fp8。前12层主要做的是基础的词法和句法分析对细节敏感不能降精度后12层和decoder做的是语义整合和条件嵌入生成对精度不那么敏感。这个方案的最终效果是T5耗时降了约35%但生成的图像和纯fp16跑出来的几乎一致。再来说并行Block。FLUX.2的并行block结构在整个Diffusion模型里算是比较特殊的。它的每个DiT block内部有两条并行的计算路径——一条是处理文本embedding的cross-attention路径另一条是处理图像tokens的self-attention路径。这两条路径在block的末端汇合做一次特征融合然后进入下一个block。这种结构的好处是信息流动更充分坏处是在kernel优化层面并行路径意味着显存带宽的竞争。因为两条路径要同时读不同的数据文本embedding和图像tokens如果缓存策略设计得不好L2 cache的命中率会直线下降。针对这个问题我们用了两个技巧。第一个是路径分离的kernel融合策略。SGLang默认会把block内所有kernel按依赖顺序做融合但这样往往会把并行路径的操作串行化。我们通过SGLang的compute graph重写接口把并行路径上的算子分别融合成两个独立的fusion group让它们在GPU上真正并发执行。这个改动对单batch的耗时影响不大但对并发场景非常有用——当多个请求同时通过同一个block时并行路径的并发执行能让GPU的utilization从70%左右升到90%以上。第二个技巧更有意思——**对time embedding的处理。**FLUX.2的time embedding层是一个非常小的MLP输入是一个标量timestep输出是一个高维向量通常是3072维或更高。就是这么一个小MLP在每个denoise step里都要被重新计算一次。我们统计过这个tiny MLP在单次推理50个step里触发了超过1000次的kernel launch。按照每一步平均3到5次launch来算光是time embedding就占了整个推理的10%到15%的kernel launch开销。这个优化很机械就是把这个MLP的输出做缓存——因为同一个请求的50个step的timestep是固定的那么time embedding的输出就只跟timestep有关可以预先算好保存下来后面每个step直接读取不用再计算。这个改动看起来小但在Agent的自动优化流程里被标记为“低成本高收益候选”实际部署后单请求延迟降了约8%。很香。FLUX.2这块另外还有一个值得注意的点就是它的显存占用模型。FLUX.2的DiT模型本身就有差不多12B参数加上T5-XXL和CLIP-L总共接近30B。这意味着在A100 80G上跑一个小batch都非常紧张。我们试过用Agent自动搜索最佳的前置KV Cache逐出策略cache eviction policy让在不同batch大小下自动决定哪些中间激活值可以提前释放以腾出显存给新的请求。这一步优化做完我们能够稳定地在A100上同时跑4个FLUX.2请求而不会OOM单卡吞吐提升了将近3倍。这个结果对线上服务来说意义是巨大的。5. 视频模型的时间维度为什么SVD和Wan的Kernel优化要完全换一套打法视频模型我这边主要跑的是Stable Video DiffusionSVD和Wan系列。说实话视频模型的推理优化比图像模型整整复杂一个数量级核心原因就四个字时间维度。图像模型是2D的——你有一张NxN的latent做50步denoise每一步计算的是NxN的分布。视频模型是3D的——你有一个NxNxF的latentF是帧数做同样50步的denoise每一步都要处理F个帧的联合分布。这就导致视频模型有两个天生的负担一是单步计算量是图像模型的F倍二是跨帧的attention计算temporal attention引入了额外的显存消耗和计算复杂度。先聊时间维度的attention问题。视频模型的DiT在每一层里除了常规的空间attention对每帧内部的token做注意力还要做时间attention对不同帧对应位置的token做注意力。时间attention的张量形状是[batch, frame, height*width, height*width]在大分辨率视频下这个矩阵的尺寸会膨胀到让人绝望——512x512、30帧的视频height*width4096时间attention矩阵就是4096x4096接近1700万个元素光这个矩阵就要占34MB的显存fp16在一个层里它还不止一个头而且时间attention要做N层。这个矩阵的显存占用是个巨大的坑因为它的生命周期太短——算完之后马上要被softmax和V矩阵乘法消耗掉但在这个过程中你必须把它完整地放在显存里。对于A100的40GB显存来说一个层吃掉几十MB内存不算什么但模型有大量层数而且并行处理4个视频请求的话情况就完全不一样了——直接在显存层面OOM。我们通过分析发现时间attention矩阵是视频模型显存OOM的头号贡献者占比超过总显存峰值的45%。对这个问题的优化做了一个非常经典的kernel级优化操作——online softmax重计算。传统实现是先把attention score矩阵完整算出来、存在显存里、然后做softmax、再乘V。online softmax的做法是不保存中间attention score矩阵而是用一个running max和running sum来增量式计算softmax每算出一行score就立刻乘V并把结果累加进输出。这样我们就把峰值显存占用从O(N²)降到了O(N)。对512x512、30帧这个配置光这一项优化就让显存峰值下降了差不多20GB。代价是实现复杂度比较高因为在SGLang的计算图上你需要把原来的attention子图完全拆开重写。这块Agent发挥的作用很大——它在搜索kernel时会自动去匹配这个“attention重计算”的模式并生成对应的Triton实现。再聊视频模型在SGLang上的调度问题。视频模型比图像模型的单请求计算时间长得多——一个5秒、30帧的视频生成一次可能要2分钟以上。这意味着经典的队列调度策略FCFS在视频场景下会引发极其严重的尾部延迟——如果队列前面排了一个大视频请求后面即使是一个只需要1秒的小图像请求也得排队等2分钟后才能跑。对线上服务来说这完全不可接受。Agent优化系统发现了一个更好的模式我把这个模式称为“时间步分片调度Timestep-Chunked Scheduling”。简单来说就是把视频模型50个denoise step划分为多个chunk比如每个chunk包含10步每个chunk作为一个独立的调度单元。当一个视频请求的chunk执行完调度器会暂停它检查队列里是否有更小更紧急的请求如果有就插队执行小请求然后回到视频请求继续跑下一个chunk。这个策略让视频请求的完成时间稍微长了一点因为有上下文切换开销大概增加2%但小请求的P99延迟直接降了70%以上。这个思路在LLM里很常见类似preemption但在Diffusion视频推理中很少有人提。对视频模型来说其实还有一个比较有前景的方向就是利用时间维度的冗余来做推理优化。视频相邻帧之间的内容高度相似这导致时间attention在很多位置上算出来的值是近似重复的。我们尝试过一个思路——在推理的前期step0到15步每两帧共享一部分时间attention的计算结果也就是把30帧看成15组每组内的2帧复用同一份attention矩阵。在实验中发现前期步骤下这种近似对最终视频质量的影响几乎无法用肉眼察觉但速度提升在低step数下很明显。6. 那些Agent踩过的坑和我事后复盘出来的教训前面讲的都是Agent帮我找到的优化点但搞Agent优化这件事并不是一帆风顺的。这一章我要专门说说这套系统在实践过程中踩过的坑以及我自己事后复盘出来的一些判断。这些经验比那些优化参数本身更值钱。第一个大坑是**“过度调优”导致的过拟合**。用Agent做autotune它会疯狂地在性能数据里找优化空间有些优化在固定的benchmark上效果好到惊人但一换输入就全废。最典型的例子是一个Qwen-Image的GEMM阻塞参数搜索——Agent在1024x1024输入下找到了一个最优的block size和num_stages组合单次推理快了18%。但当我们把输入分辨率改成2048x2048之后这个组合的性能直接倒退了15%。原因不复杂不同分辨率下矩阵乘法的维度差异决定了最适合的block size不同。所以Agent优化时必须跑多组不同输入形态的泛化验证不能只盯着一个benchmark跑。我们后来在Agent系统里加了一个“泛化测试集”包含至少5种不同的分辨率、3种不同的batch大小任何优化必须在这个测试集上全部通过才能归档。这个改动让很多看起来很美、实际不一定能用的优化项目被自动淘汰掉了。第二个坑是在线服务里显存优化不等于延迟优化。Agent系统在离线benchmark上发现某个kernel优化能让显存占用降25%觉得赚翻了——因为降显存意味着可以开更大的并发。但上了线之后P99延迟反而变高了。查了半天才发现这个优化是通过减少L2 cache的使用来降低显存占用的但这导致同一个数据被反复从HBM显存主存储器读取而HBM的带宽远低于L2 cache。结果就是显存占用的确降了但计算时间因为cache miss变多了。这是一个很典型的“A/B指标倒挂”案例。我给大家的建议是不要让Agent只看一个目标指标必须给它建立一个多指标联合约束。我们后来给Agent设了三个硬性门槛——显存峰值不超过某个阈值、P95延迟不能比优化前差、吞吐必须提升超过5%三个门槛同时满足才接受这个优化。第三个坑也最有价值是关于数值误差积累的判断。Diffusion模型是迭代结构每步的输出是下一步的输入这意味着哪怕每个step只引入了浮点噪声经过50步的迭代后噪声也会被指数级放大。我们最开始用Agent做bf16替换的时候每一层的输出误差对比fp16是几乎可以忽略的相对误差在0.1%以内但跑到第50步之后生成的图像和fp16版本之间出现了明显的色偏。后来分析原因是bf16的尾数精度太低导致某些残差连接处的信息被抹掉了。这是个教训对Diffusion模型做任何精度降低不能只看单step的误差必须看完整50步之后的最终输出误差。后来我把Agent的验证环节改了——所有优化必须在推理全部step完成后对最终生成结果做CLIP similarity和SSIM双指标对比任何一个指标低于阈值就自动打回。复盘到现在我对Agent做Diffusion Kernel优化这件事的整体判断是它现在更像是一个极其勤快的助手而不是一个敢拍板的决策者。它能一天跑几百组实验把你从枯燥的benchmark循环里解放出来但真正决定优化方向正确与否的还是你对模型架构的理解和对业务目标的把握。拿最后这句话当个结尾吧别把Agent当成能独立思考的工程师把它当成一个可以无限加班的实习生——它真的能替你打很多杂但关键的判断还是得你自己来做。
