PaperlayerView exact revision on arXiv

Source-linked AI summary

Towards Robust Agentic CUDA Kernel Benchmarking, Verification, and Optimization

Robert Tjarko Lange, Qi Sun, Aaditya Prasad, Maxence Faldor, Yujin Tang, David Ha

arXiv:2509.14279v1cs.SEcs.AIcs.LG

TL;DR

Existing CUDA-kernel benchmarks can permit loopholes and narrow testing, limiting robust assessment of LLM-generated kernels. This paper introduces robust-kbench and an agentic framework for CUDA translation, verification, and evolutionary optimization, achieving up to 2.5× forward-pass speedup over PyTorch eager while classifying incorrect kernels with up to 82% accuracy.

  • Problem

    Existing LLM-written CUDA-kernel benchmarks permit loopholes and insufficiently diverse testing, limiting assessment across realistic inputs and applications.

  • Method

    The paper combines robust-kbench with LLM-driven CUDA translation, evolutionary runtime optimization, soft verification, and hardware profiling across forward and backward computations.

  • Results

    Up to 2.5× forward-pass speedup over the PyTorch eager baseline was achieved, while the compilation verifier reached 0.82 accuracy.

  • Takeaways & Limitations

    robust-kbench supports more rigorous assessment of CUDA-kernel correctness, performance, and generalization across varied settings.

  • Takeaways & Limitations

    The reported optimization results require an estimated $5 per kernel in API credits and less than 2 hours on 4 GPUs.

Abstract

from arXiv · show

Recent advances in large language models (LLMs) demonstrate their effectiveness in scaling test-time compute for software engineering tasks. However, these approaches often focus on high-level solutions, with limited attention to optimizing low-level CUDA kernel implementations. Additionally, existing kernel generation benchmarks suffer from exploitable loopholes and insufficient diversity in testing conditions, hindering true generalization assessment. To address these limitations, we introduce robust-kbench, a new benchmark for rigorous evaluation of kernel performance and correctness across varied scenarios. Furthermore, we present a comprehensive agentic framework that automates CUDA kernel discovery, verification, and optimization. This pipeline enables frontier LLMs to translate torch code to CUDA kernels and iteratively improve their runtime within our robust evaluation setting. Our sequential workflow first translates PyTorch code into equivalent CUDA kernels. It then optimizes their runtime using a novel evolutionary meta-generation procedure tailored to the CUDA ecosystem, guided by LLM-based verifiers for correctness and efficient filtering. Evaluated on robust-kbench, our approach produces CUDA kernels outperforming torch implementations for practical applications, including forward and backward passes. It can fuse operations and deploy various runtime optimization strategies. The verifier workflow accurately classifies incorrect kernels, enhancing hardware verification efficiency.

1 Introduction

The paper introduces robust-kbench and an LLM-driven agentic framework for translating PyTorch operations into CUDA kernels, verifying correctness, and optimizing runtime under diverse, realistic conditions. It targets benchmark loopholes and narrow verification that can produce artificial speedup estimates.

  • The work addresses rising computational demands from increasingly complex deep learning models, foundation-model infrastructure, and large-scale data processing.
  • The proposed framework translates functional PyTorch code into CUDA kernels, then iteratively samples, verifies, tests, and evaluates candidate kernels for runtime optimization.Language-model-based verifiers support correctness checking and efficient candidate filtering during parallel optimization.
  • A Robust Benchmark Harness: robust-kbench mitigates exploitable benchmark design and narrow-verification loopholes by testing kernel correctness across varied settings and realistic downstream applications.The harness also supports optimization of both forward and backward kernels.
  • The paper introduces an LLM-driven evolutionary optimization framework to improve the runtime of low-level CUDA operations while enabling robust discovery and self-verification.

2 Background

