From Large to Small: Transferring CUDA Optimization Expertise via Reasoning Graph

TL;DR

This paper tackles the difficulty of leveraging GPUs for CUDA programming in privacy-sensitive settings by introducing ReGraphT, a training-free framework that transfers LLM-like multi-step reasoning to small language models. It builds a CUDA Reasoning Graph from LLM-generated optimization trajectories and employs Monte Carlo Graph Search to guide SLMs through complex CUDA code-generation tasks. A dedicated CUDAEval benchmark with difficulty tiers based on reasoning trajectories enables fine-grained evaluation. Empirical results show that ReGraphT, especially with MCGS, delivers substantial performance gains over prompting and retrieval-based baselines, bringing SLMs closer to LLM-level capabilities while preserving privacy and reducing compute overhead. The approach highlights the potential of structured reasoning graphs and retrieval-guided exploration for domain-specific, multi-step code optimization in real-world CUDA workloads.

Abstract

Despite significant evolution of CUDA programming and domain-specific libraries, effectively utilizing GPUs with massively parallel engines remains difficult. Large language models (LLMs) show strong potential in generating optimized CUDA code from sequential code. However, using LLMs in practice faces two major challenges: cloud-based APIs pose risks of code leakage, and local deployment is often computationally expensive and inefficient. These drawbacks have spurred interest in small language models (SLMs), which are more lightweight and privacy-friendly. Encouragingly, recent studies show that SLMs can achieve performance comparable to LLMs on specific tasks. While SLMs can match LLMs on domain-specific tasks, their limited reasoning abilities lead to suboptimal performance in complex CUDA generation according to our experiments. To bridge this gap, we propose ReGraphT, a training-free, retrieval-augmented generation framework that transfers LLM-level reasoning to smaller models. ReGraphT organizes CUDA optimization trajectories into a structured reasoning graph, modeling the combined CUDA optimizations as state transitions, and leverages Monte Carlo Graph Search (MCGS) for efficient exploration. We also present a CUDA-specific benchmark with difficulty tiers defined by reasoning complexity to evaluate models more comprehensively. Experiments show that ReGraphT outperforms HPC-specific fine-tuned models and other retrieval-augmented approaches, achieving an average 2.33X speedup on CUDAEval and ParEval. When paired with DeepSeek-Coder-V2-Lite-Instruct and Qwen2.5-Coder-7B-Instruct, ReGraphT enables SLMs to approach LLM-level performance without the associated privacy risks or excessive computing overhead.

Figures (11)

• Figure 1: Average number of reasoning steps and the performance of the generated code with SLMs and LLMs.
• Figure 2: Overview of the proposed ReGraphT framework.
• Figure 3: ReGraph construction based on LLM optimization trajectory.
• Figure 4: An overview of MCGS on ReGraph.
• Figure 5: CUDAEval curation process.

Theorems & Definitions (1)

• Definition 1