AMD TraceLens Agent Automates PyTorch Bottleneck Triage

AMD introduced the TraceLens Agent, an autonomous system that converts raw PyTorch profiles into prioritized optimization reports. The agent builds on TraceLens, an open-source Python library that processes profiler traces from PyTorch, JAX, and AMD RocProf. TraceLens analyzes each kernel's contribution to runtime, achieved FLOPs, and bandwidth, using a roofline model to classify kernels as memory-bound or compute-bound. The library's Trace2Tree module constructs a hierarchical event tree linking Python operations and GPU kernels back to their CPU-side dispatch, including tensor shapes and call stacks. Additional tools, NcclAnalyser and TraceDiff, handle collective communication analysis and gap analysis between workloads.

The agent addresses the manual effort required to interpret TraceLens output. A single profile can contain thousands of CPU events and kernel records, and identifying bottlenecks requires reasoning about tensor shapes, precision mode, backend selection, launch path, and call stacks. The agent orchestrates analysis through a skill file, a structured prompt loaded into context when a capability is requested. The workflow pre-processes the trace at the start, precomputing values that sub-agents reason over, including per-category kernel metrics, efficiencies, time splits across compute, idle, communication, and memory copy, and kernel fusion candidates.

Sub-Agent Architecture

The analysis fans out into specialized sub-agents, each with its own prompt and isolated context window. The compute tier assigns one analyzer per major kernel class: GEMM, attention, Mixture-of-Experts, elementwise, reduction, normalization, convolution, and Triton-generated kernels, plus a generic analyzer for uncategorized cases. Each analyzer scores kernels by their gap to achievable peak performance. The system tier handles cross-kernel opportunities: an idle CPU analyzer characterizes host-side stalls and device idle gaps, a multi-kernel analyzer examines transfer patterns and compute-communication overlap, and a kernel-fusion analyzer identifies short kernel sequences that could collapse into a single launch. System analyzers activate conditionally, only when relevant conditions are detected, such as significant idle time or exposed communication overhead.

Sub-agents exchange data through disk using JSON, CSV, and intermediate markdown files with enforced schemas. This disk-backed exchange serves as persistent workflow state, useful in longer runs where an agent may lose context from earlier steps.

Reliability Guardrails

The agent includes three guardrails. First, templates: the orchestrator and sub-agents fill fixed sections rather than generating free-form text, ensuring each run produces the same sections in the same order. Second, feedback loops with validation at both sub-agent and orchestrator levels, where agents validate their own outputs and correct local errors. An evidence guardrail confines each agent to quantities supported by trace data. Third, an end-to-end evaluation harness that captures regressions before each release, performing automated testing for workflow adherence, structural correctness, semantic accuracy, and reproducibility. The evaluation uses scripted checks and LLM-based evaluators, with results aggregated and classified to highlight actionable issues.

The final output is a single Markdown report. It opens with an executive summary describing the workload and top-level metrics, then presents findings in three tiers: Compute Kernel Optimizations, Kernel Fusion Opportunities, and System-Level Optimizations. Each finding includes Insight, Action, and Impact components, with an impact score ranking it relative to others. A Detailed Analysis section follows with supporting tables and reasoning. Structured HTML comment markers accompany ranked items and impact scores, making the report machine-consumable for automated optimization systems.

An example finding shows 18 BF16 GEMM shapes flagged as the dominant cluster, accounting for 45.7 seconds of GPU kernel time, roughly 80.6 percent of compute time. The cluster spans large MLP shapes, gate and up shapes, QKV and output shapes, and LM-head shapes, all dispatching through the Tensile backend with compute-bound status. One operation, aten::mm with shape (24576,8192) by (8192,28672) in BF16, ran 320 times, totaling 7607.463 milliseconds, achieving 68.74 percent of its roofline efficiency.