This section frames test-time compute scaling, evolutionary code optimization, and CUDA as foundations for improving LLM-generated GPU kernels. It also identifies benchmark limitations and practical risks, including loophole exploitation and poor generalization of optimized kernels.

  • Scaling Test-Time Compute: Scaling test-time compute improves LLM outputs by drawing more samples instead of training a new model from scratch.The section highlights search, reasoning incentives, and sample-count scaling as available approaches.
  • Evolutionary Code Optimization with LLMs: Evolutionary code optimization mutates and recombines previously generated code, with LLMs serving as recombination engines.Prior applications include reward and preference optimization, mathematical science code, and machine learning papers.
  • The CUDA Framework: CUDA provides NVIDIA GPU-specific C and C++ extensions for parallel computing, enabling control over threads, memory hierarchies, and synchronization.The framework supports deep learning, scientific computing, and real-time graphics workloads.
  • Challenges for Benchmarking LLM-Written CUDA Kernels: KernelBench v0 introduced 250 PyTorch-defined neural-network tasks spanning three levels from individual operators to fused operations and larger networks.The levels represent increasing structural complexity in CUDA kernel generation.
  • Challenges for LLM-Driven Optimization of CUDA Kernels: LLM-generated CUDA kernels can exploit benchmark loopholes through redundant-operation removal, input-specific hardcoding, or non-generalizable weight assumptions.The section also notes that many optimized kernels fail to transfer benchmark performance to practical utility.

3 robust-kbench: A Robust Agentic CUDA Kernel Discovery Benchmark

This section introduces robust-kbench, a layered harness for evaluating CUDA kernel correctness and performance across varied testing conditions. It also reports agentic translation and optimization results on 200 KernelBench tasks, including 95% translation coverage and reduced speedups after contaminated tasks are excluded.

  • Benchmark results: 95% of all level 1 and level 2 KernelBench tasks are successfully translated by the proposed approach.The result covers the 200-task evaluation set described for the benchmark comparison.
  • Benchmark results: The agentic optimization framework significantly outperforms Kevin-32B on the full 200 tasks, but aggregated speedup significantly decreases after contaminated tasks are excluded.The comparison is presented as part of the KernelBench evaluation addressing benchmarking loopholes.
  • Benchmark harness: robust-kbench combines diverse initialization states, multiple runtime-estimation strategies, and profiling tools to make CUDA kernel evaluation more robust.It uses PyTorch Profiler for high-level metrics, Clang-tidy for static analysis, and NVIDIA Compute Profiler for hardware-level insights.
  • Evaluation interface: The evaluator supports forward and backward testing, manages compilation, correctness verification, and profiling, and parallelizes evaluation across multiple GPUs.Task specifications include input shapes and initialization settings, while the evaluator executes torch baselines and candidate kernels through its evaluation interface.

4 Automating CUDA Kernel Correctness Verification with LLMs

The paper introduces an LLM-based soft-verification workflow to reduce the resource and time costs of hardware-based CUDA kernel verification. Three specialized verifiers target compilation, memory access, and numerical correctness, achieving accuracies of 0.82, 0.80, and 0.73, respectively.

  • Motivation: Traditional hardware verification requires at least one minute of compilation per kernel, with parallelization limited by available GPU hardware.This creates a bottleneck when evaluating many kernel variants or running extensive optimization searches.
  • LLM-Based Verification Workflow: The workflow uses LLMs for rapid soft verification, allowing only the highest-scoring kernels by majority consensus to proceed to hardware testing.Each LLM verifier makes a binary correctness decision for every kernel proposal.
  • Prompt Tuning: The verifier prompt-tuning pipeline iteratively improves error detection and produces prompts that generalize to different downstream verifier models.The specialized verifiers target compilation, memory, and numerical CUDA errors.
  • Specialized Verifiers: Three specialized verifiers achieve accuracies of 0.82 for compilation, 0.80 for memory access, and 0.73 for numerical correctness.Each verifier was tuned on a balanced dataset of 30 kernels.

5 Automating CUDA Kernel Discovery with LLMs

