PaperlayerView exact revision on arXiv

Source-linked AI summary

KernelBlaster: Continual Cross-Task CUDA Optimization via Memory-Augmented In-Context Reinforcement Learning

Kris Shengjun Dong, Sahil Modi, Dima Nikiforov, Sana Damani, Edward Lin, Siva Kumar Sastry Hari, Christos Kozyrakis

arXiv:2602.14293v1cs.LGcs.AI

TL;DR

CUDA optimization must adapt across changing workloads and GPU architectures, but existing workflows have limited cross-task memory and traditional model specialization can be costly. KernelBlaster uses memory-augmented in-context reinforcement learning with profile-guided textual feedback and a persistent knowledge base to guide optimization search. It reports improved performance across KernelBench levels and summarizes inference-time adaptation from accumulated successes and failures.

  • Problem

    CUDA optimization requires extensive hardware-specific exploration, while existing agentic workflows have limited ability to retain and reuse knowledge from prior tasks.

  • Method

    KernelBlaster combines in-context reinforcement learning, profile-guided textual feedback, and a persistent CUDA Knowledge Base for cross-task optimization.

  • Results

    KernelBlaster improves performance over PyTorch across KernelBench levels and reports 7.9× speedup over traditional compiler baselines such as IREE.

  • Takeaways & Limitations

    The framework demonstrates that LLM agents can adapt optimization strategies at inference time by learning from histories of successes and failures stored in long-term memory.

  • Takeaways & Limitations

    In the current evaluation, the agent primarily discovers algebraic simplifications, fusion, memory-locality improvements, Tensor Core use, and layout transformations, without observed Hopper-specific warp specialization.

Abstract

from arXiv · show

Optimizing CUDA code across multiple generations of GPU architectures is challenging, as achieving peak performance requires an extensive exploration of an increasingly complex, hardware-specific optimization space. Traditional compilers are constrained by fixed heuristics, whereas finetuning Large Language Models (LLMs) can be expensive. However, agentic workflows for CUDA code optimization have limited ability to aggregate knowledge from prior exploration, leading to biased sampling and suboptimal solutions. We propose KernelBlaster, a Memory-Augmented In-context Reinforcement Learning (MAIC-RL) framework designed to improve CUDA optimization search capabilities of LLM-based GPU coding agents. KernelBlaster enables agents to learn from experience and make systematically informed decisions on future tasks by accumulating knowledge into a retrievable Persistent CUDA Knowledge Base. We propose a novel profile-guided, textual-gradient-based agentic flow for CUDA generation and optimization to achieve high performance across generations of GPU architectures. KernelBlaster guides LLM agents to systematically explore high-potential optimization strategies beyond naive rewrites. Compared to the PyTorch baseline, our method achieves geometric mean speedups of 1.43x, 2.50x, and 1.50x on KernelBench Levels 1, 2, and 3, respectively. We release KernelBlaster as an open-source agentic framework, accompanied by a test harness, verification components, and a reproducible evaluation pipeline.

1 INTRODUCTION

KernelBlaster targets the scalability and adaptability limits of CUDA optimization as workloads and GPU architectures evolve. It uses memory-augmented in-context reinforcement learning to accumulate optimization experience and improve cross-task search.

  • Motivation: Evolving workloads and GPU architectures invalidate previously tuned kernels, making specialized optimization increasingly costly and difficult to scale.The paper cites FlashAttention’s performance drop after porting to H100 as an example of architecture mismatch.
  • Motivation: LLM-based CUDA optimization still has substantial opportunities to improve generalizability, learning capability, sample efficiency, and cost.
  • Approach: KernelBlaster formulates CUDA optimization as in-context reinforcement learning with textual gradient updates from profile data.The framework performs learning at inference time without directly updating model weights.
  • Approach: Its Persistent CUDA Knowledge Base stores optimization experience in a hierarchical representation of performance states for retrieval across tasks and GPU platforms.The representation categorizes code instances into performance states and supports reusable optimization knowledge.
  • Results: 1.43× geometric mean speedup over the PyTorch baseline is reported on KernelBench Level 1, while Level 2 reaches 2.50×.

2 RELATED WORK

