TL;DR: We’ve built an agentic kernel development framework that identifies model-level optimization opportunities, generates improved kernels, and validates them in our serving stack. On our current models, we’ve improved end-to-end latency by 42.3% on Qwen-Image, 15.2% on FLUX.2 and 5.5% increase in tok/s on MiniMax M3.
We’ve seen in recent years that agents have become surprisingly capable at kernel development, from ideation to generating kernels from scratch. Existing benchmarks such as KernelBench have made it easier to evaluate how well agents can optimize kernels on isolated general-purpose problems.
However, there’s a gap between winning a kernel benchmark and shipping optimizations into production.
A few reasons why:
The best kernel configuration depends on the production workload. The kernel that wins on a general benchmark may lose on a specific deployment. Optimizations such as tile shapes, warp-specialization strategy and CTA configurations respond differently to changes in tensor shape, batch size, sequence length, etc. Kernels like MoE and Attention make this especially visible.
A faster microbenchmark doesn’t necessarily translate to a faster model. Once your changes are integrated, interactions with downstream dependencies like CUDA graph capture and multi-stream execution can wipe out kernel-level gains or even result in a regression.
Optimizing kernels individually can miss higher-level opportunities. End-to-end traces often show that only a small subset of kernels have headroom for improvement. The lower-effort wins may come from restructuring the computation around them: fusing operations, eliminating redundant work, or removing pipeline bubbles.
Integrating a new kernel into a production serving engine is nontrivial. Unlike modifying a standalone torch model, serving engines have interconnected execution paths and dependencies. New kernels must be wired into the correct path, replace existing computation cleanly, and remain compatible with the surrounding runtime.
With this in mind, we’ve created a solution that bridges the gap between benchmarks and production. Given a model and serving engine, our framework is able to profile the full workload, reason about the best optimizations, then generate and ship those kernels straight to production.
The stack
The optimization stack divides into two layers:
Architecture diagram
Model-Level Optimization: Understands the full model workload, profiles where time is spent, and proposes changes such as fusion and redundant work elimination.
Per-Kernel Optimization: Takes generated and other performance-critical kernels identified in the trace, explores several implementations in parallel, and iterates on the strongest candidate.
The first layer helps expand the search space beyond one-for-one kernel improvements. Rather than only optimizing kernels in isolation, the framework can restructure the execution graph by removing redundant work, reducing intermediate materialization, or combining operations before generating and improving the underlying kernels.
Learning across optimization runs
Our framework also has a self-improving mechanism: kernels that pass correctness and end-to-end performance checks are retained as reusable candidates., while lessons from both successful and failed attempts are added to an evolving knowledge base alongside workload constraints and integration findings.
This creates a self-improvement loop where each optimization iteration starts from accumulated experience, enabling the agent to generate stronger candidates and converge faster over time.
Persistent knowledge architecture
Results + case studies
Our initial experiment targeted diffusion models, namely Qwen-Image and FLUX.2 served with SGLang on B300 GPUs. The optimizations highlighted below were identified, proposed, and implemented entirely by our agentic framework.
Improvement between baseline, model level and kernel level optimizations
Optimizations on both models
Optimization #1 - Prepacked FP8 Scales
The FP8 paths in Qwen-Image and FLUX.2 were wasting launches converting scale metadata into DeepGEMM’s required format before matrix multiplications. Constant weight scales were repeatedly repacked through sequences of small kernel launches.
We eliminate this overhead by changing the main FP8 activation producers to emit packed scales directly while also moving weight-scale packing to model load time. The numerical computation is unchanged, so outputs remain bit-identical.
For example, at FLUX.2 attention projections:
Baseline
Optimized
Another example at Qwen-Image feed-forward layer:
Baseline
Optimized
The optimization reduced end-to-end latency by 7.3% on Qwen-Image and 6.1% on FLUX.2, with these gains persisting throughout the subsequent FP8 optimizations.
Optimization #2 - Fused QKV projection and epilogue
Both models’ original attention paths compute the image query, key, and value projections independently, despite them all using the same input. This resulted in repeated activation quantization and GEMM setup throughout every attention block.
The optimization merges the three FP8 projections into one GEMM, then fuses bias addition, QK normalization, RoPE, and writes to the joint image-text attention buffers in a single Triton epilogue. NVFP4 still uses separate Q, K, and V GEMMs as each projection uses a different scale.
Baseline
Optimized
Optimization #3 - Normalization + quantization kernel fusion
In both models, normalization previously produced a large BF16 tensor that the following quantization kernel immediately reads back. Thus, the fix was to simply fuse these together, eliminating the intermediate BF16 write-and-read round trip.
Baseline
Optimized
On Qwen-Image, the fused kernel emits both the original BF16 result and pre-quantized FP8 activations for the QKV and feed-forward GEMMs. This reduces latency by 4.3%, and creates the producer path used by the packed-scale optimization.
On FLUX.2’s residual path, the fused kernel emits the normalized output, updated residual, packed E2M1 values, and swizzled E4M3 scales in one pass. This improves end-to-end latency by 0.7%.
Qwen-Image
Optimization #1 - Bias absorption
After the previous optimization, there are two standalone bias additions remaining after the attention and feed-forward output projections which account for roughly 11% of Qwen-Image's FP8 step time. To account for this, we fold each bias into the next fused operation (residual normalization scale and residual update) reducing latency by 5.2%.
Optimization #2 - CFG modulation cache
Classifier-free guidance runs two denoiser passes at the same timestep. Each pass uses different conditioning (one receives the prompt, while the other receives an empty or negative prompt $\varnothing$) The previous implementation recomputed the same timestep-only image and text modulation branches in both passes:
𝜖
c
o
n
d
=
𝐹
(
𝑥
𝑡
,
𝑡
,
𝑐
)
,
𝜖
u
n
c
o
n
d
=
𝐹
(
𝑥
𝑡
,
𝑡
,
∅
)
𝜖
C
F
G
=
𝜖
u
n
c
o
n
d
+
𝑤
(
𝜖
c
o
n
d
−
𝜖
u
n
c
o
n
d
)
The noisy latent and timestep are shared by both passes. The image and text modulation branches are functions only of the timestep embedding and fixed model parameters, not the prompt:
𝑒
𝑡
=
embed
(
𝑡
)
and so:
𝑚
i
m
a
g
e
=
𝑊
i
m
a
g
e
𝑒
𝑡
+
𝑏
i
m
a
g
e
,
𝑚
t
e
x
t
=
𝑊
t
e
x
t
𝑒
𝑡
+
𝑏
t
e
x
t
Because these modulation branches depend only on $e_t$ and fixed weights, their outputs are identical across the conditional and unconditional passes at the same timestep, making it cacheable. Prompt-dependent outputs like hidden states and attention are computed separately.
Create cache key at DiT entry (the same timestep object is passed to both CFG branches):
Cache the image and text modulation outputs in each block
This contributes to a reduction in latency of 2.1% for FP8 and 3.1% for NVFP4.
Optimization #3 - Per-kernel optimization
We then run an optimization pass on performance-critical and previously fused kernels, producing the following improvements:
Together, these per-kernel optimizations have a latency improvement of 7.6% for FP8 and 13.4% for NVFP4.
FLUX.2
Optimization #1 - Single-Block QK normalization + RoPE
FLUX.2’s single-stream transformer block didn’t use the production fused QK-normalization and RoPE kernel because of a Python contiguity guard that rejected the merged-GEMM views. The fallback ran QK RMSNorm and interleaved RoPE as separate passes which repeatedly concatenated the cosine and sine caches.
The new replacement is a per-token-CTA kernel that loads each contiguous 12 KB Q/K head tile, performs RMSNorm in FP32, rounds the result to BF16, and applies interleaved RoPE in the same pass. It reads the cosine and sine tensors directly, eliminating 48 of 60 cache concatenations per step.
Baseline
Optimized
The fused kernel offers a 2× speedup, resulting in end-to-end latency improvements of 2.3% for FP8 and 4.0% for NVFP4.
Optimization #2 - Fused SwiGLU + FP8/NVFP4 quantization
Each invocation of SwiGLU previously produced a large BF16 intermediate that a separate FP8 or NVFP4 quantization kernel read for the output projection. Some single blocks also launched another operation to join attention features with the SwiGLU output.
This optimization replaces those multi-stage paths with a single fused kernel that performs the aforementioned steps in one pass:
Baseline
Optimized
For FP8, matching production exactly requires preserving the original operation order: compute SiLU using division, round to BF16, multiply in BF16, and derive the FP8 scale from the stored BF16 result.
For NVFP4, the same fused path directly emits the packed E2M1 values and swizzled E4M3 scales required by the downstream FP4 GEMM. This eliminates the intermediate BF16 write-and-read round trip, contiguous copy, and multiple standalone kernel launches.
The fused kernel reduces latency by 2.3% for FP8 and 3.8% for NVFP4.
Optimization #3 - Gated Residual Normalization
FLUX.2’s residual path previously ran the gated residual update and layer normalization as two separate operations. The previous production stack didn’t support FLUX.2’s gate, leaving the model on the unfused path.
Baseline
Optimized
The new kernel fuses the gate multiplication, residual update, normalization, and scale/shift into one operation, reducing latency by 1.2% for FP8 and 2.3% for NVFP4.
Optimization #4 - Per-Kernel Optimization
Similar to Qwen-Image, we run another per-kernel optimization loop that identifies the following improvements:
Together, these kernels reduced NVFP4 latency by 2.8% and FP8 latency by 1.9%.
Future Direction
The framework is designed to be model and engine-agnostic, allowing the same optimization loop to be applied across various serving stacks. We have already begun expanding into LLM optimization, where kernel implementations are considerably more mature and leave less headroom for improvement. Despite this, early results show up to 5.5% tok/s improvements on models such as MiniMax M3 and GLM-5.2 on VLLM (Stay tuned!)
As the harness and production integration continue to improve, we see a path toward automatically generating kernels that are specialized for the workloads that actually matter given a specific model, hardware platform, tensor shapes, and serving patterns. Rather than relying solely on generic kernels, each deployment could continuously evolve toward the implementation best suited to its real traffic.
TL;DR:我们构建了一个智能体内核开发框架(agentic kernel development framework),能够识别模型级别的优化机会,生成改进的内核,并在我们的服务栈中验证它们。在当前的模型上,我们将 Qwen-Image 的端到端延迟降低了 42.3%,FLUX.2 降低了 15.2%,MiniMax M3 的 tok/s 提升了 5.5%。
近年来我们看到,智能体在内核开发方面展现出了惊人的能力,从构思到从零生成内核。现有的基准测试(如 KernelBench)使得评估智能体在独立通用问题上优化内核的能力变得更加容易。
然而,在内核基准测试中获胜与将优化部署到生产环境之间仍存在差距。
原因如下:
最佳的内核配置取决于生产工作负载。在通用基准测试中获胜的内核,在特定部署中可能会表现不佳。诸如 tile 形状、warp 专用化策略(warp-specialization strategy)和 CTA 配置等优化,对张量形状、批量大小(batch size)、序列长度等变化有不同的响应。像 MoE 和 Attention 这样的内核使这一点尤为明显。
更快的微基准测试(microbenchmark)不一定意味着更快的模型。一旦你的更改被集成,与下游依赖项(如 CUDA graph 捕获和多流执行)的交互可能会抹杀内核级别的增益,甚至导致性能回退。
单独优化内核可能会错过更高层次的机会。端到端追踪(traces)通常显示,只有一小部分内核有改进空间。低成本的收益可能来自于围绕它们重构计算:融合操作、消除冗余工作,或移除流水线气泡(pipeline bubbles)。
将新内核集成到生产服务引擎中并非易事。与修改独立的 torch 模型不同,服务引擎具有相互关联的执行路径和依赖关系。新内核必须被正确接入到相应的路径中,干净地替换现有计算,并与周围的运行时保持兼容。
考虑到这些,我们创建了一个弥合基准测试与生产环境之间差距的解决方案。给定一个模型和服务引擎,我们的框架能够对完整工作负载进行性能分析(profile),推理出最佳优化方案,然后生成并将这些内核直接部署到生产环境中。
技术栈
优化栈分为两层:
架构图
模型级优化(Model-Level Optimization):理解完整的模型工作负载,分析时间花费在哪里,并提出诸如融合和消除冗余工作等改进建议。
逐内核优化(Per-Kernel Optimization):对追踪中识别出的已生成内核和其他性能关键内核进行处理,并行探索多种实现方案,并在最强候选方案上迭代。
第一层有助于将搜索空间扩展到一对一的内核改进之外。框架不再仅孤立地优化内核,而是可以通过移除冗余工作、减少中间物化(intermediate materialization)或在生成和改进底层内核之前合并操作来重构执行图。
跨优化运行的学习
我们的框架还具备自我改进机制:通过正确性和端到端性能检查的内核将被保留为可复用候选方案,而从成功和失败尝试中获得的经验教训将与工作负载约束和集成发现一起被添加到不断演进的知识库中。
这创建了一个自我改进循环,每次优化迭代都从积累的经验开始,使智能体能够生成更强的候选方案并随时间推移更快地收敛。
持久化知识架构
结果与案例研究
我们最初的实验针对扩散模型(diffusion models),即在 B300 GPU 上通过 SGLang 服务的 Qwen-Image 和 FLUX.2。以下突出显示的优化完全由我们的智能体框架识别、提出并实现。
基线、模型级和内核级优化之间的改进
两个模型上的优化
优化 #1 - 预打包 FP8 缩放因子(Prepacked FP8 Scales)
Qwen-Image 和 FLUX.2 中的 FP8 路径在矩阵乘法之前浪费了启动开销,将缩放因子元数据转换为 DeepGEMM 所需的格式。常量权重缩放因子通过一系列小型内核启动被反复重新打包。
我们通过将主要 FP8 激活生产者改为直接输出打包的缩放因子,同时将权重缩放因子打包移至模型加载时间来消除这一开销。数值计算保持不变,因此输出保持位级一致(bit-identical)。
例如,在 FLUX.2 的注意力投影中:
基线
优化后
另一个例子在 Qwen-Image 的前馈层:
基线
优化后
该优化将 Qwen-Image 的端到端延迟降低了 7.3%,FLUX.2 降低了 6.1%,这些增益在后续的 FP8 优化中持续保持。
优化 #2 - 融合 QKV 投影与收尾(epilogue)
两个模型的原始注意力路径独立计算图像的 query、key 和 value 投影,尽管它们都使用相同的输入。这导致在每个注意力块中重复进行激活量化和 GEMM 设置。
该优化将三个 FP8 投影合并为一个 GEMM,然后在单个 Triton 收尾(epilogue)中融合偏置加法、QK 归一化、RoPE 以及对联合图文注意力缓冲区的写入。NVFP4 仍然使用独立的 Q、K 和 V GEMM,因为每个投影使用不同的缩放因子。
基线
优化后
优化 #3 - 归一化 + 量化内核融合
在两个模型中,归一化之前会生成一个大型 BF16 张量,随后的量化内核立即将其读回。因此,修复方法是将这两者融合在一起,消除中间的 BF16 写入和读取往返。
基线
优化后
在 Qwen-Image 上,融合内核同时输出原始 BF16 结果和用于 QKV 及前馈 GEMM 的预量化 FP8 激活。这将延迟降低了 4.3%,并创建了打包缩放因子优化所使用的生产者路径。
在 FLUX.2 的残差路径上,融合内核在一次传递中输出归一化结果、更新的残差、打包的 E2M1 值和 swizzled 的 E4M3 缩放因子。这将端到端延迟改善了 0.7%。
Qwen-Image
优化 #1 - 偏置吸收(Bias absorption)
在前一个优化之后,注意力和前馈输出投影之后仍有两个独立的偏置加法,约占 Qwen-Image FP8 步骤时间的 11%。为了解决这个问题,我们将每个偏置折叠到下一个融合操作(残差归一化缩放和残差更新)中,将延迟降低了 5.2%。
优化 #2 - CFG 调制缓存(CFG modulation cache)
无分类器引导(Classifier-free guidance)在同一时间步运行两次去噪器传递。每次传递使用不同的条件(一个接收提示词,另一个接收空或负提示词 $\varnothing$)。之前的实现在两次传递中重新计算了相同的时间步专用图像和文本调制分支:
$$\epsilon_{cond} = F(x_t, t, c), \quad \epsilon_{uncond} = F(x_t, t, \varnothing)$$
$$\epsilon_{CFG} = \epsilon_{uncond} + w(\epsilon_{cond} - \epsilon_{uncond})$$
噪声潜变量和时间步在两次传递间共享。图像和文本调制分支仅是时间步嵌入和固定模型参数的函数,而非提示词的函数:
$$e_t = \operatorname{embed}(t)$$
因此:
$$m_{image} = W_{image} e_t + b_{image}, \quad m_{text} = W_{text} e_t + b_{text}$$
由于这些调制分支仅依赖于 $e_t$ 和固定权重,它们的输出在相同时间步的条件和无条件传递中是相同的,因此可以被缓存。依赖提示词的输出(如隐藏状态和注意力)则分别计算。
在 DiT 入口处创建缓存键(相同的时间步对象被传递给两个 CFG 分支):
在每个块中缓存图像和文本调制输出
这使 FP8 的延迟降低了 2.1%,NVFP4 降低了 3.1%。
优化 #3 - 逐内核优化
然后我们对性能关键和先前融合的内核运行优化传递,产生以下改进:
这些逐内核优化合计使 FP8 延迟改善了 7.6%,NVFP4 改善了 13.4%。
FLUX.2
优化 #1 - 单块 QK 归一化 + RoPE
FLUX.2 的单流 transformer 块没有使用生产环境的融合 QK 归一化和 RoPE 内核,因为一个 Python 连续性检查(contiguity guard)拒绝了合并 GEMM 的视图。回退方案将 QK RMSNorm 和交错 RoPE 作为独立传递运行,反复拼接余弦和正弦缓存。
新的替代方案是一个逐 token-CTA 内核,加载每个连续的 12 KB Q/K 头 tile,在 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 值和 swizzled 的 E4M3 缩放因子。这消除了中间的 BF16 写入和读取往返、连续拷贝以及多个独立内核启动。
融合内核将 FP8 延迟降低了 2.3%,NVFP4 降低了 3.8%。
优化 #3 - 门控残差归一化(Gated Residual Normalization)
FLUX.2 的残差路径之前将门控残差更新和层归一化作为两个独立操作运行。之前的生产栈不支持 FLUX.2 的门控,使模型留在未融合的路径上。
基线
优化后
新内核将门控乘法、残差更新、归一化和缩放/偏移融合为一个操作,将 FP8 延迟降低了 1.2%,NVFP4 降低了 2.3%。
优化 #4 - 逐内核优化
与 Qwen-Image 类似,我们运行了另一个逐内核优化循环,识别出以下改进:
这些内核合计将 NVFP4 延迟降低了 2.8%,FP8 延迟降低了 1.9%。
未来方向
该框架设计为与模型和引擎无关,允许将相同的优化循环应用于各种服务栈。我们已经开始向 LLM 优化扩展,其中内核实现已经相当成熟,改进空间较小。尽管如此,早期结果显示在 MiniMax M3 和 GLM-5.2 等模型上,在 VLLM 上 tok/s 提升高达 5.5%(敬请期待!)
随着工具链和生产集成的持续改进,我们看到了一条通向自动生成内核的道路——这些内核专门针对给定特定模型、硬件平台、张量形状和服务模式下真正重要的工作负载进行优化。每个部署不再仅依赖通用内核,而是可以持续向最适合其真实流量的实现演进。