Saved in:
| Main Authors: | , , , , , , , |
|---|---|
| Format: | Preprint |
| Published: |
2026
|
| Subjects: | |
| Online Access: | https://arxiv.org/abs/2605.29734 |
| Tags: |
Add Tag
No Tags, Be the first to tag this record!
|
| _version_ | 1866914613299249152 |
|---|---|
| author | Zhang, Yining Yi, Mingyang Wang, Chen Xiang, Xuwen Jia, Tianhe Dan, Zedong Zong, Chengqing Wang, Yue |
| author_facet | Zhang, Yining Yi, Mingyang Wang, Chen Xiang, Xuwen Jia, Tianhe Dan, Zedong Zong, Chengqing Wang, Yue |
| contents | 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. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2605_29734 |
| institution | arXiv |
| publishDate | 2026 |
| record_format | arxiv |
| spellingShingle | HTAM: Hierarchical Transition-Attended Memory for Operator Optimization Zhang, Yining Yi, Mingyang Wang, Chen Xiang, Xuwen Jia, Tianhe Dan, Zedong Zong, Chengqing Wang, Yue Computation and Language 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. |
| title | HTAM: Hierarchical Transition-Attended Memory for Operator Optimization |
| topic | Computation and Language |
| url | https://arxiv.org/abs/2605.29734 |