用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 Cache，PD分离把prefill和decode的负载拆开调度，连续batch动态拼装请求。这些优化在token逐字生成的任务里收益巨大，因为注意力计算在decode阶段是memory-bound，KV 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 Guidance（CFG），意味着每个采样步都跑两遍网络，一遍带条件一遍不带条件。两遍网络在早期层的输入几乎一样，区别主要来自文本条件参与注入的位置。如果在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、融合norm+linear、启用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改成了稀疏索引格式，重新跑benchmark，attention耗时从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或者访存密集的gemm；FLUX.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 PCIe，CUDA 12.4，PyTorch 2.4，SGLang基于社区最新release分支，Agent运行在独立的一台CPU节点上，通过SSH下发任务。测试集是10条中文prompt加10条英文prompt，图像模型固定1024x1024分辨率，视频模型固定16帧512x512。

各项优化合入后的实测结果如下表，所有数据都是warmup 3次之后正式采样20次，取p50耗时和峰值显存。

这里要强调一点：这些数据是基于我们内部环境、我们选定的这版模型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_LIST="7.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调到了256，benchmark显示p50提升了6%，看起来是一个干净的胜利。但Verifier在质量检查时发现CLIP score从0.318下降到了0.302，再一看生成图，出现了明显的条纹伪影。

排查后的根因是：block size的变化改变了attention内部数值累加的顺序和精度，这种细微的数值差异在LLM任务里可能无所谓，但在图像/视频生成任务里会被迭代式去噪放大成可见伪影。生成模型对数值路径的敏感度远超判别式模型，这是我们在优化LLM时不太容易形成的直觉。

这次事故之后，我们在Verifier里加入了三重质量检查：结构化相似度指标、CLIP score、人工抽检。并且定了一条硬性规则：任何会改变数值路径的优化项，比如调block size、切换attention变体、改精度策略，必须同时通过质量门禁才能合入。这条规则后来写进了Memory，Agent没有再犯过同样的错误。

5.5 踩坑四：Agent日志的可读性设计

最后一个偏工程的坑。Agent跑久了之后，日志量会变得非常大，如果日志格式只适合人读，检索效率极低；如果只适合机器检索，人review的时候又很痛苦。我们的解决方案是双轨制：每个experiment生成一个JSONL格式的结构化记录，方便SQLite导入和后续检索；同时每个关键决策点会生成一段自然语言的summary，方便人类在晚上快速浏览当天Agent干了什么。

这个设计在一次排查线上问题时帮了大忙：我们想查"Agent之前是否在视频模型的CFG分支共享上做过尝试"，直接在SQLite里检索到两条记录，一条是激进方案被质量门禁拦下的完整日志，一条是保守方案的收益数据。如果没有结构化的Memory库，这些经验大概率会被淹没在日志海里，重蹈覆辙。

这套Agent体系到现在为止已经在我们内部跑了两三个月，累计沉淀了四十多条有效的kernel优化经验和二十多条"此路不通"的教训。我最大的感受是：Agent并没有替代人的判断，它把人的经验转化成可重复执行的实验循环，让每一次试错都变成可检索的资产。真正值钱的不是某一轮跑出来的性能提升，而是那个"架构感知->瓶颈定位->动作选择->质量验证->经验沉淀"的闭环本身。如果你也在做diffusion模型的serving优化，我建议不要一开始就追求全自动，先把profile、验证、记录这三个环节自动化，再逐步把决策权交给Agent。这样一步步来，踩坑的成本会低很多。