Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Integrating Performance Tools in Model Reasoning for GPU Kernel Optimization

Daniel Nichols, Konstantinos Parasyris, Charles Jekel, Abhinav Bhatele, Harshitha Menon

TL;DR

This work tackles the hard problem of GPU kernel optimization by enabling a reasoning-capable language model to interact with runtime performance tools during its inference. By fine-tuning a large model with Group Relative Policy Optimization (GRPO) and verifiable rewards, the system learns to use compiler diagnostics, profilers, and benchmarks to iteratively propose and validate kernel edits. It combines a large 70B-parameter model for training with a distilled 8B model for practical deployment, achieving kernel speedups and demonstrable end-to-end improvements in a shock-hydrodynamics application. The results show that tool-informed reasoning can surpass traditional code edits, offering a scalable path toward automated, verifiable performance optimization in HPC workloads.

Abstract

Language models are now prevalent in software engineering with many developers using them to automate tasks and accelerate their development. While language models have been tremendous at accomplishing complex software engineering tasks, there are still many areas where they fail to deliver desirable results, for instance code performance related tasks. Tasks like optimization depend on many complex data from the environment, hardware, etc. that are not directly represented in source code. Recent efforts have seen large improvements in general code modeling tasks using chain-of-thought style reasoning, but these models still fail to comprehend how the environment interacts with code performance. In this paper we propose a methodology to train language models that can interact with performance tools during their reasoning process. We then demonstrate how this methodology can be used to train a state-of-the-art GPU kernel optimization model.

Integrating Performance Tools in Model Reasoning for GPU Kernel Optimization

TL;DR

This work tackles the hard problem of GPU kernel optimization by enabling a reasoning-capable language model to interact with runtime performance tools during its inference. By fine-tuning a large model with Group Relative Policy Optimization (GRPO) and verifiable rewards, the system learns to use compiler diagnostics, profilers, and benchmarks to iteratively propose and validate kernel edits. It combines a large 70B-parameter model for training with a distilled 8B model for practical deployment, achieving kernel speedups and demonstrable end-to-end improvements in a shock-hydrodynamics application. The results show that tool-informed reasoning can surpass traditional code edits, offering a scalable path toward automated, verifiable performance optimization in HPC workloads.

Abstract

Language models are now prevalent in software engineering with many developers using them to automate tasks and accelerate their development. While language models have been tremendous at accomplishing complex software engineering tasks, there are still many areas where they fail to deliver desirable results, for instance code performance related tasks. Tasks like optimization depend on many complex data from the environment, hardware, etc. that are not directly represented in source code. Recent efforts have seen large improvements in general code modeling tasks using chain-of-thought style reasoning, but these models still fail to comprehend how the environment interacts with code performance. In this paper we propose a methodology to train language models that can interact with performance tools during their reasoning process. We then demonstrate how this methodology can be used to train a state-of-the-art GPU kernel optimization model.
Paper Structure (39 sections, 4 equations, 4 figures, 3 tables, 1 algorithm)

This paper contains 39 sections, 4 equations, 4 figures, 3 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Background
  3. Fine-tuning LMs with RL
  4. Test-time compute and LM Reasoning
  5. Chain-of-thought (CoT) and Reasoning
  6. Cost–quality trade-offs
  7. Overview
  8. GPU Kernel Data Collection
  9. Dataset Requirements
  10. HeCBench GPU Kernels
  11. Synthetic Kernels
  12. Model Design: Reasoning with Optimization Tools
  13. Model Inputs, Outputs, and Reasoning
  14. Tools Available to the Model
  15. An Example Trace
  16. ...and 24 more sections

Figures (4)

  • Figure 1: The distribution of speedups for each model across the kernel types in the RAJAPerf suite. A speedup of 0 denotes a failed result (i.e. incorrect or has syntax errors). The opt-r1 and opt-r1-mini models are consistently able to improve kernel performance across most benchmark problems.
  • Figure 2: Four traces of benchmark_code tool call results during optimization of the FEMSweep RAJAPerf kernel by opt-r1. Each plot shows the speedup of progressive invocations of the benchmark tool during the model's reasoning process.
  • Figure 3: The distribution of speedups on the RAJAPerf kernels from opt-r1 and opt-r1-mini with specified tools removed. For example, the last bar shows the distribution of speedups where the microbench_code tool is not available. We observe that the benchmark_code tool is crucial to consistently getting valid optimizations from the models.
  • Figure 4: Tool use counts per reasoning trace across all traces from the RAJAPerf evaluation. benchmark_code, file_viewer, and search are the most widely used, while most traces rarely use the microbench_code tool more than once.