Prior LLM-driven CUDA optimization systems span training, prompting, search, memory, evolutionary, and textual-gradient approaches. Their main differences concern experience retention, search scope, adaptability, sample efficiency, and cost.

  • Training-based methods: Training-based CUDA optimizers specialize models through retraining or fine-tuning, requiring reward-driven feedback and substantial samples for weight updates.
  • Static prompting: Static prompt systems use manually engineered heuristics but cannot learn from prior outcomes and may require redesign for new architectures or domains.
  • Search-based methods: Search-based methods improve within-task exploration but typically restart resource-intensive searches for each kernel without retaining reusable lessons.
  • Memory-augmented methods: Memory-augmented systems retain verified kernels, performance data, or hierarchical knowledge bases, but effectiveness depends on retention and retrieval policies.
  • Evolutionary methods: Evolutionary agents preserve history through candidate populations, while full archives can increase storage and context costs and underrepresent negative outcomes.
  • Textual-gradient methods: Textual-gradient methods update prompts through natural-language feedback rather than backpropagating gradients into model weights.

3 METHODOLOGY

KernelBlaster treats CUDA optimization as in-context reinforcement learning, using a persistent Knowledge Base to accumulate profiling-informed strategies across tasks and time. Its workflow matches kernel states to candidate optimizations, explores them through weighted search, and updates the stored policy from performance feedback.

  • Method overview: KernelBlaster uses a persistent CUDA Knowledge Base to consolidate optimization experience across multiple tasks and guide future rollouts.The Knowledge Base aggregates trajectories and retrieves predicted performance values for unseen tasks.
  • Method overview: The framework represents optimization across State, Time, and Task dimensions, with measured rollout outcomes distilled into long-term memory.State captures kernel performance signatures, while time tracks attempts and task distinguishes optimization problems.
  • Knowledge Base construction: A state matcher uses Nsight Compute performance and code signatures to classify kernels as known or newly discovered states.New states are appended to the Knowledge Base; known states retrieve associated optimization candidates.
  • Knowledge Base construction: Weighted random selection samples top-k candidate optimizations using predicted performance gains while preserving exploration beyond the best historical performer.Selected transformations are applied, tested, and profiled for correctness and performance before feedback updates the Knowledge Base.
  • In-context reinforcement learning: KernelBlaster treats the prompt-based Knowledge Base as mutable policy parameters and approximates policy gradients through LLM-based performance evaluation and gap analysis.The update incorporates natural-language analysis of discrepancies between expected and observed optimization performance.

4 EVALUATION

The evaluation compares KernelBlaster with agentic and compiler-based CUDA optimization systems across GPU architectures and benchmark levels under matched execution and profiling conditions. The experimental workflow uses iterative code-generation rollouts with profile-based rewards and replay-buffer evaluation.

  • Evaluation setup: KernelBlaster, AI CUDA Engineer, Kernelsseum, and IREE are evaluated across diverse GPU architectures and benchmark levels.The comparison includes compiler- and agent-driven pipelines under equivalent execution and profiling conditions.
  • Evaluation workflow: The LLM-based policy optimization loop samples code tasks, generates optimized code actions, transitions to the resulting code, and assigns profile-based rewards.The rollout stores state-action pairs for later evaluation and policy updating.
  • Evaluation workflow: Classical REINFORCE updates model policy parameters through estimated gradients, whereas KernelBlaster adapts the loop to strategy-guided LLM rollouts.The comparison establishes the reinforcement-learning structure underlying the in-context formulation.
  • Evaluation setup: The evaluation compares KernelBlaster with PyTorch eager execution, torch.compile, IREE, and AI CUDA Engineer as primary baselines.The baselines span default execution, JIT compilation, compiler optimization, and agentic optimization.

4.2 Evaluation Metrics

The evaluation measures both generated-code performance and optimization-process efficiency. Metrics cover speedup, correctness, high-impact optimization frequency, and token cost relative to defined baselines and attempt budgets.

  • Performance metrics: Code performance is quantified primarily by speedup over the PyTorch baseline, alongside system-level measures of optimization efficiency.The target metrics explicitly include kernel performance and system performance.
  • Performance metrics: Valid Rate measures the percentage of optimization problems passing both functionality and LLM-based soft-verification checks.Mean speedup uses the best performance among PyTorch Eager and torch.compile as the 1.0x baseline.
  • Performance metrics: fastp measures the percentage of kernels achieving at least r times the baseline speedup within k attempts while producing correct outputs.The metric characterizes how frequently an optimization system finds high-impact solutions under an attempt budget.
  • System metrics: System cost is measured by the total number of tokens consumed to optimize a kernel, compared with a minimal agentic iteration loop.This metric evaluates optimization-process efficiency rather than generated-kernel execution speed.