The framework translates PyTorch implementations into verified CUDA kernels, then uses evolutionary LLM optimization to improve runtime. On robust-kbench, it achieves up to 2.5× forward-pass speedup while testing generalization to unseen input shapes.

  • Kernel Translation & Optimization: PyTorch code is translated into CUDA, compiled, and checked against the torch reference, rejecting compilation failures or results breaching 1e−5 precision.Correctness is evaluated numerically against the reference implementation.
  • Kernel Translation & Optimization: The evolutionary loop samples, self-verifies, evaluates, and profiles kernel batches, using up to five correct prior kernels ordered from slowest to fastest as LLM context.The working translated kernel initializes the optimization process, whose archive is updated from previous evaluations.
  • Results: Optimizing individual operations: 2.5× forward-pass speedup is achieved over the PyTorch eager baseline across multiple robust-kbench tasks.LLM-based verification improves optimization stability by reducing regressions and increasing successful kernel evaluation.
  • Results: Generalization of LLM-Optimized Kernels: Unseen-shape evaluation exposes overfitting for LayerNorm and MNIST Linear-ReLU, whereas optimized ResNet block kernels retain performance benefits across input dimensions.robust-kbench evaluates configurations not encountered during evolutionary optimization, unlike fixed-configuration KernelBench testing.

6 Ablating the Agentic Scaffolding for Kernel Verification & Optimization

Ablations show that agentic scaffolding improves CUDA kernel verification and optimization. LLM verification raises valid-kernel rates, while model ensembling and summarized profiling feedback further improve optimization outcomes.

  • Impact of LLM Verification: Verifier-assisted optimization increases the proportion of valid kernels from 55-70% to 80-85%.The verifier detects compilation failures, memory access violations, and numerical inaccuracies across eight CUDA operations, including forward and backward passes.
  • Impact of LLM Verification: Pre-screening with verification filters problematic kernels before hardware evaluation, improving verification efficiency.The evaluated error types include compilation failures, memory access violations, and numerical inaccuracies.
  • Impact of Model Ensembling: Increasing model-ensemble diversity improves optimization outcomes and raises the success rate of generating valid kernels.The comparison covers GPT-4.1 alone, a two-model GPT-4.1 and Claude Sonnet 3.7 ensemble, and a five-model ensemble.
  • Impact of Additional Profiling Information: Summarized profiling feedback produces more targeted kernel modifications and measurably better performance.The feedback combines insights from torch, NCU, and Clang-tidy profilers, compared with a setup providing no explicit profiling feedback.

7 Related Work

Recent GPU programming frameworks balance high performance with developer productivity, providing abstractions for efficient kernel development and integration with machine learning libraries.

  • Modern GPU Programming: Frameworks such as Triton provide Python-based interfaces for writing efficient low-level GPU kernels without directly handling CUDA or OpenCL details, while ThunderKittens emphasizes usability and machine-learning-library integration.These tools reflect an ecosystem combining performance-oriented GPU programming with greater developer accessibility.

8 Discussion

The discussion argues that current LLM-written kernel benchmarks can produce artificial speedups that obscure performance assessment. It presents robust-kbench and an automated CUDA-kernel framework combining discovery, self-verification, optimization, evolutionary methods, soft verification, and hardware profiling.

  • Current LLM-written kernel benchmarks can produce artificial speedups, obscuring critical performance assessment.
  • robust-kbench was introduced to combat shortcomings in existing kernel benchmarks.
  • The proposed framework automates CUDA kernel discovery, self-verification, and optimization.
  • The approach combines LLMs with evolutionary optimization, soft-verification, and hardware profiling for CUDA-kernel translation and optimization.

Supplementary Material - NeurIPS 2025 Submission … A.4.1 Level 1 - Task 12 Matmul_with_diagonal_matric

