← Back to feed

Optimising LLM training and serving across a distributed hardware system is an engineering speciality. The configuration space is large, the feedback loop requires running actual workloads, and the expertise to read profiler traces and map them to the right parallelism strategy lives in a small number of engineers.

Google Research has published PROMPTS — Performance Optimization via Multi-Agent Planning for LLM Training and Serving — a system that automates that process using two agents in sequence.

The results across 8 internal production workloads: the human-validated configuration appeared in the system’s first proposal batch every time, and the top-ranked proposal matched what engineers actually adopted in 87.5% of cases.

The Two Agents

Analyzer Agent. This agent reads profiler traces from a running system and classifies the primary bottleneck as one of three types: compute-bound, memory-bound, or communication-bound. The classification step is what allows the search space to collapse. Instead of testing configurations exhaustively — thousands of possible combinations of data parallelism, model parallelism, and sequence parallelism — the agent identifies which dimension of the system is the binding constraint.

Proposal Agent. Once the bottleneck is classified, the Proposal Agent retrieves engineering knowledge and generates three targeted parallelism strategies. The proposals are grounded in the specific bottleneck classification: a memory-bound workload gets proposals that reduce memory pressure, not proposals that optimise communication topology.

The key design insight is that the Proposal Agent does not need to search the full configuration space. It searches the subspace implied by the Analyzer’s diagnosis. That reduction is what makes the system practically useful — not a large search, but a small, informed one.

Test Conditions

The paper tested PROMPTS across 8 production workloads covering:

• Dense models and MoE models
• Pre-training, post-training, and serving
• TPU systems from 2 chips to 2,048 chips

These are internal Google workloads, not toy benchmarks. The engineering knowledge the Proposal Agent retrieves was built from real TPU configuration practice.

The metric used for evaluation is whether the configuration the human engineers ultimately adopted appeared in the system’s first proposed batch — the set of configurations generated before any iterative refinement.

Result: the human-validated configuration appeared in the first batch on all 8 workloads. In 7 of the 8 cases, it was the first configuration in the batch — meaning the agent’s highest-ranked proposal was also the one the engineers chose.

Why the Bottleneck Classification Matters

LLM infrastructure optimisation has historically been treated as a large-scale search problem: enumerate candidate configurations, run them, measure performance, repeat. This approach works but requires substantial engineering time and compute budget per optimisation cycle.

PROMPTS restructures the problem. The bottleneck classification step makes the assumption that compute, memory, and communication constraints are separable and diagnosable from profiler data before an exhaustive search. If that assumption holds — and the 8/8 first-batch result suggests it does for the workloads tested — then optimisation becomes a two-step analytical problem rather than a search problem.

The practical implication: the number of configurations that need to be tested goes from exhaustive to targeted. The Proposal Agent generates three strategies per diagnosis. Three runs, not hundreds.

Scope and Limitations

The paper covers 8 workloads, which is a limited sample. All workloads are internal Google systems running on Google TPUs. The system uses engineering knowledge built from that same infrastructure, which may limit direct transferability to other hardware stacks.

The 87.5% top-match rate covers 7 of 8 workloads. In the eighth case, the human-validated configuration appeared in the first batch but not as the first item. That is still a useful result — the configuration was surfaced, just ranked second — but it distinguishes from claiming perfect precision.

What the paper establishes is a proof of concept for bottleneck-first optimisation in LLM infrastructure. The pattern — diagnose the constraint, then search the constrained subspace — is general even if the current implementation is Google-TPU-specific. Whether it transfers to other hardware stacks (Nvidia GPU clusters, AWS Trainium) is an open question the paper does not answer.

The paper is published at research.google/pubs/prompts.