4.3 Execution Harness

KernelBlaster evaluates generated CUDA kernels through compilation, numerical verification, profiling, and soft structural checks. These safeguards address reward hacking by rejecting incorrect or shortcut-based speedups.

  • Execution and verification: The C++ test harness launches optimized CUDA kernels alongside reference Torch implementations and reports numerical verification results.Compilation failures are discarded and returned to the code-lowering agent for another attempt.
  • Execution and verification: After verification, Nsight Compute profiles the optimized code to collect kernel performance information.The harness separates correctness checks from performance profiling.
  • Execution and verification: Generated kernels undergo numerical validation and LLM-based soft verification to detect removed functionality and non-native-library shortcuts.These checks were introduced to prevent illusory speedups caused by evaluation loopholes or structural deviations.

4.5 Comparison against PyTorch

KERNELBLASTER improves on PyTorch across KernelBench Levels 1 and 2, with especially pronounced benefits for Level 2 workloads that offer a larger optimization space.

  • Over 50% of optimized solutions improve upon the best-performing result between native Eager PyTorch and PyTorch Compile on both Level 1 and Level 2.
  • Level 2 workloads show stronger benefits because composed operators provide a larger search space for the agentic flow to exploit.
  • The PyTorch comparison uses the original KernelBench baseline code, with fastp measuring the percentage of kernels meeting the desired improvement target.
  • KERNELBLASTER’s optimization flow starts from a prior CUDA implementation rather than directly optimizing native PyTorch code.

4.7 Comparison Against Other Agentic Workflows

Against other agentic workflows and compiler baselines, KERNELBLASTER shows stronger performance across GPU settings while also extending to composed Level 3 workloads, with full-model scaling remaining constrained.

  • KERNELBLASTER with cuDNN yields a consistently higher percentage of kernels exceeding each speedup threshold than AI CUDA Engineer on L40S Level 1 and Level 2 problems.
  • AI CUDA Engineer reaches a 1.1x geometric mean speedup on Level 1 and 1.7x on Level 2 versus PyTorch, with an 82% validation rate.
  • On Level 3 workloads, generated CUDA reaches 2.68× over PyTorch for LeNet5 and 1.95× for SqueezeNetFireModule.
  • Knowledge learned on Level 1 and Level 2 transfers to Level 3 through fusion, memory-locality improvements, and algebraic simplification.
  • Full-model optimization is limited by processing one optimization per iteration and by verbose CUDA source that dilutes performance signals.

4.10 Cost Summary

Measured performance gains generally increase with token consumption, while KERNELBLASTER outperforms compiler and agentic baselines across the reported comparisons.

  • More tokens consumed are positively correlated with better measured speedup over the original CUDA reference code.Token counts vary with code size, profiled kernels, and optimization complexity despite equal iteration counts.
  • KERNELBLASTER achieves higher geometric-mean speedups over PyTorch than AI CUDA Engineer on KernelBench Levels 1 and 2 and composes with cuDNN.
  • KERNELBLASTER achieves a 7.9× speedup against the IREE compiler baseline.
  • For simple kernels, KERNELBLASTER performs similarly to prior agentic workflows, while Level 2 kernels show improved performance from diverse optimizations.

5 DISTRIBUTION OF OPTIMIZATION USAGE

KERNELBLASTER explores a diverse, heavy-tailed set of optimization techniques, with profiling-guided sequencing distinguishing productive transitions from repetitive micro-tuning.

  • The approximately 50 KB Knowledge Base records diverse optimizations, with no performance state exceeding 20% and an average of 5.5 states reached per kernel.
  • Over 50% of repeated instruction-level-parallelism applications and over 80% of repeated grid-size optimizations yield negligible speedups below 1.01×.
  • Preparatory transitions produce substantial median gains: shared-memory tiling before tensor-core utilization reaches approximately 2.41×.
  • Data-layout transformation before operation fusion yields approximately 1.95×, while control-flow simplification before tensor-core tuning yields approximately 1.42×.
  • Technique usage and success are heavy-tailed: broadly applicable methods dominate attempts, while specialized transformations are invoked selectively.
  • Successful applications concentrate in SIMD, grid-size optimization, instruction-level parallelism, block-size adaptation, work-per-thread increase, register-pressure reduction, fast-math, and thread coarsening.
  • Frequent techniques can also have substantial failure mass, motivating profiling-conditioned selection rather than uniform or frequency-based application.
  • The current evaluation primarily discovers algebraic simplification, fusion, and memory-locality improvements, while also leveraging Tensor Cores and layout transformations.

