Technical Summary: HTAM - Hierarchical Transition-Attended Memory for Operator Optimization

1. Problem Statement

The rapid scaling of Large Language Models (LLMs) and the evolution of specialized accelerator architectures (e.g., GPUs) have created a critical demand for high-performance operator development. While LLMs have demonstrated strong code generation capabilities, optimizing GPU kernels remains a hardware-aware search problem that is difficult to solve with one-shot generation.

The core challenge lies in the granularity mismatch of existing LLM-based optimization methods:

• Coarse-grained approaches (e.g., textual hints or bottleneck diagnoses) are reusable and easy to retrieve but are often too abstract to guide the generation of executable CUDA code.
• Fine-grained approaches (e.g., detailed implementation traces) provide actionable details but create a large, noisy search space, making retrieval inefficient and obscuring optimization bottlenecks.

Furthermore, kernel optimization requires coupled multi-stage decisions (selecting tile sizes, thread mappings, memory patterns, etc.) and the non-trivial translation of these strategies into correct, hardware-aware code. Existing methods struggle to organize optimization experience to guide this multi-step search while providing concrete evidence for code edits.

2. Methodology: HTAM Framework

To address the granularity mismatch, the paper proposes HTAM (Hierarchical Transition-Attended Memory), a coarse-to-fine framework that organizes optimization experience into a structured memory bank. HTAM mirrors the decision process of human experts: first selecting a high-level global direction, then refining it with concrete local strategies.

2.1 Hierarchical Transition Graph (HTG)

HTAM structures its memory as a Hierarchical Transition Graph $M = (V_G, V_{GL}, E_G)$:

• Global Nodes ($V_G$): Represent coarse optimization directions (e.g., "Memory Access," "Data Reuse," "Throughput"). These nodes store high-level intent, static priors (goals, triggers, risks), and runtime statistics (success rates, improvement trends).
• Local Nodes ($V_{GL}$): Represent specific, executable strategies under a parent global node (e.g., "aligned vectorized loads" under "Memory Access"). These nodes store implementation-level knowledge, including edit patterns, verification checks, and positive/negative evidence from past executions.
• Transition Edges ($E_G$): Connect global nodes, storing path-dependent knowledge about transitioning from one global direction to another. These edges record when a transition is applicable, its rationale, risks, and trajectory-level evidence (performance gains, compilation outcomes).

2.2 Optimization Process

The optimization proceeds as a $T$-step evolution process where the implementation $c_t$ is updated iteratively:

1. State Summarization: The current implementation and history are summarized into a compact state $s_t$ containing task context, feedback (compile, correctness, runtime), and recent optimization paths.
2. Global Direction Selection: HTAM selects a global direction $g_t$ based on the current state and recent history. It uses a transition-aware scoring mechanism that aggregates evidence from the recent prefix of global directions ($p^g_t$) using position-aware weights (inspired by attention mechanisms). This avoids the high cost of storing full-chain memories while capturing multi-step dependencies.
   • The score is calculated as a weighted sum of transition utilities from previous nodes to the candidate node.
   • A Softmax-based policy balances exploitation of high-reward transitions with exploration.
3. Local Strategy Selection: Once a global direction is chosen, the search space is restricted to associated local nodes. An LLM evaluates candidate local strategies based on the state, current implementation, and the specific global/local node memories to select the best strategy $l_t$.
4. Implementation & Update: The LLM generates the next implementation $c_{t+1}$ guided by the selected local strategy. Evaluation feedback is then used to update the HTG: global node statistics, local node evidence, and transition edge feedback are refreshed to inform future steps.

3. Key Contributions

• Coarse-to-Fine Framework: HTAM decomposes LLM-based operator optimization into global direction selection and local strategy refinement, addressing the granularity mismatch in existing methods.
• Structured Memory Organization: It introduces a Hierarchical Transition Graph that organizes experience into global nodes, local nodes, and transition edges. This structure reduces the search space and storage complexity compared to full-chain memory or flat retrieval.
• Transition-Aware Decision Making: The framework employs a scoring mechanism that aggregates recent optimization history via edge memories, guiding multi-step decisions without the exponential cost of storing full trajectories.
• Empirical Validation: Extensive experiments demonstrate that HTAM significantly outperforms baselines in correctness, speed of finding fast solutions, and geometric mean speedup.

4. Experimental Results

Experiments were conducted on KernelBench, a suite of 250 tasks (Level-1 single operators, Level-2 fused operators, Level-3 model-level tasks), using DeepSeek-R1 as the backbone model on NVIDIA A100 GPUs.

• Performance Gains: HTAM achieved 98.4% correctness, 84.0% Fast@1 (percentage of tasks with a correct and faster solution on the first try), and a 1.978× geometric mean speedup.
• Comparison to Baselines:
  • Compared to the strongest reproduced baseline (Flat Memory Retrieval), HTAM improved correctness by 8.8 percentage points, Fast@1 by 25.2 percentage points, and speedup from 1.464× to 1.978×.
  • It significantly outperformed vanilla LLM generation and feedback-only refinement methods.
• Generalization:
  • Backend Transfer: HTAM remained effective across different LLM backends (DeepSeek-R1, GPT-4o, DeepSeek-V4-Flash), though performance correlated with the backend's code generation capability.
  • Cross-Benchmark Transfer: When initialized with memory built from KernelBench and applied to Robust-KBench tasks, HTAM showed improved speedups (e.g., +1.76× for LayerNorm), indicating transferable optimization experience, though gains were operator-dependent.
• Ablation Studies: Removing the hierarchical structure or local memory caused significant drops in performance (e.g., speedup dropping to 0.974× without local memory), confirming that structured global-local experience is critical. Similarly, removing prefix-aware transition scoring reduced performance, validating the need for history-dependent decision making.

5. Significance and Claims

The paper claims that structured memory serves not merely as a repository of past cases but as a decision prior for selecting the next useful optimization direction. By organizing experience hierarchically and modeling transitions between optimization steps, HTAM enables LLMs to navigate the complex, multi-stage search space of GPU kernel optimization more effectively than flat retrieval or one-shot generation.

The authors emphasize that HTAM provides a practical path to automating high-performance operator development, achieving near-expert levels of correctness and speedup while maintaining a manageable online API cost (approximately $0.21 per task for a 6-step optimization using DeepSeek-R1).

Limitations Acknowledged:

• The evaluation is limited by computational and API costs, focusing primarily on KernelBench and a subset of Robust-KBench.
• The system relies on executable validation (compilation, correctness, runtime) rather than formal verification, meaning correctness is not guaranteed for all possible shapes or data distributions.
• Generalization to substantially different hardware or operator families may require adapting the memory schema.