论文
HTAM：用于操作员优化的分层转换参与内存
HTAM: Hierarchical Transition-Attended Memory for Operator Optimization
摘要
高性能 GPU 内核对于高效 LLM 部署至关重要，但优化它们仍然需要大量专业知识。最近基于 LLM 的代码生成使得自动 GPU 算子生成很有前景，但算子优化仍然是一个硬件感知的搜索问题。现有的基于 LLM 的方法面临着粒度不匹配的问题：粗略提示可重用，但难以执行，而详细记忆是可操作的，但会扩大搜索空间并模糊优化瓶颈。因此，关键的挑战是以适当的粒度组织优化经验。为了解决这个问题，本文提出了 HTAM（分层过渡参与内存），这是一种用于基于 LLM 的算子优化的从粗到细的框架。 HTAM 构建了一个两级层次转换图（HTG）来组织粗略的全局方向、详细的局部策略以及优化步骤之间的转换经验。在每个进化步骤中，HTAM从当前状态和最近的优化历史中选择一个全局方向，检索相应的局部策略内存，并用它来指导具体的CUDA代码生成。在完整 KernelBench 套件上进行的实验表明，与基于 LLM 的基线相比，HTAM 持续提高了正确性、快速解决率和加速速度，而后端和 Robust-KBench 研究表明结构化内存具有可转移的优势。