XLA from the Trenches: Correctness on TPU Workloads at Scale
Ignacio Cano and colleagues at Google describe how they verify and triage the XLA:TPU compiler at production scale, combining fuzzing, differential testing on real workloads and human-supervised swarms of AI agents that generate stress tests, reproduce failures and draft fixes.
Ask this paper
Problem. Aggressive compiler optimizations can introduce silent numerical miscompilations and buffer corruption that degrade multi-million-dollar training runs and live serving without raising any compiler error.
Fuzzing. Template-guided synthetic programs are combined with opportunistic fuzzing on real production graphs. Performance autotuners already compile many variants of each graph, so differential checks run alongside them at little extra cost.
Agent swarms. Autonomous agents generate stress tests, reproduce failures and write fixes. The team deliberately steered them away from loud crashes toward silent data corruption, and the paper shows minimal JAX reproducers the agents found and patched end to end.
Isolation. A hierarchical pipeline bisects a full workload graph down to the single failing operation inside a fused kernel.
Results. The framework has resolved over 142 compiler bugs, with 103 more findings under review. Lessons include that reviewer burden grows faster than fix complexity and that adversarial pre-filtering of agent fixes reduces false-positive fatigue.
Abstract
Accelerating modern foundational models has pushed domain-specific compilers toward increasingly complex, multi-tiered optimization pipelines. However, this heavy optimization focus introduces delicate tradeoffs between peak execution performance and correctness. Subtle codegen defects can trigger numerical degradation and memory or buffer corruption, invalidating multi-million-dollar training runs and degrading live serving quality without raising obvious compiler errors. In this paper, we present a battle-tested methodology for verifying and triaging Google's XLA:TPU compiler at hyperscale. We leverage fuzzing techniques, combining synthetic program generation with real production traces, while orchestrating AI agent swarms to autonomously generate stress tests, reproduce failures, and draft fixes. To pinpoint root causes in full workloads, we pair this with a hierarchical differential pipeline that bisects large graphs down to individual failing operations inside fused kernels. Deployed in production, our framework has already driven the resolution of over 142 compiler bugs (with 103 additional findings under review), from which we distill pragmatic lessons learned.