生产环境中的智能内核
构建智能体内核开发框架，通过模型级与逐内核两层优化，将 Qwen-Image 端到端延迟提升 42.3%、FLUX.2 提升 15.2%、MiniMax M3
TL;DR: 我们构建了一个智能体内核开发框架，能够识别模型层面的优化机会、生成改进后的内核，并在我们的服务栈中进行验证。针对当前模型，我们已将 Qwen-Image 的端到端延迟提升了 42.3%，FLUX.2 提升了 15.2%，MiniMax M3 的每秒处理 token 数（tok/s）提升了 5.5%。
近年来，我们看到智能体在内核开发方面变得出人意料地强大，从构思到从零生成内核皆能胜任。诸如 KernelBench 之类的现有基准测试，使得评估智能体在孤立通用问题上优化内核的能力变得更加容易。
然而，在赢得内核基准测试与将优化成果投入生产之间，仍存在一道鸿沟。
原因如下：
- 最佳内核配置取决于生产工作负载。 在通用基准测试中胜出的内核，可能在特定部署场景下表现不佳。诸如 tile 形状、warp 专用化策略和 CTA 配置等优化手段，对张量形状、批大小、序列长度等变化会有不同的响应。MoE 和 Attention 这类内核尤其能体现这一点。
- 更快的微基准测试并不一定意味着更快的模型。 一旦你的改动被集成，与下游依赖项（如 CUDA graph 捕获和多流执行）的交互可能会抵消内核层面的收益，甚至导致性能回退。
- 单独优化内核可能会错失更高层面的机会。 端到端的追踪结果往往显示，只有一小部分内核存在改进空间。更省力的收益可能来自围绕这些内核重构计算：融合操作、消除冗余工作，或移除流水线气泡。
- 将新内核集成到生产级服务引擎中并非易事。 与修改独立的torch模型不同，服务引擎具有相互关联的执行路径和依赖关系。新内核必须接入正确的路径，干净地替换现有计算，并与周围的运行时保持兼容。
基于此，我们创建了一个解决方案，弥合了基准测试与生产环境之间的差距。给定模型和服务引擎，我们的框架能够分析完整工作负载，推理出最佳优化方案，然后生成并将这些内核直接部署到生产环境中。
技术栈
优化栈分为两层：
- 模型级优化： 理解完整模型工作负载，分析时间消耗所在，并提出融合和冗余工作消除等更改建议。
- 逐内核优化： 获取生成的及跟踪中识别出的性能关键内核，并行探索多种实现方案，并对最强候选进行迭代优化。
第一层有助于将搜索空间扩展到一对一的核改进之外。框架不仅限于孤立地优化内核，还可以通过消除冗余工作、减少中间物化或合并操作来重构执行图，然后再生成和改进底层内核。
跨优化运行的学习机制
我们的框架还具备自我改进机制：通过正确性与端到端性能检查的内核会被保留为可复用候选，同时，无论成功还是失败尝试中的经验教训，都会连同工作负载约束与集成发现，一并纳入不断演进的知识库。
由此形成一个自我提升的闭环，每次优化迭代都从累积的经验出发，使智能体能够生成更强的候选方案，并随着时间推移更快地收敛。
结果与案例研究
我们的初步实验聚焦于扩散模型，即由 SGLang 在 B300 GPU 上服务的 Qwen-Image 与 FLUX.2。以下重点展示的优化方案，完全由我们的智能体框架识别、提出并实施。
两个模型的优化
优化 #1 - 预打包 FP8 缩放因子
Qwen-Image 和 FLUX.2 中的 FP8 路径在矩阵乘法前，浪费了大量内核启动时间，用于将缩放元数据转换为 DeepGEMM 所需的格式。常量权重缩放因子通过一系列小型内核启动被反复重新打包。
我们通过让主要的 FP8 激活生成器直接输出打包后的缩放因子，并将权重缩放因子的打包移至模型加载时，从而消除了这一开销。数值计算保持不变，因此输出结果仍与之前完全一致。
例如，在 FLUX.2 的注意力投影中：
另一个例子是在 Qwen-Image 的前馈层中：
这一优化使 Qwen-Image 的端到端延迟降低了 7.3%，FLUX.2 降低了 6.1%，并且这些收益在后续的 FP8 优化中得以保持。
优化 #2 - 融合的 QKV 投影与尾声处理
两个模型的原始注意力路径分别计算图像查询、键和值投影，尽管它们都使用相同的输入。这导致在每个注意力块中重复进行激活量化和 GEMM 设置。
该优化将三个 FP8 投影合并为一个 GEMM，然后在单个 Triton 尾声处理中融合偏置加法、QK 归一化、RoPE，并写入联合图像-文本注意力缓冲区。NVFP4 仍使用独立的 Q、K、V GEMM，因为每个投影使用不同的缩放因子。
| 模型 | FP8 延迟 | NVFP4 延迟 |
|---|---|---|
| Qwen-Image | 降低 15.8% | 降低 1.3% |
| FLUX.2 | 降低 1.3% | 降低 1.1% |
优化 #3 - 归一化 + 量化内核融合
在两个模型中，归一化此前会生成一个大型 BF16 张量，紧随其后的量化内核会立即读回该张量。因此，修复方法就是简单地将两者融合，消除中间 BF16 的写入与读取往返开销。
在 Qwen-Image 上，融合内核同时输出原始 BF16 结果以及用于 QKV 和前馈 GEMM 的预量化 FP8 激活。这使延迟降低了 4.3%，并创建了打包缩放优化所使用的生产者路径。
在 FLUX.2 的残差路径上，融合内核在一次遍历中输出归一化结果、更新后的残差、打包的 E2M1 值以及交错排列的 E4M3 缩放因子。这使端到端延迟改善了 0.7%。
Qwen-Image
优化 #1 - 偏置吸收
经过上述优化后，在注意力与前馈输出投影之后，仍有两个独立的偏置加法操作，它们约占 Qwen-Image FP8 步骤时间的 11%。为了解决这个问题，我们将每个偏置折叠到下一个融合操作（残差归一化缩放和残差更新）中，使延迟降低了 5.2%。
优化 #2 - CFG 调制缓存
无分类器引导在同一时间步运行两次去噪过程。每次过程使用不同的条件（一次接收提示词，另一次接收空提示或负提示 \varnothing）。此前的实现在两次过程中都重新计算了仅依赖时间步的图像和文本调制分支：
噪声潜变量和时间步长在两次前向传播中是共享的。图像和文本调制分支仅依赖于时间步长嵌入和固定的模型参数，而非提示词：
因此：
由于这些调制分支仅依赖于 e_t 和固定权重，它们在相同时间步长下的条件与无条件前向传播中输出完全一致，因此可以缓存。而依赖提示词的输出，如隐藏状态和注意力，则分别计算。
在DiT入口处创建缓存键（相同的时间步长对象传递给两个CFG分支）：
def _cfg_cache_optimization_enabled(active) -> bool:
return active
# QwenImageTransformer2DModel.forward
if _cfg_cache_optimization_enabled() and isinstance(timestep, torch.Tensor):
# Both CFG branches receive the same timestep tensor.
# Keep a reference to it so tensor identity can be used safely.
cache_key = {
"timestep": timestep,
"version": version_if_available(timestep),
}在每个块中缓存图像和文本调制输出。
# QwenImageTransformerBlock.forward
cached = getattr(self, "modulation_cache", None)
cache_hit = (
cache_key is not None
and cached is not None
and cached["timestep"] is cache_key["timestep"]
and cached["version"] == cache_key["version"]
)
if cache_hit:
# Second CFG pass: reuse the cached outputs.
image_modulation = cached["image_modulation"]
text_modulation = cached["text_modulation"]
else:
# First CFG pass: compute the modulation outputs.
image_modulation = image_modulation_GEMM(timestep_embedding)
text_modulation = text_modulation_GEMM(timestep_embedding)
# Cache them for the second CFG pass.
if cache_key is not None:
self.modulation_cache = {
"timestep": cache_key["timestep"],
"version": cache_key["version"],
"image_modulation": image_modulation,
"text_modulation": text_modulation,
}这有助于将延迟降低 FP8 为 2.1%，NVFP4 为 3.1%。
优化 #3 - 每内核优化
随后，我们对性能关键且先前已融合的内核执行优化过程，产生以下改进：
| 原始内核 | 功能描述 | 新内核 |
|---|---|---|
swiglu_fp4_quant | 将 SwiGLU 与 NVFP4 量化融合，适配生产级形状 — 优化 #2 内核 | 1.43× / 1.41× |
resnorm_quant | 融合残差归一化与量化，加速 FP8 残差归一化链 — 优化 #3 内核 | 1.65× |
qknorm_rope | 融合双块 QK 归一化与 RoPE — 优化 #1 内核 | 2.0× |
token_cat | 加速剩余的图像-文本拼接 — 优化 #2 内核 | 4.6× |
norm_out | 融合边界层归一化、缩放乘法与残差加法 — 优化 #3 内核 | 3.35× |
swiglu | 优化剩余未量化的 SwiGLU 位置 — 优化 #2 内核 | 1.24× |
gate_res_norm | 优化剩余的门控残差归一化路径 — 优化 #3 内核 | 1.27× |
综合来看，这些逐内核优化在 FP8 上实现了 7.6% 的延迟改进，在 NVFP4 上实现了 13.4% 的延迟改进。
FLUX.2
优化 #1 - 单块 QK 归一化 + RoPE
FLUX.2 的单流 Transformer 块未使用生产级融合 QK 归一化和 RoPE 内核，原因是 Python 连续性检查拒绝了合并 GEMM 视图。回退方案将 QK RMSNorm 和交错 RoPE 作为独立步骤运行，反复拼接余弦和正弦缓存。
新的替代方案是一个按 token 分配的 CTA 内核，它加载每个连续的 12 KB Q/K 头块，在 FP32 中执行 RMSNorm，将结果四舍五入为 BF16，并在同一遍中应用交错 RoPE。它直接读取余弦和正弦张量，每步消除了 60 次缓存拼接中的 48 次。
融合内核提供了 2 倍加速，端到端延迟改善分别为 FP8 的 2.3% 和 NVFP4 的 4.0%。
优化 #2 - 融合 SwiGLU + FP8/NVFP4 量化
每次调用 SwiGLU 之前都会产生一个大型 BF16 中间结果，由单独的 FP8 或 NVFP4 量化内核读取以用于输出投影。某些单块还会启动另一个操作，将注意力特征与 SwiGLU 输出合并。
此优化将多阶段路径替换为单个融合内核，该内核在一遍中执行上述步骤：
对于 FP8，要精确匹配生产环境，必须保留原始操作顺序：使用除法计算 SiLU，四舍五入为 BF16，在 BF16 中相乘，并从存储的 BF16 结果中推导 FP8 缩放因子。
对于 NVFP4，相同的融合路径直接输出下游 FP4 GEMM 所需的打包 E2M1 值和交错 E4M3 缩放因子。这消除了中间 BF16 写入和读取的往返、连续复制以及多次独立内核启动。
融合内核将 FP8 延迟降低 2.3%，NVFP4 延迟降低 3.8%。
优化 #3 - 门控残差归一化
FLUX.2 的残差路径此前将门控残差更新和层归一化作为两个独立操作执行。之前的生成堆栈不支持 FLUX.2 的门控，导致模型走的是未融合路径。
新内核将门控乘法、残差更新、归一化以及缩放/平移融合为单一操作，使 FP8 延迟降低 1.2%，NVFP4 延迟降低 2.3%。
优化 #4 - 逐内核优化
与 Qwen-Image 类似，我们运行了另一轮逐内核优化循环，识别出以下改进：
| 原始内核 | 功能 | 新内核 |
|---|---|---|
swiglu_fp4_quant | 将 SwiGLU 与 NVFP4 量化融合，适配生产对齐形状 — 优化 #2 内核 | 1.43× / 1.41× |
resnorm_quant | 将残差归一化与量化融合，并加速 FP8 残差归一化链 — 优化 #3 内核 | 1.65× |
qknorm_rope | 融合双块 QK 归一化与 RoPE — 优化 #1 内核 | 2.0× |
token_cat | 加速剩余的图像-文本拼接 — 优化 #2 内核 | 4.6× |
norm_out | 融合边界层归一化、缩放乘法与残差加法 — 优化 #3 内核 | 3.35× |
swiglu | 优化剩余未量化的 SwiGLU 位置 — 优化 #2 内核 | 1.24× |
gate_res_norm | 优化剩余的门控残差归一化路径 — 优化 #3 内核 | 1.27× |
这些内核共同使 NVFP4 延迟降低 2.8%，FP8 延迟降低 1.9%。
未来方向
该框架设计为模型与引擎无关，使得同一优化循环能够应用于各种服务栈。我们已开始拓展至LLM优化领域，其中内核实现已相当成熟，改进空间相对有限。尽管如此，早期结果显示，在VLLM上针对MiniMax M3和GLM-5.2等模型，吞吐量提升可达5.5%（敬请期待！）
随着测试平台与生产集成的持续改进，我们预见一条路径，能够针对特定模型、硬件平台、张量形状及服务模式，自动生成专门化的内核。不再单纯依赖通用内核，每次部署都能持续演进，趋向于最适合其实际流量特征的实现方案。