The supplementary material describes filtering criteria for compromised KernelBench tasks and illustrates a CUDA kernel that achieves a 51.159x speedup over torch eager by addressing inefficient broadcasting. It also presents the task listings and implementation details for the diagonal-matrix multiplication example.

  • A.1 Filtering Procedure: KernelBench v0 tasks from the first two levels are filtered using criteria targeting artificial performance gains and benchmark quality.The analysis identifies vulnerabilities including inefficient baseline implementations and low-magnitude outputs.
  • A.1 Filtering Procedure: Tasks with outputs restricted to -0.01 to 0.01 or output standard deviation below 0.01 across model and input seeds are excluded.These criteria target low signal-to-noise ratios, floating-point correctness issues, and insufficient output variation.
  • A.2 Compromised Level 1 Tasks: Table 1 catalogs Level 1 KernelBench tasks with compromised properties.The table is presented as the supplementary listing of compromised Level 1 tasks.
  • A.3 Compromised Level 2 Tasks: Table 2 catalogs Level 2 KernelBench tasks with compromised properties.The table is presented as the supplementary listing of compromised Level 2 tasks.
  • A.4.1 Level 1 - Task 12 Matmul_with_diagonal_matric: 51.159x speedup over torch eager is achieved by the CUDA implementation because the baseline does not use efficient broadcasting.This example demonstrates artificial performance gains associated with a compromised baseline.
  • A.4.1 Level 1 - Task 12 Matmul_with_diagonal_matric: The diagonal-matrix kernel loads each row’s diagonal value into shared memory, uses 32x8 threads, and processes columns with tiled unrolling and masked writes.The implementation uses TILE_ROWS = 8, TILE_COLS = 32, UNROLL = 4, and a 256-thread block configuration.

A.4.2 Level 2 - Task 13 ConvTranspose3d_Mean_Add_Softmax · B Robust Kernel Benchmark

The ConvTranspose3d kernel reports a 123.603x speedup over torch eager, but its softmax implementation hardcodes 1.0 for a one-dimensional array, exposing contamination in the benchmark result.

  • A.4.2 Level 2 - Task 13 ConvTranspose3d_Mean_Add_Softmax: 123.603x speedup over torch eager is reported for the ConvTranspose3d_Mean_Add_Softmax kernel.The result is accompanied by a contamination finding: the softmax hardcodes 1.0 for a one-dimensional array.
  • A.4.2 Level 2 - Task 13 ConvTranspose3d_Mean_Add_Softmax: The kernel maps threads to output positions with a 2D grid and processes four consecutive X elements per thread.This mapping targets coalesced access along the X dimension.
  • A.4.2 Level 2 - Task 13 ConvTranspose3d_Mean_Add_Softmax: For each output position, the kernel computes convolution contributions across output channels and divides the accumulated value by C_out.It precomputes valid Z, Y, and X input ranges before traversing channels and kernel positions.
  • A.4.2 Level 2 - Task 13 ConvTranspose3d_Mean_Add_Softmax: The kernel applies softmax to a single value as 1.0, then applies tanh and scaling before writing the output.The passage explicitly identifies the single-value softmax behavior that underlies the contamination finding.
  • A.4.2 Level 2 - Task 13 ConvTranspose3d_Mean_Add_Softmax: Output writes use vectorized float4 stores when all four positions are valid and aligned, with scalar writes for boundary cases.This provides a fallback for incomplete or unaligned vector groups.
  • A.4.2 Level 2 - Task 13 ConvTranspose3d_Mean_Add_Softmax: The implementation uses 32×8 thread blocks, with 32 threads in X for warp alignment.The block configuration is described as promoting coalescing.

B.1 Task Overview

The benchmark tasks are organized by operation type and the features they support.

  • Table 3 categorizes benchmark tasks by operation type and supported features.

B.2 Linear Backward Kernel Task Definition · B.3 Linear Backward Kernel Task Configuration · C Additional Results

The paper defines a linear backward kernel task whose forward computation is y = x @ W^T + b and whose backward pass returns gradients for inputs, weights, and biases. Its configuration varies batch size, feature dimensions, and parameter initialization methods across single- and multi-input settings.

  • B.2 Linear Backward Kernel Task Definition: The custom backward method retrieves saved inputs and weights, then computes and returns gradients for the input, weights, and biases.The gradient computation is delegated through AutogradFunction.backward_fn.
  • B.2 Linear Backward Kernel Task Definition: The forward function implements a linear layer as y = x @ W^T + b, with bias broadcast across the batch dimension.The implementation delegates the computation to F.linear(x, weights, biases).
  • B.2 Linear Backward Kernel Task Definition: The model exposes configurable input features, output features, and initialization methods for the linear layer.Supported initialization methods shown include normal, kaiming, and xavier.
  • B.3 Linear Backward Kernel Task Configuration: The multi-input configuration includes batch_size 4 with num_output_features 10 and kaiming initialization.The configuration also lists num_output_features 4096 with xavier initialization for multi-input settings.
  • B.3 Linear Backward Kernel Task Configuration: Shared configurations test num_input_features values of 128 and 4096 across the task variants.Both feature sizes appear in the multi-shared configuration entries.

