Source-linked AI summary
KernelSkill: A Multi-Agent Framework for GPU Kernel Optimization
Qitong Sun, Jun Han, Tianlin Li, Zhe Tang, Sheng Chen, Fei Yang, Aishan Liu, Xianglong Liu, Yang Liu
TL;DR
GPU kernel optimization needs more interpretable and efficient strategy selection than implicit LLM heuristics provide. KernelSkill addresses this with a multi-agent framework using reusable long-term expert skills and short-term trajectory memory. On KernelBench Levels 1–3, it achieves 100% success and average speedups of 5.44×, 2.82×, and 1.92× over Torch Eager, respectively.
Problem
Existing LLM-based kernel optimizers often rely on implicit method-selection heuristics, while noisy feedback, shifting bottlenecks, and structural preconditions make reusable, traceable decisions difficult.
Method
KernelSkill coordinates multiple agents with structured long-term optimization knowledge and short-term per-task trajectories for evidence-grounded, iterative kernel refinement.
Results
100% success and average speedups of 5.44×, 2.82×, and 1.92× over Torch Eager were achieved on KernelBench Levels 1–3, respectively.
Takeaways & Limitations
Explicit expert skills and trajectory memory make optimization-method selection more auditable and refinement more stable across kernel tasks.
Takeaways & Limitations
Performance depends on long-term-memory coverage, and baseline comparisons are not uniformly reproduced because some systems are unavailable or lack reported KernelBench results.
Abstract
from arXiv · showhide
Improving GPU kernel efficiency is crucial for advancing AI systems. Recent work has explored leveraging large language models (LLMs) for GPU kernel generation and optimization. However, existing LLM-based kernel optimization pipelines typically rely on opaque, implicitly learned heuristics within the LLMs to determine optimization strategies. This leads to inefficient trial-and-error and weakly interpretable optimizations. Our key insight is to replace implicit heuristics with expert optimization skills that are knowledge-driven and aware of task trajectories. Specifically, we present KernelSkill, a multi-agent framework with a dual-level memory architecture. KernelSkill operates by coordinating agents with long-term memory of reusable expert skills and short-term memory to prevent repetitive backtracking. On KernelBench Levels 1-3, KernelSkill achieves a 100% success rate and average speedups of 5.44x, 2.82x, and 1.92x over Torch Eager on Levels 1, 2, and 3, respectively, outperforming prior baselines. Code is available at https://github.com/0satan0/KernelMem/.
1. Introduction
KernelSkill addresses the cost and opacity of GPU kernel optimization by replacing implicit LLM heuristics with reusable, trajectory-aware expert skills. Its multi-agent system combines long-term knowledge and short-term task memory, achieving strong KernelBench results.
- GPU kernel optimization remains costly because it requires architecture-aware coordination, profiling-guided diagnosis, targeted edits, and repeated validation.
- KernelSkill replaces implicit model-internal heuristics with knowledge-driven, trajectory-aware optimization skills.
- The framework externalizes expert knowledge for traceable method selection and maintains explicit optimization state to stabilize profiling-guided refinement.
- Its architecture coordinates generation, verification, profiling, planning, and repair in a closed-loop optimizer.
- Two-level memory combines reusable long-term optimization skills with short-term per-task optimization and repair history.
- 100% success and average speedups of 5.44×, 2.82×, and 1.92× over Torch Eager were achieved on KernelBench Levels 1–3, respectively.
2. Related Works
Prior GPU kernel optimization relies on expert libraries, compiler/DSL ecosystems, or LLM-based training and closed-loop refinement. These approaches remain costly, difficult to generalize, or dependent on implicit learned guidance.
- Expert-maintained libraries and compiler/DSL ecosystems can be effective but require substantial engineering effort and continual hardware- and workload-specific upkeep.
- Training-based LLM methods adapt specialized models or learn hierarchical policies for iterative kernel improvement.
- Training-centric pipelines can struggle to generalize across large search spaces and have difficulty incorporating per-kernel runtime feedback.
- Inference-time approaches refine kernels through compilation checks, correctness tests, and profiling feedback across kernel, DSL, evolutionary, and end-to-end settings.
3. Motivation
The paper identifies implicit optimization-method selection and missing trajectory tracking as central weaknesses in LLM-based kernel optimization. KernelSkill responds by making expertise auditable and preserving per-task state across refinement rounds.
- Implicit method selection is error-prone because profiling signals are noisy, bottlenecks shift, and optimization actions depend jointly on runtime evidence and code structure.
- A motivating GEMM example achieved only 0.032× speedup when refinement repeatedly pursued fusion instead of the primary efficiency bottleneck.
- The challenge is accurate, traceable method selection that identifies the dominant bottleneck and avoids repeatedly spending iterations on low-yield directions.
- KernelSkill externalizes expert knowledge into retrievable, auditable decisions grounded in profiling evidence and maintains explicit per-task optimization state.
- Long-term and short-term memories together reduce oscillations, support coupled edits, and enable larger speedups with fewer refinement rounds under profiling feedback.
4. KernelSkill
KernelSkill is a closed-loop, multi-agent optimizer that combines expert knowledge retrieval with per-task trajectory memory. It generates, verifies, diagnoses, repairs, plans, and profiles kernels while using static features and runtime evidence to guide method selection.
- KernelSkill overview: KernelSkill integrates a multi-agent refinement pipeline with long-term optimization knowledge and short-term per-task trajectories.
- KernelSkill overview: The system includes Generator, Reviewer, Feature Extractor, Retrieval, Planner, Optimizer, Diagnoser, Repairer, and two memory modules.
- Refinement pipeline: Starting from PyTorch, the Generator creates seed kernels and the Reviewer evaluates them before closed-loop refinement begins.
- Refinement pipeline: Compilation or correctness failures trigger diagnosis and repair, whereas successful kernels enter feature-based retrieval and planning for optimization.
- Refinement pipeline: The loop runs for at most N rounds and returns the best-performing kernel.
- Agent roles: The Generator prioritizes correctness and creates diverse operator-level seed kernels for later refinement.
- Feature extraction: Static code features capture structural optimization signals that profiling alone may weakly reflect or leave ambiguous.
- Feature extraction: Feature extraction uses 18 fine-grained feature types retrieved through rule-based matching and LLM-based inference, although the set is expandable.
5. Experiments
KernelSkill is evaluated on KernelBench Levels 1–3 using success, speedup, and Fast1 metrics against prior training-based and agentic baselines. It achieves perfect success and the highest average speedup across all levels, while ablations show that both memory components improve reliability and optimization.
- Benchmark and metrics: KernelBench Levels 1–3 comprise 250 tasks spanning single-kernel operators, multioperator workloads, and full model architectures.The evaluation follows prior work by reporting results on Levels 1–3 rather than Level 4.
- Experimental results: 100% success across all three levels makes KernelSkill one of the methods achieving reliable end-to-end kernel synthesis.STARK is the only compared baseline also matching 100% success on every level.
- Experimental results: 5.44×, 2.82×, and 1.92× average speedup over Torch Eager are the best results on Levels 1, 2, and 3, respectively.KernelSkill improves over STARK by 79.5%, 4.8%, and 21.5% on Levels 1, 2, and 3, respectively.
- Experimental results: KernelSkill outperforms profiling-grounded multi-agent baselines, reaching 2.82× versus CudaForge’s 2.10× on Level 2 and 1.92× versus 1.28× on Level 3.On Level 1, it reaches 5.44×, compared with approximately 1.45–1.49× for Astra, PRAGMA, and CudaForge.
- Experimental results: KernelSkill maintains perfect success and the highest speedups across levels, whereas training-centric baselines degrade sharply on harder tasks.QiMeng falls from 1.00 success and 2.20× speedup on Level 1 to 0.70 success and 0.73× speedup on Level 3.
- Experimental results: 15 refinement rounds yield higher average speedups and per-round gains than STARK after 30 rounds, with gains attributed largely to cross-task long-term memory.KernelSkill’s per-round gains are 0.36/0.19/0.13 across Levels 1/2/3, versus STARK’s 0.10/0.09/0.05.
- Ablation: 96%, 98%, and 94% success without short-term memory rises to 100% across Levels 1–3 when short-term memory is enabled.Short-term memory tracks per-task attempts, prevents repeated failures, and supports coupled multi-step fixes.
- Ablation: Long-term memory produces larger average-speedup gains by narrowing the method search space and reducing trial-and-error.Removing long-term memory yields only limited performance improvement relative to the memory-free baseline.
6. Conclusion
KernelSkill is a memory-augmented multi-agent framework for GPU kernel optimization that achieves full success on KernelBench Levels 1–3. It reports progressively lower average speedups across the three levels and depends on long-term memory coverage.
- KernelSkill achieves a 100% success rate on KernelBench Levels 1–3.
- 5.44×, 2.82×, and 1.92× average speedups over Torch Eager are reported for Levels 1, 2, and 3, respectively.
- Performance depends on long-term memory coverage, with LLM-only evidence-based method selection used when no matching case is retrieved.
A. Algorithm
Algorithm 1 provides the pseudocode for KernelSkill.
- Algorithm 1 contains the pseudocode for KernelSkill.
- The algorithm is presented as a procedural description of KernelSkill.
- The pseudocode serves as the formal algorithmic reference for KernelSkill.
B. Long-Term Memory Schema.
KernelSkill’s long-term memory organizes profiling, runtime, code, and derived features into structured decision fields. It also stores rules, bottleneck mappings, candidate methods, and implementation guidance for traceable optimization selection.
- Feature schema: Field mapping, run features, code features, and derived fields standardize profiling, runtime, structural, and composite kernel information.
- Decision control: Headroom tiers and bottleneck priority rules quantify remaining optimization potential and resolve competing bottleneck diagnoses.
- Decision control: NCU predicates and global forbidden rules provide reusable profiling conditions and veto constraints for unsafe or invalid optimizations.
- Method selection: The decision table maps bottleneck types, headroom tiers, and kernel characteristics to candidate optimization methods.
- Method selection: LLM assistance supplies method-specific rationales and implementation cues for methods selected by the deterministic policy.
C. Long-Term Memory Decision Workflow.
The decision workflow transforms kernel measurements and code characteristics into permitted optimization methods. Deterministic normalization, feature derivation, bottleneck matching, case gating, and global veto rules precede LLM-assisted planning.
- Input processing: NCU metrics, runtime features, and code features are aggregated as the input for the current kernel.
- Input processing: Raw NCU metric keys are normalized into standardized fields, and composite indicators are deterministically computed for optimization decisions.
- Diagnosis and matching: Optimization headroom is assigned to discrete tiers based on performance indicators.
- Diagnosis and matching: The workflow identifies bottlenecks from profiling signatures and matches decision cases using headroom, kernel properties, and gating predicates.
- Filtering and planning: Global veto constraints filter unsafe or invalid optimizations before the pipeline returns the permitted methods for the matched case.
- Filtering and planning: LLM-assisted planning adds non-binding explanations, implementation guidance, and expected benefits for the selected methods.
D. Case
The case implements a fused matrix-multiplication kernel that scales and clamps its output, with supporting PyTorch and CUDA code defining the task and launch configuration.
- Provided code: The paper separately provides the task code and an optimized implementation without KernelSkill.These are identified as Algorithms 2 and 3.
- Reference model: The PyTorch model applies a linear transformation, scaling, residual addition, clamping, LogSumExp, and Mish activation.The documented input and output shapes are (batch_size, input_size) and (batch_size, hidden_size).
- CUDA implementation: The task is implemented as a fused matrix-multiplication and scale-clamp CUDA kernel.The kernel receives X, W, B, output, dimensions, scale, and clamp parameters.
- CUDA implementation: Each thread computes a row-column output element when its indices fall within the batch and output-feature bounds.Thread coordinates derive from block and thread indices, followed by bounds checking.
- CUDA implementation: The kernel accumulates products over input features, adds the bias, and writes the result to the output tensor.The output is stored using row-major indexing over batch and output-feature dimensions.
- Kernel launch: The extension allocates an output tensor and launches the fused kernel with a 16×16 thread block and dimension-based grid.The launch passes tensor pointers, dimensions, scaling, and clamp parameters to the kernel.
E. Retrieval method and generate optimization plan.
Figure 4 presents the retrieval method and the generation of an optimization plan.
- Retrieval: Figure 4 depicts the retrieval method.
- Planning: Figure 4 depicts optimization-plan generation.
- Overview: The figure combines retrieval and plan generation in one pipeline view.