6 ABLATION STUDIES

The ablations show that persistent knowledge reuse, detailed profiling, and guided search jointly improve CUDA optimization, while search breadth and depth exhibit diminishing returns. The Knowledge Base also supports faster reuse across optimization passes and GPU platforms.

  • Knowledge Base and profiling: Constructing a Knowledge Base from scratch is initially expensive, but later optimization passes achieve faster coverage of optimization strategies.A Knowledge Base trained on an A6000 can also be reused across different GPU platforms.
  • Knowledge Base and profiling: Persistent Knowledge Base reuse is necessary for transferring successful optimization strategies across kernels; profiling alone provides meaningful but limited gains.The no-memory agent, despite full Nsight Compute profiling, achieves 1.67x slower results than the full system.
  • Search hyperparameters: Search breadth has diminishing returns beyond 8 trajectories, while additional trajectories still benefit lower-25th-percentile kernels.Search depth similarly shows diminishing returns beyond 4 optimization steps.
  • Profiling feedback: Removing non-cycle signals degrades performance across most kernels, indicating that feedback beyond raw cycle counts is needed for effective optimization.The cycle-only agent achieves 1.22x speedup over PyTorch on Level-2, compared with 1.57x when using NCU data.
  • Efficiency comparison: The minimal agent uses 2.4x more tokens than KERNELBLASTER and achieves better performance per token in only 0.379x relative terms, while KERNELBLASTER wins in 71% of cases.The comparison uses 10 trajectories of length 10.

7 CONCLUSIONS & FUTURE WORK

KERNELBLASTER applies in-context reinforcement learning and long-term memory to CUDA optimization, achieving reported speedups while exposing workflow characteristics and identifying knowledge-base management and cost-related directions for future work.

  • Conclusions: KERNELBLASTER reports a 1.32x geometric-mean speedup on KernelBench Level 1 and 2 problems over the PyTorch baseline.The conclusion attributes inference-time adaptation to learning from stored histories of successes and failures.
  • Conclusions: The study analyzes optimization diversity, learning ability, search hyperparameters, and complex optimization sequences.These analyses complement the agentic optimization workflow with observations about workflow characteristics.
  • Future work: Future Knowledge Base management should address storage overhead and bias toward early entries through randomized sampling and periodic updates.These strategies are presented as future directions rather than evaluated components of the reported system.
  • Future work: Model heterogeneity is proposed to reduce the quality, cost, and latency burden of LLM calls by matching model capability to agent complexity.The paper also identifies phase ordering as a challenge because transformation effectiveness depends on application sequence.
  • Future work: Aggregated structured experience could eventually support parameter-level model refinement beyond a single agent’s in-context database.The passage describes distilling short-term adaptations into durable capability as a future possibility.

8 APPENDIX

The appendix illustrates KERNELBLASTER’s kernel-level optimizations through algebraic simplification, fusion, shared-memory tiling, instruction-level parallelism, reductions, and Tensor Core execution. These implementations target lower memory traffic and launch overhead while preserving correctness.

  • KernelBench Level 2 Q18: 20.17× speedup over the PyTorch baseline is achieved on KernelBench Level 2 Q18 through exact removal of redundant logsumexp operations.After reductions, the tensor dimension has size one, so logsumexp(x, dim = 1) = x.
  • Fused kernel design: Kernel fusion directly produces the final scalar output, while unrolled accumulation and warp-shuffle reductions improve execution efficiency.These choices reduce global memory traffic and kernel-launch overhead while preserving correctness.
  • Fused kernel design: The fused implementation stages input tiles in shared memory, assigns output work across threads, and uses multiple accumulators to increase instruction-level parallelism.The code also performs cooperative tile loading, unrolled processing, and block-level reduction.
  • Tensor Core implementation: The GEMM implementation uses FP16×FP16→FP32 WMMA, four warps per 32 × 32 block tile, and split-K partitioning across grid.z.Full tiles use direct global loads, while boundary tiles use zero-padded shared-memory buffers.
  • Tensor Core implementation: Weights stored row-major are consumed as column-major matrix-B fragments through pointer and leading-dimension adjustments, avoiding explicit transposition.Split-K partial accumulations use atomicAdd before a separate epilogue applies bias, ReLU, scaling, and FP16 casting.
Loading 2602.14293v1…