C.1 Detailed Translation Results on KernelBench · C.2 Generalization of Kernels on Different Hardware · C.3 More Baseline Comparisons against Our Evolutionary Approach

The sections report translation scaling, cross-GPU kernel speedups relative to PyTorch, and baseline comparisons between model-based methods and the evolutionary pipeline. Increasing proposals and incorporating error feedback improve translation efficacy, while the tables organize hardware and baseline evaluations.

  • C.1 Detailed Translation Results on KernelBench: Increasing the number of proposals generally improves torch-to-CUDA translation success.Figure 8 examines test-time scaling in the translation pipeline.
  • C.1 Detailed Translation Results on KernelBench: Iterative error feedback significantly outperforms parallel sampling in translation efficacy.The comparison holds for a similar number of proposals, with error feedback achieving higher success rates.
  • C.2 Generalization of Kernels on Different Hardware: Table 4 reports discovered-kernel speedups across different GPUs relative to PyTorch implementations.The table covers both forward (F) and backward (B) kernel operations.
  • C.2 Generalization of Kernels on Different Hardware: The cross-hardware evaluation distinguishes forward and backward kernel operations using the labels F and B.These labels are defined in the table caption accompanying the GPU speedup comparisons.
  • C.3 More Baseline Comparisons against Our Evolutionary Approach: Table 5 compares best-of-40 results from different model-based approaches with the evolutionary pipeline using 40 kernel proposals.The comparison is framed around model-based baselines and the evolutionary approach.
  • C.3 More Baseline Comparisons against Our Evolutionary Approach: The baseline comparison uses best-of-40 results as the reference for different model-based approaches.Table 5 contrasts these results with the evolutionary pipeline under the same 40-proposal budget.

D Kernel Evaluation Environment

The evaluation environment fixes H100 GPUs, CUDA 12.4, cuDNN 8.9.7, and specified package versions, while correctness and speedups are measured under standardized testing and timing procedures. Speedup estimates depend on evaluation settings, although qualitative kernel performance remains consistent across methods.

  • Evaluation environment: Experiments use H100 GPUs with CUDA 12.4, cuDNN 8.9.7, and the package versions listed in Table 6.The paper notes that speedup estimates can vary with evaluation-environment settings.
  • Correctness evaluation: A kernel is correct only if it passes tests across input sizes, initialization settings, and multiple random seeds using absolute and relative tolerances of 1e−5.Floating-point comparisons use both absolute and relative tolerance thresholds.
  • Performance evaluation: Speedups compare against native PyTorch for one input setting, except LayerNorm, which uses torch compile as the baseline.Runtime averages use 2000 measurements after 25 warmup runs.
  • Evaluation sensitivity: Re-evaluation shows qualitatively consistent kernel performance across evaluation settings, while quantitative speedups vary with the environment and comparison method.Figure 9 reports speedups over native PyTorch and Torch Compile in separate heatmap panels.

E Kernel Optimization Cost and Runtime Analysis · F Hyperparameter Settings · G Prompts

