Technical Summary: Arbor – Tree Search as a Cognition Layer for Autonomous Agents

Problem Statement

Recent advances in Large Language Model (LLM)-based code optimization have demonstrated that autonomous agents can discover high-quality solutions for isolated computational targets (e.g., single kernels or specific operators). However, production software performance arises from the complex interaction of multiple layers (application, framework, compiler, kernel, hardware). Optimizing one layer often exposes bottlenecks or regressions in another, creating a stateful, large action space where interventions can invalidate prior progress.

Existing approaches face three primary limitations:

1. Isolated Scope: Compiler autotuners and multi-agent frameworks typically operate within a single layer or on a single target, failing to address cross-layer dependencies.
2. Lack of Stateful Memory: Current systems do not maintain a shared working memory that evolves with every measurement, treating failures merely as stop signals rather than diagnostic data.
3. Inability to Handle Cascading Failures: When an optimization causes a system crash or regression, existing systems lack the mechanism to diagnose the root cause, propagate constraints back to the search, and reshape subsequent exploration.

The core challenge shifts from candidate generation to candidate selection: navigating a dynamic landscape where each intervention reshapes the bottleneck distribution, requiring a system that can detect shifts, diagnose failures across layers, and expand the search frontier autonomously.

Methodology: The Arbor Framework

Arbor formulates full-stack performance optimization as a heuristic-scored tree search over a dynamically expanding action space. It introduces a cognition layer that maintains an explicit search tree serving as shared working memory across agents.

1. Search Procedure and State Management

• Tree Structure: The optimization state is a tree $T = (V, E)$ rooted at a profiled baseline. Each node represents an action and its outcome (kept, reverted, or crashed). The path from root to node defines the cumulative sequence of interventions.
• Dynamic Expansion: The tree is not static. After a successful action, the system re-profiles to reveal previously invisible bottlenecks, generating new branches that did not exist at initialization.
• Search Loop: The Orchestrator executes a depth-first loop:
  1. Profile the current state to identify bottlenecks and GPU time shares.
  2. Score candidate actions using a heuristic function.
  3. Dispatch the highest-scored action to a Domain Specialist.
  4. Gate results through accuracy verification and end-to-end benchmarking.
  5. Update the tree: kept actions become the new baseline; reverted actions record diagnostic annotations; crashed actions trigger root-cause analysis.
  6. Re-score remaining candidates as the bottleneck distribution shifts.

2. Heuristic Scoring

Candidates are scored by balancing expected gain against cost and risk, incorporating an exploration bonus:
$$h(a) = \frac{g(a)}{c(a) \cdot (1 - r_{acc}) \cdot (1 - r_{crash}) \cdot m_{gap}} + C_s \frac{\ln N}{1 + n_{cat}}$$

• $g(a)$: Expected gain, estimated from the target kernel's GPU time share and category-specific speedup priors (derived from a persistent knowledge base).
• $c(a)$: Estimated wall-clock cost (implementation, restart, benchmarking).
• $r_{acc}, r_{crash}$: Empirical failure rates for accuracy and crashes, refined as the campaign progresses.
• $m_{gap}$: Urgency multiplier scaling with the remaining gap to the throughput target.
• UCB Term: Encourages exploration of under-sampled action categories.

3. Multi-Agent Architecture with Checks and Balances

Arbor structures its agents to mimic effective engineering organizations, decomposing capabilities into Hard Skills (domain expertise) and Soft Skills (coordination protocols).

• The Orchestrator (The "Tech Lead"):

  • Drives the optimization loop by profiling, scoring, and delegating.
  • Dynamically constructs Domain Specialists at runtime for specific tasks (kernel, framework, communication, compiler, dispatch).
  • Composes prompts using hardware context, knowledge base entries, and failure histories.
  • Does not implement code itself but decides what to try and who implements it.

• The Critic (The "Quality Assurance"):

  • Safeguards stability and measurement integrity.
  • Performs Root-Cause Analysis (RCA) on crashes, distinguishing implementation errors from unproductive directions.
  • Enforces Guardrails to detect benchmark anomalies.
  • Conducts Introspection to evaluate if a reverted idea retains merit under different constraints.
  • Monitors long-horizon system health, publishing constraints to prevent revisiting unproductive regions.

• Domain Specialists:

  • Constructed dynamically per task.
  • Operate asynchronously across layers, performing iterative refinement with local validation before returning merge-ready patches.
  • Adhere to constraints published by the Critic.

• Persistent Knowledge Base:

  • Accumulates strategies and failure modes across sessions.
  • Enables warm-start transfer to new models and hardware, informing scoring priors and reducing run-to-run variance.

Key Contributions

1. Stateful Tree Search Formulation: The paper formulates cross-layer optimization as a tree search where the tree serves as shared working memory. Failures propagate as constraints on child nodes, and successes trigger re-profiling to expand the search frontier.
2. Checks-and-Balances Multi-Agent Architecture: Introduces a dual-agent system where an Orchestrator drives aggressive optimization while a Critic safeguards stability. Neither agent can unilaterally drive the system, ensuring that short-term throughput sacrifices are made only when necessary for session viability.
3. Skill Decomposition: Decomposes agent capabilities into hard skills (domain expertise) and soft skills (coordination protocols), backed by a persistent knowledge base for cross-session transfer.
4. Dynamic Specialist Construction: Moves away from fixed agents, constructing Domain Specialists dynamically at runtime based on the specific optimization target, hardware context, and historical data.

Results

Arbor was evaluated on LLM inference serving across six production models on AMD Instinct GPUs (MI355X and MI300X).

• Throughput Improvements:
  • On MI355X, Arbor achieved +40% to +193% throughput improvement over vendor-optimized baselines after days of fully autonomous operation.
  • On MI300X (cross-generation validation), improvements ranged from +62% to +99%.
• Reproducibility: Run-to-run variance was within 2 percentage points, demonstrating hardware-agnostic reproducibility.
• Ablation/Comparison: A single agent without the Arbor harness (the "harness" being the multi-agent tree search framework) plateaued at only +33% throughput improvement and crashed irrecoverably within hours.
• Case Study (gpt-oss-120b): In a specific search tree for a 120B MoE model, the system explored 30 actions. 9 were kept, 16 were reverted with diagnostic insights, and 3 crashed but were recovered via RCA. The system reached 6,515 tok/s/GPU (a 48% improvement over the 4,405 tok/s/GPU baseline), utilizing insights from introspection (e.g., identifying dead code paths) and RCA (classifying crashes as dependency issues rather than algorithmic flaws).

Significance and Claims

The paper claims that Arbor represents a shift in autonomous optimization from isolated candidate generation to navigating stateful, large action spaces. By treating failures as diagnostic signals that reshape exploration and maintaining a shared working memory, Arbor enables fully autonomous multi-day campaigns that would otherwise be impossible due to system instability.

The framework demonstrates that:

1. Cross-layer optimization is feasible without human intervention, provided the system can detect bottleneck shifts and diagnose failures across the stack.
2. Checks-and-balances architectures are critical for sustained autonomy, preventing aggressive exploration from destabilizing the system.
3. Knowledge transfer across sessions and hardware generations is achievable through persistent memory, lowering the barrier to bringing up new hardware or models.

Ultimately, Arbor aims to reduce the operating cost and carbon intensity of LLM deployments by automating the coordinated effort traditionally required from engineering teams across the full software and hardware stack.