Skip to content

PRAGMA: A Profiling-Reasoned Multi-Agent Framework for Automatic Kernel Optimization

🕒 Published (v1): 2025-11-09 12:01 UTC · Source: Arxiv · link

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

PRAGMA is a multi-agent LLM framework for automatic GPU/CPU kernel optimization that closes the generation-verification-profiling loop using hardware performance counters (Nsight Compute, Linux Perf) rather than coarse execution-time feedback alone. A Conductor Agent interprets fine-grained profiling metrics, classifies bottlenecks, and directs a Coder Agent to iteratively refine kernels while retaining the historically best implementation. On KernelBench, PRAGMA achieves 2.81× and 2.30× average speedup over PyTorch on CPU and GPU respectively.

Problem

Existing LLM-based kernel generation systems (e.g., Caesar) rely solely on correctness checks or overall runtime as feedback, providing no visibility into hardware-level bottlenecks such as memory throughput, occupancy, or instruction-level inefficiencies. This coarse feedback causes iterative refinements to be heuristic and unstable, preventing LLMs from reasoning about why a kernel is slow.

Method

PRAGMA organizes four agents in a closed feedback loop:

  1. Coder Agent — generates and incrementally refines Triton (GPU) or C++ (CPU) kernel implementations based on Conductor instructions.
  2. Verifier Agent — compiles and runs each candidate, comparing outputs against a reference; passes error logs to the Conductor. Non-LLM component.
  3. Profiler Agent — collects hardware metrics via NVIDIA Nsight Compute (GPU: occupancy, register usage, memory throughput, tensor core cycles) or Linux Perf (CPU: top-down microarchitecture metrics — front-end/back-end bound, retired instructions). Raw profiler documentation is auto-crawled and LLM-summarized into an offline compendium that the Conductor uses for semantic reasoning about metric meaning.
  4. Conductor Agent — aggregates Coder output, Verifier logs, and Profiler data; classifies performance bottlenecks; and formulates targeted optimization instructions. Crucially, it maintains a persistent record of the historically best kernel and its profiling snapshot, enabling regression detection and counterfactual reasoning across iterations.

The Conductor's prompt includes: source code, compiler/runtime error causes, profiler metric documentation, current profiling results, and the historical best code+profile. On each iteration, if correctness fails the loop repairs; once correct, the profiler runs and the Conductor drives performance refinement. Adding a new backend requires only hardware specs and profiler documentation URL plus a collection-API implementation.

Key Contributions

  • Profile-guided multi-agent architecture (Coder / Verifier / Profiler / Conductor) with a closed optimization loop integrating hardware counters into LLM reasoning.
  • Automated profiling knowledge consolidation: two-stage LLM pipeline crawls official profiler docs, summarizes individual entries, then merges them into a unified semantic compendium used at inference time.
  • Bottleneck-aware reasoning module: maps low-level metrics to high-level optimization strategies; Conductor can dynamically request additional metrics beyond a default set.
  • Cross-platform support: Triton for GPU (NVIDIA, AMD, Ascend NPU), C++/PyBind for CPU, extensible to new backends with minimal configuration.
  • Historical best preservation: retains best code + profiling data across iterations for regression detection, enabling stable, evidence-grounded refinement rather than random search.

Results

  • CPU (Intel Xeon Gold 6230R): PRAGMA achieves 2.81× average speedup over Torch; 10.95× improvement over N-PRAGMA (ablation without profiling feedback).
  • GPU (NVIDIA A100 40GB): PRAGMA achieves 2.30× average speedup over Torch; 1.90× over N-PRAGMA; 4.50× over Caesar.
  • Success rate vs. Caesar: Caesar 49.0%, N-PRAGMA 92.0%, PRAGMA 92.0%.
  • Fast1 (fraction meeting/beating Torch): Caesar 14.3%, N-PRAGMA 32.2%, PRAGMA 59.8%.
  • Max reduction case study (GPU): PRAGMA improves monotonically from 31.41% to 120% of Torch over 5 attempts; N-PRAGMA peaks at 85.97% then degrades to 58.06%.
  • Conv 1D case study (CPU): Conductor correctly identifies that higher IPC does not imply better performance (conditional branches degrading SIMD vectorization); after targeted fix, speedup rises from 2.93× to 4.71×.
  • Convolution category: limited gains across all methods due to loop-nest complexity and diverse tensor shapes.

Limitations

  • Convolution kernels show consistently weaker gains; high implementation complexity (deep loop nests, tensor layout transforms) makes correctness and performance hard to achieve simultaneously within 15 attempts.
  • CPU matrix multiplication lags behind because C++ lacks Triton's built-in autotune mechanism; finding optimal tiling for BLAS-competing performance across diverse shapes remains difficult.
  • Maximum 15 attempts per task; no analysis of failure modes when the budget is exhausted.
  • Evaluation uses a single LLM backbone (DeepSeek-R1) — generalizability to other models is untested.
  • Profiler documentation crawl and consolidation is offline; stale or incomplete vendor docs could degrade the Conductor's metric reasoning.
  • No evaluation on AMD GPUs or Ascend NPU despite claiming multi-backend support.

Relevance to Harnesses / Meta-Harnesses

PRAGMA is a concrete multi-agent meta-harness whose outer loop orchestrates specialized sub-agents (Coder, Verifier, Profiler, Conductor) through structured information exchange rather than ad-hoc prompting — directly instantiating the harness design pattern of role decomposition + feedback routing. The automated profiling knowledge consolidation module is itself a sub-harness: an LLM pipeline that transforms raw external documentation into a queryable semantic layer injected into the main loop at runtime, illustrating how meta-harnesses can self-configure their domain knowledge. The historical-best tracking mechanism demonstrates a stateful harness design where the orchestrator accumulates and compares cross-iteration evidence rather than treating each step as stateless, a pattern generalizable to any iterative refinement harness. For researchers tracking harness design, PRAGMA provides a worked example of how domain-specific feedback signals (hardware profiling) can be integrated as a first-class harness component rather than an afterthought.