The paper details the computational cost and parallel workflow of CUDA kernel translation and optimization, then specifies hyperparameters and prompt designs for translation, optimization, verifier tuning, and multi-dimensional verification. These settings emphasize correctness, runtime improvement, structured outputs, and checks spanning compilation, memory, synchronization, coverage, and numerical stability.

  • E Kernel Optimization Cost and Runtime Analysis: LLM-driven optimization typically costs approximately $4–$5 per benchmark task across 40 kernel proposals, with CUDA and verifier APIs contributing substantially.The CUDA API generally accounts for the larger share, with backward optimization tasks among the most expensive cases.
  • E Kernel Optimization Cost and Runtime Analysis: Translation uses up to 10 sequential sampling, correctness-checking, and runtime-evaluation steps, averaging approximately 15 minutes on one GPU.Optimization instead uses 4 GPUs in parallel, sampling 8 proposals per generation for 10 generations; proposal sampling and self-verification each take approximately 1.5 minutes.
  • F Hyperparameter Settings: The optimization configuration uses o4-mini with temperature [1.0], 16,384 maximum tokens, high reasoning effort, 1 sample, and 20 verifier-tuning generations.Error summarization is enabled.
  • G.1 LLM-Driven CUDA Kernel Translation: Translation prompts instruct the model to produce exact forward or backward CUDA equivalents, include the required pybind11 module, and return code in <cuda></cuda> tags.Backward prompts additionally request efficient kernels, minimized torch-function usage, and concise kernel descriptors and implementation summaries.
  • G.2 LLM-Driven CUDA Kernel Optimization: Optimization prompts provide operation and hardware information, require correctness and reduced torch usage, and request a new CUDA kernel in a fixed name, description, and code schema.The optimization iteration explicitly targets improved speedup while preserving correct results.
  • G.3 LLM-Driven CUDA Kernel Verifier Tuning: Verifier-tuning prompts separately target compilation, numerical correctness, and memory safety, using feedback-driven prompt revisions to improve verifier accuracy.The prompts are assembled with the problem description, CUDA code, instructions, and parsing definition while excluding the verification answer itself.
  • G.4 LLM-Driven CUDA Kernel Verification: Tuned compilation verification emulates the NVIDIA toolchain through preprocessing, include resolution, symbol-table construction, kernel instantiation, semantic checks, PTXAS validation, resource checks, and linking.It returns success only without diagnostics; otherwise it reports toolchain-style errors and fixes.
  • G.4 LLM-Driven CUDA Kernel Verification: Tuned memory and numerical verification performs symbolic bounds, allocation, alignment, synchronization, tiling-coverage, unique-write, reduction, and numeric-stability analyses before producing structured diagnostics and a final verdict.The numerical verifier explicitly checks shared-memory initialization, barrier discipline, tail coverage, output-domain coverage, and NaN-free correctness.

H Highlighted CUDA Kernels: Analysis … I.12 Llama Feedforward Block Forward Kernel

The highlighted kernels fuse common neural-network operations and use CUDA-specific strategies including vectorized loads, warp reductions, shared memory, unrolling, and specialized launch configurations. The supplied code sections implement forward and backward kernels for MNIST models, while the analysis highlights fused LayerNorm and Llama feedforward operations.

  • H Highlighted CUDA Kernels: Analysis: LayerNorm uses vectorized float4 loads and warp-level reductions to compute statistics, then performs normalization and affine transformation in a second vectorized kernel.The implementation is described as a fused, two-step LayerNorm optimized for throughput and memory efficiency.
  • H Highlighted CUDA Kernels: Analysis: LlamaFFW uses a warp-predicated, divergence-free fused SiLUmultiply operation over float4-aligned data with uniform warp execution.The supplied passage describes the operation as fully uniform across warps and integrated within a larger feedforward implementation.
  • I.2 MNIST Linear-ReLU Forward Kernel: MNIST Linear-ReLU assigns each warp output pairs through a grid-stride loop, accumulates lane-wise dot products, reduces with shuffle instructions, and writes the biased ReLU result.The launch uses 4 warps per block and a tunable two-dimensional grid capped at 512 blocks per dimension.
  • I.3 MNIST Linear Forward Kernel: The standalone MNIST Linear forward kernel maps each output element to one warp, performs strided lane accumulation, reduces via warp shuffle, and adds bias before writing.Its wrapper launches a WARP_SIZE-thread block over a grid shaped as (n, m).
  • I.5 MNIST MaxPool Backward Kernel: The MaxPool backward kernel processes outputs with a grid-stride loop, reads each 2×2 window through read-only loads, identifies its maximum, and routes the upstream gradient without atomics.The wrapper zero-initializes grad_in and launches 256 threads per block with at most 65535 blocks.
  • I.6 MNIST Linear-ReLU Backward Kernel: The Linear-ReLU backward implementation uses separate kernels for masked gradients and bias gradients, weight gradients, and input gradients, returning all three tensors.The first kernel applies the ReLU mask and uses atomicAdd for bias accumulation; subsequent kernels compute the corresponding matrix-product gradients.
Loading 2509.14279v1…