Compiler-Grounded Hierarchical Diagnosis for LLM-Based Triton Kernel Optimization

Technical Summary: Compiler-Grounded Hierarchical Diagnosis for LLM-Based Triton Kernel Optimization

1. Problem Statement

While Large Language Models (LLMs) have advanced automated kernel generation and optimization, existing approaches predominantly rely on "surface signals" such as compilation errors, runtime metrics, and profiling data. These signals indicate that a kernel is slow but rarely explain why the backend compiler fails to realize a profitable optimization. This is particularly acute on emerging accelerators like NPUs (e.g., Ascend), where the compiler stack is evolving, target-specific, and opaque to developers.

The core challenge identified is a semantic gap between operator-level source code and backend performance behavior. Performance failures often emerge only after the code passes through intermediate representation (IR) stages, scheduling decisions, and backend-specific compiler passes. Consequently, LLM agents, typically pretrained on GPU-centric patterns, behave as trial-and-error optimizers: they propose plausible source edits but lack the evidence to map runtime symptoms to specific compiler-level causes.

2. Methodology

The paper proposes a compiler-grounded, hierarchical optimization framework that treats kernel optimization as a progressive cross-layer diagnosis problem. Instead of treating compiler grounding as an always-on feature, the system escalates from lightweight evidence to deep compiler analysis only when necessary.

The Hierarchical Workflow:
The system operates through a closed-loop pipeline organized into four escalating levels:

• L1: Pattern Triage: The system first checks for strong, pattern-backed hypotheses using a curated index of optimization patterns, heuristics, and expert rules (including internal heuristics distilled from prior runs). If a low-risk rewrite is justified by existing evidence, the process stops here.
• L2: Profiling Diagnosis: If patterns fail, the system gathers runtime evidence (latency, throughput, hardware counters) to localize bottlenecks (e.g., memory movement, scalar activity, synchronization). This identifies the symptom but not necessarily the mechanism.
• L3: IR Attribution: When profiling is insufficient, the system inspects Intermediate Representations (IR) to translate runtime symptoms into structural diagnoses (e.g., non-unit-stride loads, missed fusion, layout-induced vectorization failures). It reconstructs the operator's lowering trajectory to identify stage-to-stage transformations.
• L4: Compiler-Source Escalation: As a last resort, the system invokes compiler-grounded analysis to explain why a diagnosed pattern blocks backend optimization. It consults compiler knowledge (pass behaviors, lowering rules, constraints) to recover the source-level conditions a successful rewrite must satisfy.

System Architecture & State Management:

• Agent-Centric Design: Implemented as an agent-centric environment (integrated into the CannBot skills repository) rather than a bespoke compiler service. It uses a thin execution harness to manage validation, correctness, and evidence collection, while domain-specific "skills" handle the optimization logic.
• Stateful Rounds: The optimization process is organized into explicit, auditable rounds. Each round records its hypothesis, evidence sources, and validation outcome. This stateful design allows the system to be resumable and ensures that deeper analyses (IR/Compiler) are attached as reusable evidence rather than ephemeral observations.
• Retrospective Synthesis: Beyond the online loop, the system performs a slower, cross-operator synthesis pass. It maps validated optimization histories to existing pattern families or creates new reusable patterns, allowing the "cheapest" level (Pattern Triage) to improve over time.

3. Key Contributions

1. Problem Formulation: The authors formulate LLM-based kernel optimization as a progressive cross-layer diagnosis problem, identifying the lack of compiler-grounded evidence as a primary reason why profiling-only agents fail on backend-specific optimizations.
2. Hierarchical Workflow: They describe a workflow that escalates from pattern triage to profiling, IR attribution, and compiler-source escalation, proposing evidence-backed source-level rewrites only when deeper evidence is required.
3. Evaluation Evidence: They present standardized, distribution-aware evaluation evidence on a NPUKernelBench-derived Ascend 950 benchmark, reporting aggregate gains, best-round accounting, and representative cross-layer traces.

4. Experimental Results

The system was implemented on Triton for Ascend NPUs and evaluated on 37 successfully converted operators from a standardized Ascend 950 benchmark (derived from NPUKernelBench). Each operator was tested on 5 cases (185 total cases).

• Performance Gains:
  • Geometric Mean Speedup: 4.35× (from initial to optimized Triton kernel).
  • Median Speedup: 2.73×.
  • Distribution: 22/37 entries exceeded 2× speedup; 13/37 exceeded 5×. The results show a long-tailed distribution, with substantial gains coexisting with near-baseline entries.
• Process Accounting:
  • The median best-round index was 10, indicating that optimal results often required multiple iterations.
  • Evidence Depth: While 33 of the selected "best" rounds were attributed to Pattern Triage, the system demonstrated the utility of deeper layers: 32 trajectories reached Profiling Diagnosis, 3 reached IR Attribution, and 4 reached Compiler-Source Escalation.
  • Case Study: A representative trace (Hyena FFT-size padding) showed how profiling identified a scalar bottleneck, IR attribution identified specific synchronization costs in the lowering, and compiler-source analysis explained the constraint, leading to a robust 2.08× speedup.

5. Significance and Claims

The paper claims that compiler grounding should be selective rather than always-on. By treating optimization as a progressive diagnosis, the system makes deep compiler reasoning practical, invoking it only when lighter-weight evidence cannot produce a sufficiently grounded rewrite.

The authors emphasize that their results demonstrate substantial headroom on a subset of operators but do not claim uniform improvement across all cases. They explicitly state that the current evaluation does not provide causal isolation (e.g., comparing progressive escalation against profiling-only or always-deep strategies under a frozen toolchain). The paper positions itself as a "systems argument" supported by transparent benchmark and process evidence, reserving causal decomposition of the hierarchy for future controlled experiments.

The work highlights that the relationship between source code and runtime performance is mediated by invisible compiler transformations, and that bridging this gap with a hierarchical, evidence-driven approach is essential for optimizing kernels on emerging, opaque accelerators like NPUs.