Paper deep dive
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
Intelligence
Status: succeeded | Model: google/gemini-3.1-flash-lite-preview | Prompt: intel-v1 | Confidence: 95%
Last extracted: 3/13/2026, 1:04:02 AM
Summary
KernelSkill is a multi-agent framework designed for GPU kernel optimization that replaces implicit LLM heuristics with a dual-level memory architecture. It utilizes long-term memory for storing reusable expert optimization skills and short-term memory for tracking task-specific trajectories, achieving significant speedups on KernelBench.
Entities (6)
Relation Signals (3)
KernelSkill â evaluatedon â KernelBench
confidence 100% · On KernelBench Levels 1-3, KernelSkill achieves a 100% success rate
KernelSkill â utilizes â dual-level memory architecture
confidence 100% · Specifically, we present KernelSkill, a multi-agent framework with a dual-level memory architecture.
Generator â partof â KernelSkill
confidence 95% · KernelSkill consists of key components, including the Generator...
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
Abstract: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 this https URL.
Tags
Links
- Source: https://arxiv.org/abs/2603.10085v1
- Canonical: https://arxiv.org/abs/2603.10085v1
Trouble viewing inline? Open PDF directly â
Full Text
57,964 characters extracted from source content.
Expand or collapse full text
KernelSkill: A Multi-Agent Framework for GPU Kernel Optimization Qitong Sun 1 2 Jun Han 1 Tianlin Li 1 2 Zhe Tang 2 Sheng Chen 2 Fei Yang 2 Aishan Liu 1 Xianglong Liu 1 Yang Liu 3 2 Abstract 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. How- ever, 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 optimiza- tions. Our key insight is to replace implicit heuristics with expert optimization skills that are knowledge-driven and aware of task trajec- tories. Specifically, we present KernelSkill, a multi-agent framework with a dual-level memory architecture. KernelSkill operates by coordinat- ing 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.44Ă, 2.82Ă, and 1.92Ă over Torch Eager on Levels 1, 2, and 3, respectively, outperforming prior baselines. Code is avail- able athttps://github.com/0satan0/ KernelMem/. 1. Introduction Efficient GPU kernels sit at the core of modern AI sys- tems because they largely determine the throughput of fundamental operators such as GEMM, convolution, and data movement. This importance is amplified by todayâs software stacks: mainstream frameworks and compilation toolchains are tightly coupled with GPU backends (Paszke et al., 2019; Tay et al., 2022; Pandey et al., 2022; Ansel et al., 2024), so using these frameworks effectively means inheritingâand depending onâthe efficiency of the un- 1 School of Computer Science and Engineering, Beihang Uni- versity, China 2 Zhejiang Lab, China 3 Nanyang Technological University, Singapore. Correspondence to: Tianlin Li<tian- lin001@buaa.edu.cn>. derlying kernels. Achieving high performance, however, remains highly expert-driven and architecture-aware, requir- ing careful coordination of memory access patterns, paral- lel decomposition, synchronization, and numerics-sensitive instruction choices. As a result, expert-driven kernel opti- mization typically involves long, profiling-guided iteration cycles, bottleneck diagnosis, targeted edits, and repeated validation, making development costly and slow. Against this backdrop, recent advances in large language models (LLMs) have opened a promising avenue for auto- mated GPU kernel generation and optimization. Existing LLM-based approaches typically follow two paradigms: (i) training-based adaptation to encode kernel-related priors (Li et al., 2025b; Baronio et al., 2025), and (i) inference- time, closed-loop refinement where the model iteratively revises kernels using compilation/correctness checks and profiling feedback (Zhang et al., 2025; Wei et al., 2025; Lei et al., 2025; Dong et al., 2025a). However, important gaps remain before such systems meet the needs of real kernel development. First, training-based adaptation is not fully aligned with how kernel engineers iterate in practice, which is driven by a profiling-to-edit feedback loop as bot- tlenecks shift. Additionally, training requires substantial data curation and training compute, resulting in high cost and long training time. Second, multi-round refinement often suffers from imprecise optimization-method selection: the system may pick mismatched strategies for the kernelâs true bottleneck or fail to update method choices as profiling signals evolve, leading to inefficient exploration and wasted iterations. Moreover, optimization decisions are frequently weakly interpretable, making it difficult to justify why a strategy was selected and to reuse that experience across kernels. Our key insight is to replace implicit, model-internal heuris- tics with kernel optimization skills that are both knowledge- driven and trajectory-aware. These skills encode reusable ex- pert optimization expertise accumulated across kernels (e.g., recurring bottleneckâmethod correspondences). Based on this, we propose KernelSkill, a memory-augmented multi- agent framework that (i) externalizes expert optimization know-how to support traceable, audit-friendly decisions, which selects targeted kernel optimization skills tailored to the context of each task, and (i) maintains explicit opti- 1 arXiv:2603.10085v1 [cs.LG] 10 Mar 2026 KernelSkill: A Multi-Agent Framework for GPU Kernel Optimization mization state across rounds to stabilize refinement under profiling feedback. It couples structured long-term opti- mization knowledge with task-specific short-term trajectory memory. This design enables interpretable method selection and robust, feedback-driven iterative improvement across heterogeneous kernels. Our main contributions are as fol- lows: âąKernelSkill: a memory-augmented multi-agent op- timizer. We present a closed-loop agent architec- ture that coordinates generation, verification, profiling, planning, and repair for iterative kernel optimization. âąTwo-level memory for stability and interpretabil- ity. We distill a substantial corpus of expert-level GPU optimization skills into a structured, retrievable long- term knowledge memory to support reusable and ex- plainable method selection, and complement it with a short-term trajectory memory that tracks per-task op- timization and repair history to stabilize multi-round refinement. âąStrong empirical validation on KernelBench. Through extensive experiments on KernelBench Levels 1â3, we demonstrate the effectiveness and robustness of KernelSkill; with ChatGPT-5.1, it achieves a 100% success rate and average speedups of 5.44Ă, 2.82Ă, and 1.92Ă over Torch Eager. 2. Related Works 2.1. Compilers and Autotuning To reduce the cost of expert kernel tuning, prior work has pursued automated GPU kernel optimization via compil- ers/DSLs and search-based autotuning (Tillet et al., 2019; Chen et al., 2018b). In practice, two paradigms dominate: (i) expert-maintained vendor libraries (e.g., cuDNN (Chetlur et al., 2014)) and (i) compiler/DSL ecosystems that expose transformations and scheduling primitives (e.g., TVM (Chen et al., 2018a), Triton (Tillet et al., 2019), ThunderKittens (Spector et al., 2024)). While effective, both require sub- stantial engineering effort and continual upkeep as hardware and workloads evolve. 2.2. Training-Based LLM Methods Recent progress in large language models (LLMs) has made them a viable tool for GPU kernel synthesis and optimiza- tion, leveraging their strong ability in code generation and transformation (Dong et al., 2025b; Jiang et al., 2026). Ex- isting LLM-based approaches broadly fall into two lines. The first line focuses on training or adapting specialized models, often via reinforcement learning to improve ker- nels through iterative optimization signals, exemplified by CUDA-L1 and Kevin (Li et al., 2025b; Baronio et al., 2025). MTMC further learns hierarchical policy guidance that can be executed step-by-step by general-purpose LLMs (Zhu et al., 2025). Related efforts learn optimization behaviors across the stack, from Triton-level synthesis to low-level scheduling (Li et al., 2025a; Woo et al., 2025; He & Yoneki, 2025). While promising, training-centric pipelines are costly and can struggle to generalize in an enormous search space (Zhai et al., 2024); moreover, they are typically driven by offline supervision or learned priors, making it harder to in- corporate per-kernel runtime feedback (e.g., profiling shifts) to close the profilingâdiagnosisâedit loop for targeted online strategy revision. 2.3. Agentic Optimization A second line of studies inference-time, closed-loop opti- mization, where LLMs iteratively refine kernels through compilation checks, correctness tests, and, when available, profiling feedback. This line spans kernel-level multi-agent refinement (Wei et al., 2025; Lei et al., 2025; Zhang et al., 2025; Dong et al., 2025a), DSL-oriented synthesis and tun- ing (Wang et al., 2025), evolutionary improvement under timing-only feedback (Andrews & Witteveen, 2025), and end-to-end PyTorch inference optimization (Nagaitsev et al., 2025). At the kernel level, several systems follow a shared profiling-grounded refinement paradigm: they use profil- ing signals (e.g., GPU specifications and NCU metrics) to diagnose bottlenecks and then trigger targeted edits in a multi-round loop (Zhang et al., 2025; Wei et al., 2025; Lei et al., 2025). CudaForge instantiates this loop with a lightweight CoderâJudge workflow (Zhang et al., 2025), while Astra generalizes it into a multi-agent pipeline with specialized roles (e.g., generation, testing, profiling, and planning) (Wei et al., 2025); PRAGMA further strengthens the bottleneck-to-action mapping by explicitly associating profiling evidence with concrete optimization steps (Lei et al., 2025). STARK complements closed-loop refinement with grounded instruction, dynamic context management, and strategic search for more systematic exploration (Dong et al., 2025a). Despite strong empirical gains, existing closed-loop optimizers often provide limited support for stabilizing refinement across rounds and for making method choices traceable, and they typically lack mechanisms to persist optimization state and reuse experience across ker- nels. For example, STARK introduces within-task memory to reduce repeated attempts (Dong et al., 2025a); in contrast, KernelSkill separates cross-task reusable knowledge from per-task trajectory state, enabling more traceable method se- lection and more stable iterative refinement under profiling feedback. 2 KernelSkill: A Multi-Agent Framework for GPU Kernel Optimization Memory Bank Long-term Memory Short-term Memory Multi-Agent Refinement Seed Generation Retrieval Pytorch Reference Seed Kernel Reviewer Deterministic Decision Policy Method Knowledge Repair Memory Optimization Memory Planner DiagnoserRepairer Optimizer __global__ void Rep_kernel(...)... __global__ void Opt_kernel(...)... Repaired Kernel Optimized Kernel Generator Feature Extractor CompilerVerifierProfiler Code Feature Refined kernel Profiler Info Rule-basedLLM-based class Model(n.Module): def __init__(self, ...): ... def forward(self, ...): ... Seed 1 Seed 2 Seed 3 __global__ void TiledSharedMem(...)... __global__ void VectorizedLoad(...)... __global__ void FusedEpilogue(...)... Method Set & Knowledge Optimization History Compiler Info Verifier Info Repair History Base Kernel New Repair Plans Base Kernel New Optimization Plans Nsight Compute Nsight Systems Speed up Tree of Optimization Trajectories Optimization Plan & Feedback Compilation status: Success or Error ? Compilation log Chain of Repair Trajectories Repair Plan & Feedback PyTorch reference Kernel outputs Consistent? Generation Rule 1.Translates PyTorch reference to CUDA kernels. 2.Ensuring correctness (compilation & output match). 3.Guided by formatting examples. 4....... KernelValue Semantics Judgment CriteriaKey Method Intent Requirements Expected Metric Change Case Field Mapping Run Features Schema Derived Fields Code FeaturesDecision Table ...... Forbidden Patterns ...... Has K LoopData Reuse Tail Handling Overhead Vector Load Store Source Counters Warp State Stats Memory Workload Analysis Occupancy ...... Pointwise Speed Of Light ...... Key Kernel Pattern Rules Machine Check Figure 1. Overview of KernelSkill. 3. Motivation As reviewed in Section 2, both training-based and agentic LLM optimizers often leave optimization-method selection to implicit, model-internal heuristics rather than explicit, reusable expertise. In real kernel development, however, profiling signals can be noisy and high-dimensional, bot- tlenecks may shift across rounds, and many optimization actions have strong preconditions that depend jointly on runtime evidence and code structure. These factors make method selection particularly error-prone: without an ex- plicit mechanism to encode reusable optimization expertise and to track per-task refinement history, an optimizer can easily spend its limited budget on low-yield directions and produce decisions that are difficult to justify or reuse. We next illustrate this failure mode with a concrete KernelBench example. Motivating example: imprecise method selection.Con- sider the KernelBench task in Appendix D, which applies a large linear projection (1024 Ă 8192by8192 Ă 8192) followed by lightweight elementwise operators. Without explicit optimization knowledge and state tracking, the opti- mizer first pursued operator fusion and generated a ker- nel that fusesLinear(GEMM), scaling, residual add, and clamping into a single CUDA kernel, while leaving logsumexpandmishunfused. However, this kernel im- plements GEMM as a naive global-memory dot-product loop without shared-memory tiling, vectorized loads, or tensor-core utilization, so the dominant GEMM bottleneck remains under-optimized. As a result, the achieved speedup is only 0.032Ă, and subsequent refinements keep allocating iterations to further fusion attempts rather than addressing the primary GEMM efficiency issue. This example high- lights that the key challenge is not merely âmore refinement rounds,â but accurate and traceable optimization-method selection: the optimizer must identify which bottleneck mat- ters most, choose high-leverage methods accordingly, and preserve the refinement state to avoid repeatedly spending the budget on low-yield directions. Design principles and KernelSkill. These observations motivate two principles: (i) externalize expert optimiza- tion knowledge into a retrievable and auditable form that grounds method selection in profiling evidence and makes bottleneck-to-action decisions traceable, empowering the system with the ability to apply kernel optimization skills across tasks; (i) maintain explicit per-task optimization state to stabilize multi-round refinement and support multi-step coupled improvements. KernelSkill instantiates these prin- ciples via a two-level memory design within a multi-agent closed loop. A structured long-term knowledge memory enables evidence-grounded and auditable method selection, while a short-term trajectory memory preserves per-task optimiza- tion state to reduce oscillations and support multi-step cou- pled edits. Together, the two memories improve both deci- sion quality and refinement stability, allowing KernelSkill to achieve larger speedups with fewer refinement rounds under profiling feedback. 4. KernelSkill We now present KernelSkill, a memory-augmented multi- agent framework for reliable GPU kernel optimization un- 3 KernelSkill: A Multi-Agent Framework for GPU Kernel Optimization der profiling feedback. We first introduce KernelSkillâs multi-agent refinement pipeline and the responsibilities of each agent. We then detail the two-level memory mech- anism that is integrated into this pipeline: the long-term memory integrates a library of targeted kernel-optimization skills distilled from prior GPU optimization literature(Hijma et al., 2023) and accumulated kernel engineering experience, thereby grounding kernel optimization in expert knowledge and guiding optimization-method selection; meanwhile, the short-term memory records per-task trajectories to condition planning and repair, including what is stored, how retrieval is performed, and how memory is updated across rounds. 4.1. Multi-Agent Optimization 4.1.1. FRAMEWORK OVERVIEW An overview of KernelSkill is shown in Figure 1. Algo- rithm 1 in Appendix A summarizes the execution loop of KernelSkill. KernelSkill consists of key components, including the Gen- erator, Reviewer, Feature Extractor, Long-term Memory (expert optimization knowledge base), Retrieval (method retrieval over long-term memory), Planner, Optimizer, Short- term Memory (per-task optimization/repair trajectories), Di- agnoser, and Repairer. Starting from a PyTorch reference implementation, Ker- nelSkill first uses the Generator to produce a small set of seed kernels, which are then evaluated by the Reviewer; the best-performing seed is selected as the initial solution. KernelSkill then iterates a closed-loop refinement process. At each round, the Reviewer produces three types of feed- back: compilation, correctness, and (if applicable) profiling signals. Here, compilation means the generated kernel can be successfully built/compiled, and correctness means the kernelâs output matches the PyTorch reference under the same inputs within the benchmarkâs validation criteria. At each round, KernelSkill follows a two-branch control flow. If compilation or correctness fails, the Diagnoser uses the failure signals and short-term repair memory to produce a repair plan, which the Repairer applies to generate a re- vised kernel. Otherwise, KernelSkill enters the optimization branch: the Feature Extractor derives static code features, Retrieval queries the long-term memory with code features and profiling feedback to obtain candidate methods, and the Planner uses these methods together with short-term opti- mization memory to produce an optimization plan executed by the Optimizer. The loop runs for at mostNrounds, and returns the best- performing kernel. We then introduce the key components of KernelSkill in detail. 4.1.2. GENERATOR The Generator agent translates a given PyTorch reference program into an equivalent implementation augmented with custom CUDA kernels. At this stage, its goal is correctness, including successful compilation and output equivalence to the PyTorch reference. To provide diverse starting points for later refinement, the agent attempts to materialize as many operator-level kernels as possible, covering most compute steps in the original program and yielding a broad set of seed kernels. Notably, the Generator agent does not optimize for speed; performance improvements are deferred to the subsequent profiling-driven refinement loop. 4.1.3. FEATURE EXTRACTOR The Feature Extractor derives a set of static code features from the kernel source, i.e., signals obtained purely by source inspection without executing the kernel or collect- ing runtime profiles. We introduce static features because optimization-method selection is not solely determined by profiling evidence: many strategies depend on structural properties of the code (e.g., memory-access patterns and precision/intrinsic usage), which are either weakly reflected or ambiguous in profiling metrics. For instance, an optimizer may reorder loops/accumulation to increase data reuse and reduce redundant global loads, or switch to TensorCore- /intrinsic-friendly math paths when the code structure and precision constraints allow. Therefore, static features com- plement profiling feedback by capturing what the kernel is (and what edits are feasible), while profiling indicates where it is slow. These features serve as retrieval keys for selecting candidate optimization methods from the long-term memory. We cur- rently define 18 feature types that characterize optimization opportunities at a fine granularity; this set is not exhaustive and can be expanded as we observe new kernel patterns and optimization needs. Feature extraction follows a hybrid design with two mecha- nisms:â¶rule-based pattern matching over the source code, andâ·LLM-based inference for features that are difficult to capture reliably with syntax patterns alone. Rule-based extraction is used for features with stable lexical/syntactic signatures (e.g., explicit API/intrinsic usage, fixed idioms, or unambiguous code patterns), providing deterministic out- puts and strong controllability. LLM-based extraction is used for features whose surface forms vary substantially across implementations (e.g., semantically equivalent but syntactically diverse indexing logic, tiling schemes ex- pressed in different styles, or implicit assumptions about layout and reuse), where rigid patterns would be brittle and lead to low recall. In the LLM mode, the model is prompted with the feature definition, the allowed value range, and the raw kernel code, and then outputs the corresponding fea- 4 KernelSkill: A Multi-Agent Framework for GPU Kernel Optimization ture value. This division improves robustness across diverse kernel styles while retaining stability for well-structured pat- terns, enabling more reliable method retrieval and planning. 4.1.4. REVIEWER KernelSkill relies on a Reviewer module to produce the execution feedback that drives both repair and profiling- guided optimization; the Reviewer consists of a Compiler, a Verifier, and a Profiler. The Compiler builds the generated kernel and returns compilation status together with warn- ings and error messages. The Verifier checks functional cor- rectness by comparing kernel outputs against the PyTorch reference under a numerical tolerance. The Profiler collects runtime performance signals via profiling tools (e.g., nsys and ncu), including execution time and resource usage (e.g., register usage and shared-memory footprint) to characterize performance bottlenecks. These signals are consumed by downstream agents to decide the next refinement action. 4.1.5. DIAGNOSER The Diagnoser agent is triggered when compilation fails or the generated kernel violates output consistency. Given the diagnostic messages from the Compiler and the Veri- fier, it infers likely root causes and proposes candidate fixes. Kernel repair is typically multi-step: fixing one error often exposes new ones or reintroduces past regressions. Accord- ingly, KernelSkill treats repair as an iterative process that progressively searches for a correct variant under compiler and verifier feedback. A recurring failure mode in this pro- cess is cyclic repair. Because each repair step is conditioned on the latest feedback, the agent may alternate between a small set of faulty variants (e.g., fixing one error while reintroducing another), leading to oscillation rather than convergence. To mitigate such patterns, we maintain a short- term repair memory that records recent repair attempts and their outcomes, and conditions the Diagnoser on this history to discourage repeated revisits of known-failing edits. We detail this short-term memory mechanism in Section 4.2.2. 4.1.6. PLANNER Once a kernel passes compilation and correctness verifica- tion, KernelSkill uses its static code features and profiling feedback to retrieve a small set of candidate optimization methods from the long-term memory. The Planner receives the retrieved methods together with their rationale and imple- mentation patterns, and then selects a method and produces a concrete, stepwise optimization plan for the current base kernel. To stabilize multi-round refinement, the Planner maintains a task-specific short-term memory that records previously attempted methods and their observed outcomes. This memory discourages repeating unproductive strategies, helps prioritize more promising directions, and supports co- ordinated application of coupled optimizations while keep- ing the refinement process method-by-method. 4.1.7. OPTIMIZER AND REPAIRER The Optimizer and the Repairer execute the plans produced by upstream agents. Given the current kernel, they trans- late the selected optimization or repair steps into concrete code edits, aiming to faithfully implement the prescribed strategy while preserving correctness. They also enforce environment-specific constraints (e.g., kernel signatures, code structure, and formatting) to ensure that the result- ing kernels are directly compilable and executable in our evaluation pipeline. 4.2. Memory Bank We find that a multi-agent optimizer without explicit mem- oryâi.e., relying only on the current kernel and raw runtime feedbackâexhibits systematic failure modes that undermine both stability and traceability. First, optimization method selection becomes biased: when profiling outputs are ver- bose, the model may over-attend to noisy or tool-suggested signals (e.g., heuristic hints in NCU) and miss more suitable expert strategies. Second, the refinement process tends to be myopic, focusing on locally simple edits and overlook- ing coupled transformations such as fusion or coordinated memory-access restructuring. Third, without a historical state, multi-round refinement can oscillateâreverting and reapplying ineffective changesâor fragment coupled strate- gies into isolated steps that lose effectiveness. These issues motivate an explicit memory design thatâ¶externalizes reusable expert knowledge for traceable method selection andâ·preserves task-specific trajectory state to stabilize refinement and achieve larger speedups with fewer rounds under profiling feedback. 4.2.1. LONG-TERM MEMORY Expert knowledge sources. The long-term memory is built from expert kernel-optimization knowledge distilled by us. Concretely, we start from a large-scale GPU optimiza- tion survey (Hijma et al., 2023) and translate its taxonomy into operational decision knowledge by following a three- step curation procedure. (1) Scenario abstraction: we group optimization techniques into recurring, task-agnostic scenarios (e.g., memory-bound access patterns, underuti- lized parallelism), and for each scenario we identify the ob- servable decision factors that practitioners use in bottleneck identification. (2) Evidence formalization: we specify how each decision factor is measured from our available signals, i.e., normalized profiling metrics (viafieldmapping), runtime features, and static code features, and we define de- terministic composite indicators (derivedfields) so that the same concept can be evaluated consistently across 5 KernelSkill: A Multi-Agent Framework for GPU Kernel Optimization Kernel #1Kernel #2Kernel #4Kernel #3 Feedbacks Repair Plan FeedbacksFeedbacks Repair Plan Feedbacks Repair Plan Kernel #5 Optimization Plan Short-term memory -- repair Figure 2. The short-term memory for the current repair round. Kernel #1Kernel #2Kernel #4Kernel #3 FeedbacksFeedbacksFeedbacksFeedbacks Kernel #5 Base #1 Speed upâŹBase#2 Optimization Plan Optimization Plan Optimization Plan Optimization Plan Short-term memory -- optimization Figure 3. The short-term memory for the current optimization round. kernels and tool versions. (3) Rule materialization: we en- code scenario-to-method criteria into an auditable determin- istic decision policy (predicates, priority rules, and global veto rules) and a decision table that maps matched evidence patterns to a candidate method set (allowedmethods); the companionllmassiststore provides method-level rationales and implementation cues. This design makes the knowledge reusable because it is expressed as task-independent scenarios and standardized evidence predicates that can be instantiated for different ker- nels using the same feature/profiling interface. It is auditable because method selection proceeds through deterministic gating: for any recommendation, we can record which nor- malized fields and predicates were satisfied, which rule/case in the decision table was matched, and which global con- straints (if any) vetoed alternatives. As a result, KernelSkill performs traceable method selectionâthe selected method set is justified by explicit evidence-to-rule matches and ac- companied by method-specific rationalesârather than ad- hoc trial-and-error driven by implicit prompt preferences. Concretely, long-term memory consists of two complemen- tary parts: Deterministic Decision Policy, which performs rule-based gating, prioritization, and method candidate se- lection from profiling/code evidence; and Method Knowl- edge, which provides method-specific rationales and imple- mentation cues to make the selected actions interpretable and easy to execute. A full schema of the long-term memory fields is provided in Appendix B. Given profiling feedback and code features, KernelSkill first applies the Deterministic Decision Policy to filter and prior- itize feasible method candidates, and then consults Method Knowledge to provide rationales and concrete implementa- tion cues for planning and execution. Optimization-Method Retrieval. KernelSkill imple- ments a structured retrieval pipeline to select appropriate kernel optimization methods from long-term memory. The workflow takes as input the kernelâs profiling signals and structural characteristics, including static code features, NCU metrics, and runtime features. To ensure robustness and consistency across kernels and profiling environments, the retrieval procedure proceeds through a sequence of de- 6 KernelSkill: A Multi-Agent Framework for GPU Kernel Optimization terministic normalization, feature derivation, and rule-based matching steps, followed by method interpretation for plan- ning. The complete long-term memory decision workflow is provided in Appendix C. Overall, this design separates deterministic method screen- ing (Deterministic Decision Policy) from method explana- tion and implementation cues (Method Knowledge). Such separation improves the rationality, stability, and inter- pretability of optimization-method selection, while still leveraging LLM flexibility for plan synthesis. 4.2.2. SHORT-TERM MEMORY Short-term memory maintains an explicit per-task refine- ment trace, which enables consistent diagnosis and plan re- vision across rounds. It coordinates multi-agent iteration by conditioning the Diagnoser/Planner on past plans, edits, and outcomes. After each round, the short-term memory records: (i) the Plannerâs optimization plan and the Diagnoserâs re- pair plan, (i) the Compiler/Verifierâs outcomes (compilation status and correctness check against the PyTorch reference), and (i) the Profilerâs performance feedback (e.g., latency and selected NCU/NSYS signals), together with the result- ing kernel version. This trace is injected into the prompts of the Planner and the Diagnoser in subsequent rounds, so they can condition the next plan on what has already been tried and what feedback it produced. Kernel optimization is a multi-round process: improvements typically emerge through successive, evidence-guided edits rather than a single modification. Without an explicit per- task trace, the Planner and Diagnoser may lose track of prior decisions and outcomes, leading to repeated ineffective attempts or reverting progress during later fixes. Short- term memory prevents this by maintaining a clear planâ result history, enabling more consistent plan revision across rounds and steadier refinement under profiling feedback. Repair Stage. During the repair stage, the Diagnoser utilizes short-term memory to generate repair plans. At this stage, the short-term memory may consist of multiple chained memory segments. Each chain begins with a kernel that first fails compilation or correctness verification. As illustrated in Figure 2, the chain starts from kernel #2, and in each subsequent iteration, the most recently generated kernel is used as the base kernel for the next repair step. Importantly, although each repair iteration operates on the latest kernel version, the repair plan itself is generated based on the entire history of repair attempts and their correspond- ing outcomes within the chain. For example, kernel #5 is generated by the Repairer based on a repair plan pro- duced by the Diagnoser, which in turn is derived from the accumulated feedback spanning from kernel #2 to kernel #4. Thus, the repair process is guided jointly by historical repair knowledge and the most recent execution feedback, enabling more informed and stable error correction. Optimization Stage. The use of short-term memory in the optimization stage differs from that in the repair stage, particularly in how the base kernel is selected. During repair, the base kernel is always the most recent kernel in the chain. In contrast, during optimization, the base kernel must first pass compilation and correctness verification. Whether the base kernel should be updated is determined by two hyperparameters: a relative speedup threshold and an absolute speedup threshold. The relative speedup measures the percentage improvement of the newly generated kernel over the current base kernel, while the absolute speedup measures the absolute difference in execution performance. If either threshold is exceeded, the base kernel is updated accordingly. As illustrated in Figure 3, all optimization methods pre- viously applied to the same base kernel, along with their execution results, are stored in the short-term memory and incorporated into the context for generating the next opti- mization plan. Based on this plan, the Optimizer produces a new kernel by iterating on the current base kernel, which is then evaluated in the next optimization round. 5. Experiments 5.1. Benchmark and Metrics KernelBench is an open-source benchmark for evaluating whether LLMs can generate correct and efficient GPU ker- nels for PyTorch workloads (Ouyang et al., 2025). It or- ganizes tasks into four difficulty levels: Level 1 contains single-kernel operators (100 tasks), Level 2 contains multi- operator workloads that require implementing and coordinat- ing multiple operators/kernels (100 tasks), Level 3 includes full model architectures (50 tasks), and Level 4 targets end- to-end optimization of Hugging Face model architectures. Since prior LLM-based kernel generation and optimization methods are primarily evaluated on Levels 1â3, we follow prior work and report results on the 250 tasks from Levels 1â 3. KernelBench additionally provides metrics (e.g., fast p ) that quantify the fraction of generated kernels that are both func- tionally correct and faster than a baseline by a specified threshold. In this work, we evaluate KernelSkill and com- pare against prior approaches using the following metrics: â¶Success: the percentage of tasks for which the method generates a kernel that compiles and passes correctness verification. â·Speedup (vs. Torch Eager): the average execution speedup across tasks, measured as the runtime ratio relative 7 KernelSkill: A Multi-Agent Framework for GPU Kernel Optimization Table 1. Success and Speedup Results. Success is the percentage of tasks for which a method generates a kernel that compiles and passes correctness verification. Speedup is the average runtime speedup over Torch Eager. KernelSkill achieves the best overall performance, matching the highest success rate while delivering the largest speedup across all levels among all compared baselines. MethodLevel 1Level 2Level 3 SuccessâSpeedupâSuccessâSpeedupâSuccessâSpeedupâ Kevin-32B0.831.180.921.740.460.32 Astra0.951.480.980.990.930.90 PRAGMA0.951.490.981.020.940.92 CudaForge0.961.451.002.100.961.28 QiMeng1.002.200.991.220.700.73 STARK1.003.031.002.691.001.58 KernelSkill1.005.441.002.821.001.92 Table 2. Ablation Results. We ablate memory components: Long-term memory stores retrievable kernel optimization knowledge for reusable and explainable method selection, while Short-term memory records per-task optimization/repair trajectories to stabilize multi-round refinement. Removing either memory reduces Fast 1 and Speedup, confirming both components are critical to KernelSkillâs improvements. MethodLevel 1Level 2Level 3 SuccessâFast 1 âSpeedupâSuccessâFast 1 âSpeedupâSuccessâFast 1 âSpeedupâ w/o memory0.960.461.520.980.751.050.940.360.95 w/o Short term memory0.960.513.530.980.852.420.940.701.78 w/o Longterm memory1.000.551.771.001.001.181.000.791.04 KernelSkill1.000.625.441.001.002.821.000.821.92 Table 3. Fast 1 Results. Fast 1 is the percentage of tasks whose generated kernel is at least as fast as the Torch baseline. MethodLevel 1Level 2Level 3 Kevin-32B0.160.610.02 Astra0.430.730.35 PRAGMA0.450.740.36 CudaForge0.540.890.68 QiMeng0.590.660.40 STARK0.711.000.87 KernelSkill0.621.000.82 to Torch Eager (the unoptimized PyTorch implementation). âžFast 1 : the percentage of tasks for which the generated kernel is at least as fast as the Torch baseline. 5.2. Baselines We compare KernelSkill against representative training- based and agentic baselines that have been widely evaluated (or can be reproduced) on KernelBench Levels 1â3. Training-based kernels. Kevin trains an LLM with a multi- turn RL recipe that explicitly models iterative refinement during kernel generation (Baronio et al., 2025). QiMeng proposes a macro-to-micro hierarchical paradigm: it learns optimization-policy guidance at a high level (âmacro think- ingâ) and executes it via stepwise low-level code realization (âmicro codingâ) to improve correctness and efficiency (Zhu et al., 2025). Agentic optimization with feedback. CudaForge is a training-free CoderâJudge workflow that iteratively gen- erates and refines CUDA kernels, where the Judge leverages hardware feedback (e.g., NCU metrics / GPU specs) to di- agnose bottlenecks and guide targeted edits (Zhang et al., 2025). Astra builds a multi-agent pipeline with specialized roles to optimize existing CUDA kernels through iterative collaboration (Wei et al., 2025). PRAGMA further strength- ens feedback-driven optimization by integrating richer pro- filing signals and a bottleneck-aware reasoning module that links profiling evidence to concrete optimization strategies across iterative refinement (Lei et al., 2025). STARK intro- duces a multi-agent refinement framework with grounded instruction, dynamic context management, and strategic search, and incorporates within-task memory to organize past attempts and reduce redundant exploration (Dong et al., 2025a). 5.3. Experiments Setting All experiments are conducted on an NVIDIA A100 GPU (80GB). We use Python 3.11.13 and PyTorch 2.6.0+cu124 (CUDA 12.4). The CUDA runtime stack includes cuDNN 9.1.0 and cuBLAS 12.9. We use the official KernelBench benchmarking framework to measure runtime. For stable timing, we follow the evaluation procedure in prior work (Zhang et al., 2025) by warming up with fixed input shapes before using CUDA events for timing; we use 25 warm- up iterations and 100 timed iterations and report the mean 8 KernelSkill: A Multi-Agent Framework for GPU Kernel Optimization latency. Profiling signals used by KernelSkill are collected with NVIDIA Nsight Compute (ncu) and Nsight Systems (nsys). KernelSkill starts from seed kernels sampled by the Gen- erator agent. For each task, we sample 3 candidate seeds and pick the best one (compiles, passes verification, and has the lowest mean latency) to start refinement. We run up to 15 refinement rounds per task, with temperature set to 1.0. All agent calls use ChatGPT-5.1 (OpenAI, 2025) as the base model. Following Algorithm 1, we promote a new base kernel only if it achieves at least 30% relative speedup over the current base (rt=0.3) or an absolute speedup gain of at least 0.3 (at=0.3), to avoid unstable base updates from small fluctuations. Result sources and reproducibility.Since Astra, PRAGMA, QiMeng, and STARK are not open-sourced, we cannot run all baselines under an identical codebase. We thus report STARK and QiMeng using the best Kernel- Bench results reported in their original papers. For Astra and PRAGMA, whose papers do not report KernelBench num- bers, we implement the pipelines following their descrip- tions and evaluate them in our experimental setup. These distinctions are noted to clarify result provenance and com- parability. 5.4. Experimental Results and Analysis Table 1 reports results on KernelBench Levels 1â3. Over- all, KernelSkill achieves a 100% success rate across all three levels, indicating reliable end-to-end kernel synthesis under our closed-loop refinement. Among the compared baselines, STARK is the only method that matches 100% success on all levels; other approaches exhibit non-trivial failure rates, especially on harder tasks (e.g., Level 3 suc- cess: Kevin-32B 0.46, QiMeng 0.70, Astra 0.93, PRAGMA 0.94, CudaForge 0.96). In terms of performance, KernelSkill delivers the best av- erage speedup on every level: 5.44Ă, 2.82Ă, and 1.92Ă over Torch Eager on Levels 1, 2, and 3, respectively. Com- pared with the strongest baseline in speedup (STARK), Ker- nelSkill improves average speedup by +79.5% on Level 1 (5.44 vs. 3.03, +2.41Ăabsolute), +4.8% on Level 2 (2.82 vs. 2.69, +0.13Ăabsolute), and +21.5% on Level 3 (1.92 vs. 1.58, +0.34Ăabsolute). This shows that KernelSkill not only preserves perfect correctness, but also consistently pushes kernels further toward higher efficiency, with the most pronounced gains on the easiest and hardest regimes. KernelSkill also substantially outperforms prior profiling- grounded multi-agent baselines (Astra, PRAGMA, Cud- aForge). On Level 1, these systems remain aroundâŒ1.45â 1.49Ăspeedup with<1.0 success, whereas KernelSkill reaches 5.44Ă(3.65Ăhigher than PRAGMA in average speedup). On Level 2 and Level 3, KernelSkill improves over the strongest among them (CudaForge) from 2.10Ă to 2.82Ă(1.34Ăhigher) and from 1.28Ăto 1.92Ă(1.50Ă higher), respectively, indicating stronger optimization effec- tiveness under the same closed-loop setting. Finally, training-centric baselines show clear brittleness on harder tasks. For example, QiMeng achieves competitive performance on Level 1 (1.00 success, 2.20Ăspeedup) but degrades markedly on Level 3 (0.70 success, 0.73Ă speedup). Kevin-32B exhibits an even sharper drop (Level 3: 0.46 success, 0.32Ăspeedup). In contrast, KernelSkill main- tains uniformly perfect success while achieving the high- est speedups across all levels, suggesting that the proposed memory-augmented, profiling-feedback-driven multi-agent refinement provides both robustness and stronger optimiza- tion capability, especially when task difficulty increases. As reported in Tables 1 and 3, while STARK achieves a higher Fast 1 rate on Levels 1 and 3, KernelSkill attains consistently larger average speedups across all three levels. Notably, KernelSkill reaches these gains with only 15 refine- ment rounds, surpassing the average speedups obtained by STARK after 30 rounds. Beyond the final speedups, Kernel- Skill is also markedly more refinement-efficient: when mea- sured by mean speedup divided by the number of refinement rounds, STARKâs per-round gains are 0.10/0.09/0.05 on Levels 1/2/3, whereas KernelSkill achieves 0.36/0.19/0.13, respectively. This contrast is notable because STARK al- ready employs a sophisticated tree-structured memory to record and reuse within-task search trajectories; in com- parison, KernelSkillâs two-level memory design (cross-task long-term memory plus per-task short-term trajectory state) yields more decisive method selection and more stable multi- round refinement, thereby improving the average speedup per refinement roundâlargely due to cross-task long-term memoryâand lifting the overall mean speedup. 5.5. Ablation Table 2 evaluates three variants to isolate the impact of mem- ory: (i) w/o memory, (i) w/o long-term memory, and (i) w/o short-term memory. Removing both memories leads to degraded reliability and weaker optimization, as the system lacks (a) cross-task reusable guidance for method selection and (b) per-task state to avoid ineffective oscillations. Without short-term memory, kernel repair and refinement can become unstable: even after 15 rounds, some tasks fail to reach correctness, yielding success rates of 96%, 98%, and 94% on Levels 1â3. Enabling short-term memory eliminates these failure cases within the same round budget, achieving 100% success across all levels and improving Fast 1 , consistent with its role in tracking per-task attempts, preventing repeated failures, and supporting coupled multi- step fixes. 9 KernelSkill: A Multi-Agent Framework for GPU Kernel Optimization Long-term memory provides cross-task reusable optimiza- tion knowledge that narrows the method search space and reduces trial-and-error. Consequently, introducing long- term memory yields a substantial gain in average speedup over the memory-free baseline, whereas removing long-term memory results in only limited performance improvement. This confirms that long-term memory is critical for selecting higher-impact optimization strategies and achieving larger speedups. 6. Conclusion We propose KernelSkill, a memory-augmented multi-agent framework for GPU kernel optimization. On KernelBench, it achieves a 100% success rate and average speedups of 5.44Ă, 2.82Ă, and 1.92Ă over Torch Eager on Levels 1, 2, and 3, respectively, outperforming prior methods. However, the performance depends on the coverage of the long-term memory; when no matching case is retrieved, KernelSkill may fall back to LLM-only evidence-based method selec- tion. Impact Statement This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here. References Andrews, M. and Witteveen, S. Gpu kernel scientist: An llm- driven framework for iterative kernel optimization, 2025. URL https://arxiv.org/abs/2506.20807. Ansel, J., Yang, E., He, H., Gimelshein, N., Jain, A., Voz- nesensky, M., Bao, B., Bell, P., Berard, D., Burovski, E., et al. Pytorch 2: Faster machine learning through dynamic python bytecode transformation and graph compilation. In Proceedings of the 29th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 2, p. 929â947, 2024. Baronio, C., Marsella, P., Pan, B., Guo, S., and Alberti, S. Kevin: Multi-turn rl for generating cuda kernels, 2025. URL https://arxiv.org/abs/2507.11948. Chen, T., Moreau, T., Jiang, Z., Zheng, L., Yan, E., Shen, H., Cowan, M., Wang, L., Hu, Y., Ceze, L., et al.TVM: An automatedEnd-to-Endoptimizing compiler for deep learning. In 13th USENIX Symposium on Operating Sys- tems Design and Implementation (OSDI 18), p. 578â594, 2018a. Chen, T., Zheng, L., Yan, E., Jiang, Z., Moreau, T., Ceze, L., Guestrin, C., and Krishnamurthy, A. Learning to opti- mize tensor programs. Advances in Neural Information Processing Systems, 31, 2018b. Chetlur, S., Woolley, C., Vandermersch, P., Cohen, J., Tran, J., Catanzaro, B., and Shelhamer, E. cudnn: Efficient primitives for deep learning, 2014. URLhttps:// arxiv.org/abs/1410.0759. Dong, J., Yang, Y., Liu, T., Wang, Y., Qi, F., Tarokh, V., Rangadurai, K., and Yang, S. Stark: Strategic team of agents for refining kernels, 2025a. URLhttps:// arxiv.org/abs/2510.16996. Dong, Y., Jiang, X., Qian, J., Wang, T., Zhang, K., Jin, Z., and Li, G. A survey on code generation with llm- based agents, 2025b. URLhttps://arxiv.org/ abs/2508.00083. He, G. and Yoneki, E. Cuasmrl: Optimizing gpu sass schedules via deep reinforcement learning, 2025. URL https://arxiv.org/abs/2501.08071. Hijma, P., Heldens, S., Sclocco, A., Van Werkhoven, B., and Bal, H. E. Optimization techniques for gpu programming. ACM Computing Surveys, 55(11):1â81, 2023. Jiang, J., Wang, F., Shen, J., Kim, S., and Kim, S. A survey on large language models for code generation. ACM Transactions on Software Engineering and Method- ology, 35(2):1â72, January 2026. ISSN 1557-7392. doi: 10.1145/3747588. URLhttp://dx.doi.org/10. 1145/3747588. Lei, K., Yang, H., Zhang, H., You, X., Zhang, K., Luan, Z., Liu, Y., and Qian, D. Pragma: A profiling-reasoned multi- agent framework for automatic kernel optimization, 2025. URL https://arxiv.org/abs/2511.06345. Li, S., Wang, Z., He, Y., Li, Y., Shi, Q., Li, J., Hu, Y., Che, W., Han, X., Liu, Z., and Sun, M. Autotriton: Auto- matic triton programming with reinforcement learning in llms, 2025a. URLhttps://arxiv.org/abs/ 2507.05687. Li, X., Sun, X., Wang, A., Li, J., and Shum, C. Cuda-l1: Im- proving cuda optimization via contrastive reinforcement learning, 2025b. URLhttps://arxiv.org/abs/ 2507.14111. Nagaitsev, K., Grbcic, L., Williams, S., and Iancu, C. Op- timizing pytorch inference with llm-based multi-agent systems, 2025. URLhttps://arxiv.org/abs/ 2511.16964. OpenAI. Gpt-5.1: A smarter, more conversational chatgpt. https://openai.com/index/gpt-5-1/, 2025. Accessed: 2025-11-12. 10 KernelSkill: A Multi-Agent Framework for GPU Kernel Optimization Ouyang, A., Guo, S., Arora, S., Zhang, A. L., Hu, W., R Ì e, C., and Mirhoseini, A. Kernelbench: Can llms write efficient gpu kernels?, 2025. URLhttps://arxiv. org/abs/2502.10517. Pandey, M., Fernandez, M., Gentile, F., Isayev, O., Tropsha, A., Stern, A. C., and Cherkasov, A. The transforma- tional role of gpu computing and deep learning in drug discovery. Nature Machine Intelligence, 4(3):211â221, 2022. Paszke, A., Gross, S., Massa, F., Lerer, A., Bradbury, J., Chanan, G., Killeen, T., Lin, Z., Gimelshein, N., Antiga, L., et al. Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems, 32, 2019. Spector, B. F., Arora, S., Singhal, A., Fu, D. Y., and R Ì e, C. Thunderkittens: Simple, fast, and adorable ai ker- nels, 2024. URLhttps://arxiv.org/abs/2410. 20399. Tay, Y., Dehghani, M., Bahri, D., and Metzler, D. Efficient transformers: A survey, 2022. URLhttps://arxiv. org/abs/2009.06732. Tillet, P., Kung, H.-T., and Cox, D. Triton: an intermediate language and compiler for tiled neural network computa- tions. In Proceedings of the 3rd ACM SIGPLAN Interna- tional Workshop on Machine Learning and Programming Languages, p. 10â19, 2019. Wang, J., Joshi, V., Majumder, S., Chao, X., Ding, B., Liu, Z., Brahma, P. P., Li, D., Liu, Z., and Barsoum, E. Geak: Introducing triton kernel ai agent & evalua- tion benchmarks, 2025. URLhttps://arxiv.org/ abs/2507.23194. Wei, A., Sun, T., Seenichamy, Y., Song, H., Ouyang, A., Mirhoseini, A., Wang, K., and Aiken, A. Astra: A multi-agent system for gpu kernel performance optimiza- tion, 2025. URLhttps://arxiv.org/abs/2509. 07506. Woo, J., Zhu, S., Nie, A., Jia, Z., Wang, Y., and Park, Y. Tritonrl: Training llms to think and code triton without cheating, 2025. URLhttps://arxiv.org/abs/ 2510.17891. Zhai, Y., Yang, S., Pan, K., Zhang, R., Liu, S., Liu, C., Ye, Z., Ji, J., Zhao, J., Zhang, Y., et al. Enabling tensor language model to assist in generatingHigh-Performancetensor programs for deep learning. In 18th USENIX Symposium on Operating Systems Design and Implementation (OSDI 24), p. 289â305, 2024. Zhang, Z., Wang, R., Li, S., Luo, Y., Hong, M., and Ding, C. Cudaforge: An agent framework with hardware feedback for cuda kernel optimization, 2025. URLhttps:// arxiv.org/abs/2511.01884. Zhu, X., Peng, S., Guo, J., Chen, Y., Guo, Q., Wen, Y., Qin, H., Chen, R., Zhou, Q., Gao, K., Wu, Y., Zhao, C., and Li, L. Qimeng-kernel: Macro-thinking micro- coding paradigm for llm-based high-performance gpu kernel generation, 2025. URLhttps://arxiv.org/ abs/2511.20100. 11 KernelSkill: A Multi-Agent Framework for GPU Kernel Optimization Algorithm 1 Multi-Agent Kernel Optimization with Memory Input: PyTorch reference P , Relative threshold rt, Absolute threshold at, max rounds N Output: Best optimized kernel K n K 0 â GENERATOR(P) (boolc 0 , f eedbackc 0 )â COMPILER(K 0 ) (boolv 0 , f eedbackv 0 )â VERIFIER(K 0 ) (speedup 0 , f eedbackp 0 )â PROFILER(K 0 ) Set base kernelâ K 0 Set bestkernelâ K 0 Set baseâ 0 for i = 1 to N do if not (boolc i and boolv i ) then Repair planâ DIAGNOSER(f eedbackc iâ1 , f eedbackv iâ1 , Repairmemory) K i â REPAIRER(Repairplan, K iâ1 ) else Code f eature base â FEATUREEXTRACTOR(basekernel) M ethodsâ RETRIEVAL(Codef eature base , f eedbackp base ) Optimizationplanâ PLANNER(M ethods, f eedbackp base , Optimizationmemory) K i â OPTIMIZER(Optimizationplan, basekernel) (boolc i , f eedbackc i )â COMPILER(K i ) (boolv i , f eedbackv i )â VERIFIER(K i ) (speedup i , f eedbackp i )â PROFILER(K i ) Update Repairmemory Update Optimization memory if speedup i > speedup base then bestkernelâ K i if ((speedup i /speedup base ) > (1 + rt))or((speedup i â speedup base ) > at) then basekernelâ K i baseâ i end if end if end if end for Return bestkernel A. Algorithm We provide the pseudocode of KernelSkill in Algorithm 1. B. Long-Term Memory Schema. The long-term memory mainly includes the following fields: ⶠfieldmapping: maps raw NCU metrics to standardized fields for consistent downstream processing. â· runfeaturesschema: defines runtime features extracted from Nsight Systems (e.g., kernellaunchcount). âž code features: kernel-structure features extracted via deterministic scanning or LLM-assisted classification. âč derivedfields: deterministic fields derived from fieldmapping, codefeatures, and runfeatures. âș headroomtiers: discretizes optimization headroom (e.g., High/Medium/Low) based on performance indicators. â» bottleneck priorityrules: resolves conflicts when multiple bottlenecks are detected. âŒncupredicates: a library of reusable Boolean predicates over standardized NCU fields. 12 KernelSkill: A Multi-Agent Framework for GPU Kernel Optimization âœglobalforbiddenrules: global veto rules that prevent unsafe or invalid optimizations. ⟠decisiontable: maps bottleneck types, headroom tiers, and kernel characteristics to candidate optimization methods. âżllmassist (Method Knowledge): Provides method-specific rationales and implementation cues for the candidate methods selected by the deterministic policy, enabling traceable decisions and faithful code-level execution. C. Long-Term Memory Decision Workflow. Specifically, the workflow is as follows: ⶠInput aggregation. Collect NCU metrics, runtime features, and code features for the current kernel. â· Metric normalization. Map raw NCU metric keys to standardized field names via fieldmapping, making downstream decisions robust to tool-version-specific naming. âžDerived-field computation. Composite indicators are deterministically computed via derivedfields, producing higher- level features that are more directly aligned with optimization decisions. âčHeadroom tier assignment. The optimization headroom is categorized into discrete tiers (e.g., High/Medium/Low) using headroomtiers, which quantify the remaining optimization potential based on performance indicators. âș Bottleneck identification. The bottleneck type is identified by matching the kernelâs profiling signature against the ncusignature patterns defined in the decisiontable. â»Case matching. Given the bottleneck type and headroom tier, the system matches a specific decision case using conditions such as headroom tier, kernel structural properties, and additional gating predicates (gatewhen). âŒGlobal rule enforcement. Before finalizing the recommendation, the system applies globalforbiddenrules as global veto constraints to prevent unsafe or invalid optimizations, regardless of local decision matches. ✠Method set retrieval. The pipeline returns the final list of permitted optimization methods, i.e., the allowedmethods associated with the matched decision case after global filtering. âŸLLM-assisted planning. After deterministic gating, the system consults the llmassist knowledge base to provide non-binding explanations, implementation guidance, and expected benefits. This information helps the LLM interpret the selected methods and generate a concrete optimization plan. D. Case We provide the code of task in Algorithm 2. Algorithm 2. The code of task. import torch import torch.n as n class Model(n.Module): """ Model that performs a matrix multiplication, scales the result, adds a residual connection, clamps the output, applies LogSumExp, and finally applies the Mish activation function. """ def __init__(self, input_size, hidden_size, scale_factor, clamp_min, clamp_max): super(Model, self).__init__() self.matmul = n.Linear(input_size, hidden_size) self.scale_factor = scale_factor self.clamp_min = clamp_min self.clamp_max = clamp_max def forward(self, x): """ 13 KernelSkill: A Multi-Agent Framework for GPU Kernel Optimization Args: x: Input tensor of shape (batch_size, input_size). Returns: Output tensor of shape (batch_size, hidden_size). """ x = self.matmul(x) x = x * self.scale_factor x = x + x x = torch.clamp(x, self.clamp_min, self.clamp_max) x = torch.logsumexp(x, dim=1, keepdim=True) x = x * torch.n.functional.mish(x) # Mish activation return x We provide the optimized code (without KernelSkill) of task in Algorithm 3. Algorithm 3. The optimized code (without KernelSkill) of task. __global__ void fused_matmul_scale_clamp_kernel( const float * __restrict__ X, const float * __restrict__ W, const float * __restrict__ B, float * __restrict__ Out, int batch, int in_features, int out_features, float scale_factor, float clamp_min, float clamp_max ) int row = blockIdx.x * blockDim.x + threadIdx.x; int col = blockIdx.y * blockDim.y + threadIdx.y; if(row < batch && col < out_features) float acc = 0.0f; for(int k = 0; k < in_features; k++) acc += X[row * in_features + k] * W[col * in_features + k]; acc += B[col]; acc = acc * scale_factor * 2.0f; acc = fminf(fmaxf(acc, clamp_min), clamp_max); Out[row * out_features + col] = acc; torch::Tensor fused_matmul_scale_clamp( torch::Tensor X, torch::Tensor W, torch::Tensor B, float scale_factor, float clamp_min, float clamp_max ) const int batch = X.size(0); 14 KernelSkill: A Multi-Agent Framework for GPU Kernel Optimization const int in_features = X.size(1); const int out_features = W.size(0); auto Out = torch::empty(batch, out_features, X.options()); dim3 block(16,16); dim3 grid((batch + block.x -1)/block.x, (out_features + block.y -1)/block.y); fused_matmul_scale_clamp_kernel<<<grid, block>>>( X.data_ptr<float>(), W.data_ptr<float>(), B.data_ptr<float>(), Out.data_ptr<float>(), batch, in_features, out_features, scale_factor, clamp_min, clamp_max ); return Out; E. Retrieval method and generate optimization plan. Feedbacks Kernel Code FeatureMethod Plan CaseïŒ id headroom kernel_structure gate_when allowed_methods forbidden_methods bottleneck optimisation method method_name modification plan evidence headroom Figure 4. Retrieval method and generate plan. 15