XLA from the Trenches: Correctness on TPU Workloads at Scale
Ignacio Cano, Amir Saadat, Rishipal Singh Bhatia, Samuel Ginzburg, Patrick C. Toulme, Ioannis Zarkadas, Praveen Narayanan, Pratik Bhatu, Nicolas Basile, Alexander Dietmüller, Ruishan Wang, Jacob Wahlgren, Victor Barros, Subhankar Shah, Michael Kuperstein, Farzin Houshmand, Emma Sevastian, Peter Gavin, Karthik Srinivasa Murthy, Berkin Ilbeyi, Amit Sabne, Eli Bendersky
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.