用模板+搜索优化GPU内核,速度提升超3倍且更稳定
A Two-Stage GPU Kernel Tuner Combining Semantic Refactoring and Search-Based Optimization
- 将内核改写为可参数化的模板,再通过搜索自动调优
- 在真实内核上实现最高3倍以上加速,效果更稳定
- 适合需要自动化高性能优化的科研与工程场景
GPU代码优化是高性能计算及大模型训练推理中的关键性能瓶颈。尽管编译器优化和手工编写内核部分缓解了该问题,但接近硬件极限的性能仍高度依赖人工代码重构与参数调优。近年来基于大语言模型代理的内核生成与优化取得进展,但多数方法仅直接重写代码,参数选择常隐含且难以控制,或需人工干预,导致性能提升不稳定。本文提出在代理驱动的迭代流程之上引入基于模板的语义重构层:将内核语义重构为显式可参数化的模板,并通过基于搜索的自动调优优化参数,在硬件资源约束下实现高效搜索。实验以SGLang中的代表性CUDA内核为测试目标,所提代理调优器迭代执行模板化、测试、分析与规划,利用性能剖析反馈进行受限参数搜索。相比仅用代理直接重写的方法,模板+搜索设计显著降低优化过程的随机性,提升可解释性,实现更系统化的高性能配置。该方法还可扩展至OpenCL、HIP等后端,支持生产级工作负载的自动化性能优化。
原文摘要 · Abstract (English)
GPU code optimization is a key performance bottleneck for HPC workloads as well as large-model training and inference. Although compiler optimizations and hand-written kernels can partially alleviate this issue, achieving near-hardware-limit performance still relies heavily on manual code refactoring and parameter tuning. Recent progress in LLM-agent-based kernel generation and optimization has been reported, yet many approaches primarily focus on direct code rewriting, where parameter choices are often implicit and hard to control, or require human intervention, leading to unstable performance gains. This paper introduces a template-based rewriting layer on top of an agent-driven iterative loop: kernels are semantically refactored into explicitly parameterizable templates, and template parameters are then optimized via search-based autotuning, yielding more stable and higher-quality speedups. Experiments on a set of real-world kernels demonstrate speedups exceeding 3x in the best case. We extract representative CUDA kernels from SGLang as evaluation targets; the proposed agentic tuner iteratively performs templating, testing, analysis, and planning, and leverages profiling feedback to execute constrained parameter search under hardware resource limits. Compared to agent-only direct rewriting, the template-plus-search design significantly reduces the randomness of iterative optimization, making the process more interpretable and enabling a more systematic approach toward high-performance configurations. The proposed method can be further extended to OpenCL, HIP, and other backends to deliver automated performance optimization for real production workloads.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。