arXiv:2605.29734cs.CL2026-05被引 1

用分层记忆提升大模型生成高效GPU代码的准确性与速度

HTAM: Hierarchical Transition-Attended Memory for Operator Optimization

论文配图:HTAM: Hierarchical Transition-Attended Memory for Operator Optimization
图 1 · 摘自论文原文
  • 构建两级分层转移图,整合全局方向与局部策略
  • 在KernelBench上提升正确率、快速求解率和加速比
  • 适合需要自动化优化GPU算子的研发人员

高性能GPU内核对大语言模型部署至关重要,但其优化仍依赖专家经验。基于大模型的代码生成虽具潜力,但算子优化仍是需考虑硬件特性的搜索问题。现有方法存在粒度不匹配:粗粒度提示可复用但难执行,细粒度记忆可操作但扩大搜索空间并掩盖瓶颈。本文提出HTAM(分层转移注意力记忆)框架,通过两级分层转移图组织全局优化方向、局部策略及步骤间转移经验。在每个演化步骤中,HTAM从当前状态与近期历史中选择全局方向,检索对应局部策略记忆,指导具体CUDA代码生成。在完整KernelBench套件上的实验表明,HTAM在正确性、快速求解率和加速比方面均优于基线方法;后端与Robust-KBench研究也验证了结构化记忆的可迁移优势。

原文摘要 · Abstract (English)

High-performance GPU kernels are essential for efficient LLM deployment, yet optimizing them remains expertise-intensive. Recent LLM-based code generation makes automatic GPU operator generation promising, but operator optimization remains a hardware-aware search problem. Existing LLM-based methods face a granularity mismatch: coarse hints are reusable but hard to execute, whereas detailed memories are actionable but enlarge the search space and obscure optimization bottlenecks. The key challenge is therefore to organize optimization experience at an appropriate granularity. To address this issue, this paper proposes HTAM (Hierarchical Transition-Attended Memory), a coarse-to-fine framework for LLM-based operator optimization. HTAM builds a two-level Hierarchical Transition Graph (HTG) to organize coarse global directions, detailed local strategies, and transition experience between optimization steps. During each evolution step, HTAM selects a global direction from the current state and recent optimization history, retrieves the corresponding local strategy memory, and uses it to guide concrete CUDA code generation. Experiments on the full KernelBench suite demonstrate that HTAM consistently improves correctness, fast-solution rate, and speedup over LLM-based baselines, while backend and Robust-KBench studies indicate transferable benefits from structured memory.

GPU优化大模型生成分层记忆CUDA

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。