Source-linked AI summary
HIERA: Workload-Aware Planning Across Implementation Spaces for GPU Kernel Optimization
Jinghao Wang, Qiqi Gu, Chenpeng Wu, Jianguo Yao, Haibing Guan, Xijun Li
TL;DR
LLM-based GPU kernel optimization is limited by fixed implementation spaces as workloads and hardware diversify. HIERA constructs contract-augmented specifications, selects workload-appropriate spaces, and guides refinement with profiling and expert knowledge. Across KernelBench workloads and base models, it improves validity, sample efficiency, and optimization performance over training-free methods while remaining competitive with CUDA-L1 without additional training.
Problem
Existing LLM-based GPU optimization methods typically use fixed implementation spaces, limiting flexibility or search efficiency as workloads diversify.
Method
HIERA combines contract-augmented task specifications with hierarchical implementation-space planning and feedback-driven refinement guided by profiling and expert knowledge.
Results
HIERA achieves the best or tied-best result in 22 of 27 comparisons and ranks first in fast1 across workload levels and base models among inference-time methods.
Takeaways & Limitations
HIERA provides structured coarse-to-fine optimization across heterogeneous workloads while balancing implementation reliability and optimization flexibility under limited search budgets.
Takeaways & Limitations
Experiments use NVIDIA A100 GPUs, and the scientific-computing case study covers only one stencil operator and configuration.
Abstract
from arXiv · showhide
High-performance GPU kernels underpin modern deep learning and scientific computing. As workloads become increasingly diverse and GPU hardware evolves rapidly, developing efficient methods for automated GPU kernel generation and optimization has become increasingly important. Existing LLM-based methods typically optimize within a fixed implementation space, limiting either optimization flexibility or search efficiency. We propose \textsc{HIERA}, a hierarchical search-space planning framework for GPU kernel optimization. \textsc{HIERA} constructs contract-augmented task specifications, selects an appropriate implementation space across PyTorch operators, CUDA libraries, and custom CUDA kernels, and uses profiling feedback and expert knowledge to guide structured iterative refinement. Experiments on KernelBench across multiple various workload levels and base LLMs show that \textsc{HIERA} delivers stronger overall implementation validity, sample efficiency, and optimization performance than existing training-free methods, while remaining competitive with the training-based CUDA-L1 without additional model training. A case study on a specialized stencil operator from scientific computing further achieves a \(1.53\times\) speedup over cuDNN, demonstrating the potentiality of the general framework beyond standard machine-learning workloads.
Introduction
GPU kernel optimization remains difficult for diverse workloads because existing abstractions and LLM methods rely on fixed implementation spaces. HIERA makes implementation-space selection workload-aware and combines it with structured refinement.
- Motivation: Specialized workloads and evolving hardware leave GPU kernel development dependent on substantial manual expertise and engineering effort.Existing libraries and framework primitives do not consistently cover increasingly diverse workload requirements.
- Motivation: Existing LLM-based methods commonly optimize within predetermined implementation spaces, trading optimization flexibility against search efficiency.High-level operators and libraries preserve semantics, while custom CUDA offers freedom but exposes a large, weakly structured decision space.
- Motivation: Under limited candidate budgets, low-level exploration can waste resources on implementations that fail compilation or correctness checks or provide little speedup.The problem is especially acute when models must navigate many interacting low-level decisions.
- HIERA: HIERA treats implementation-space selection as an explicit workload-aware optimization decision.It selects an appropriate space before structuring refinement within that space using profiling feedback and expert domain knowledge.
- HIERA: HIERA combines contract-augmented specifications with hierarchical planning across PyTorch operators, CUDA libraries, and custom CUDA kernels.The specification fixes interfaces, parameter semantics, compilation rules, and verification behavior, while planning guides refinement across implementation regimes.
Related Work
Prior systems improve GPU kernel optimization through compilers, benchmarks, feedback, or training, but generally remain tied to predefined representations or implementation spaces. HIERA instead plans implementation granularity explicitly and refines candidates within a workload-appropriate space.
- Tensor Compilers and GPU DSLs: Tensor compilers and GPU DSLs provide structured optimization spaces but remain constrained by predefined compiler representations, scheduling languages, or DSL models.TVM and Ansor search compiler schedules and transformations, while Triton provides a high-level custom-kernel language and compiler.
- GPU Kernel Generation Benchmarks: KernelBench evaluates whether generated GPU implementations are functionally correct and faster than their PyTorch references.The benchmark contains 250 workloads across three levels and reports speedup-threshold validity through fastp.
- Agentic, Feedback-Driven CUDA Optimization: CUDAForge uses correctness results, GPU specifications, and Nsight Compute metrics for iterative refinement but restricts exploration to a fixed custom-CUDA space.This restriction can be inefficient for composite and model-level workloads that require reconstructing dependencies and intermediate data flows.
- Reinforcement Learning for CUDA Optimization: Reinforcement-learning methods encode optimization preferences in model parameters and require training, whereas HIERA uses inspectable planning decisions without retraining.The comparison concerns how optimization knowledge is represented and applied.
- HIERA: HIERA combines contract-augmented specifications, workload-aware cross-granularity planning, and expert-guided optimization-direction pruning.This design enables structured coarse-to-fine optimization across heterogeneous workloads.
Method
HIERA formulates implementation-space selection as a workload-aware planning decision, then guides iterative kernel refinement with contracts, profiling feedback, and expert knowledge. Its pipeline prunes optimization choices, verifies generated candidates, and feeds measured results into subsequent iterations.
- Contract-Augmented Specification: HIERA converts each task into a contract-augmented specification that fixes interfaces, compilation rules, and parameter semantics while focusing generation on optimization targets.Fixed wrappers, bindings, and correctness references reduce the need to regenerate task infrastructure.
- Workload-Aware Planning: Implementation-space selection is explicit because high-level composition can incur dispatch and materialization overheads, whereas custom CUDA may consume budget reconstructing complex workload logic.The appropriate granularity therefore depends strongly on workload characteristics.
- Workload-Aware Planning: At each planning step, SSDA selects among nested pure-CUDA, CUDA-library, and CUDA-library-plus-PyTorch spaces with increasing permissiveness.The selected space defines the admissible implementation choices for the current refinement step.
- Direction Planning: Within the selected space, SSDA scores optimization directions using the task specification, current candidate, profiling feedback, expert knowledge, and a fixed rubric.The direction taxonomy covers control-flow specialization, parallelism, memory transactions, reuse and pipelining, and Tensor Core or instruction utilization.
- Structured Refinement: SPA translates the selected direction into concrete transformation steps and constraints, after which OA generates the next candidate kernel set within the admissible space and direction.This separates direction selection from strategy construction and candidate generation.
- Evaluation and Feedback: Candidates are compiled and correctness-checked before valid implementations are benchmarked and profiled with Nsight Compute for latency and hardware bottlenecks.Measured results and intermediate decisions are recorded, strong candidates are retained as parents, and feedback guides later refinement.
Experimental Setup
The evaluation spans KernelBench workloads, multiple base LLMs, fixed-space variants, ablations, limited candidate budgets, and a scientific-computing stencil case study. It uses standardized validity and speedup metrics under controlled protocols, with scope limited mainly to A100 GPUs and a single stencil configuration.
- Benchmark and Baselines: HIERA is evaluated on 250 KernelBench workloads spanning operator-level, fused-operator, and model-level tasks.The benchmark includes 100 Level-1, 100 Level-2, and 50 Level-3 workloads.
- Benchmark and Baselines: The comparison includes KernelBench-Caesar, CUDAForge, and the training-based CUDA-L1 across three base LLMs.The base models are DeepSeek-V3.2, Qwen3.6-Plus, and Gemini-3.6-Flash.
- Variants and Protocols: Implementation-space experiments compare Pure CUDA, CUDA Libraries, CUDA Libraries + PyTorch, and dynamically selected HIERA spaces on 90 sampled tasks.The sample contains 30 tasks from each of KernelBench Levels 1–3.
- Variants and Protocols: Ablations remove either contract augmentation or hierarchical planning to quantify the contributions of HIERA’s principal components.The contract ablation requires generating supporting files, while the planning ablation uses a fixed optimization prompt.
- Variants and Protocols: HIERA, KernelBench-Caesar, and CUDAForge generate up to B = 18 candidates through three refinement rounds with six candidates per round.Model versions, temperature 0.3, candidate budgets, and evaluation settings are aligned across methods.
- Variants and Protocols: Limited-budget results report the best valid implementation among the first B candidates at B ∈ {1, 6, 12, 18}.The budget points are cumulative and extracted from the same generation trajectories.
- Case Study: The scientific-computing case study uses a radius-3 2D box stencil with a 7 × 7 neighborhood, 10240 × 10240 input, and 50 candidate evaluations.Latency is amortized per step and compared with cuDNN convolution.
Results
HIERA improves optimization quality and sample efficiency across KernelBench workloads, with benefits from adaptive implementation-space selection, contract augmentation, and hierarchical planning. It also finds a stencil implementation that outperforms cuDNN.
- Limited-budget evaluation: At B = 18, HIERA reaches 92.4% fast0 and 60.4% fast1, exceeding KernelBench at 85.2%/29.6% and CUDAForge at 85.6%/44.8%.At B = 1, it already achieves 71.6% fast0 and 32.4% fast1.
- Overall performance: HIERA achieves the best or tied-best result in 22 of 27 comparisons against LLM-conditioned baselines and outperforms CUDA-L1 on averaged Level 1 and Level 2 metrics.It ranks first among inference-time methods in fast1 across workload levels and base models.
- Implementation-space selection: 1.42× mean and 1.15× median speedups make HIERA the strongest aggregate result, while its variance is 23.4% lower than Pure CUDA.Pure CUDA reaches a 9.32× maximum but has a 1.00× mean, 0.62× median, and variance of 1.79.
- Ablation study: Removing contract augmentation reduces validity and optimization performance more severely as workload complexity increases, including a 73-point fast0 drop on Level 2.On Level 1, fast0/fast1/fast2 fall from 97/68/18% to 86/40/0%.
- Ablation study: Removing hierarchical planning sharply reduces acceleration while only modestly affecting validity, with Level 2 fast1/fast2 falling from 62/20% to 6/6%.This indicates generic refinement can produce valid implementations without reliably finding productive implementation spaces or optimization directions.
- Stencil case study: 4.71 ms in the fifth iteration is 34.8% below cuDNN's 7.23 ms reference, corresponding to a 1.53× speedup on the 2D box stencil.The runtime decreases from 13.71 ms in the first iteration to 4.71 ms in the fifth.
Conclusion
HIERA explicitly plans implementation spaces and optimization directions for GPU kernel optimization, combining contract-augmented specifications with feedback-driven refinement.
- HIERA balances implementation reliability and optimization flexibility under limited search budgets.The framework combines contract-augmented task specifications with feedback-driven refinement.
- Future work will extend HIERA across GPU architectures, precision formats, and implementation backends.
- Future work will also explore self-improving planning from accumulated optimization trajectories.