XLA from the Trenches: Correctness on TPU Workloads at Scale
Organizations: Google
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.
Figures & tables
| Compiler Subsystem | Defects | Primary Defect Categories |
|---|---|---|
| TensorCore LLO & Emitters | Sublane store mask corruption, unsigned comparison mismatches, narrow-integer unpack guards, multi-tile loop revisit and chunk-packing bugs | |
| HLO Simplifiers & Rewrites | Nested dynamic-slice index overflow, degenerate scatter indexing, invalid loop trip-count and dilated-pad folding, space-to-batch halo underflow | |
| Precision & Emulation Rewriters | Emulated -bit integer divide-by-zero, IEEE-754 bitcast and narrowing cast omissions, sign-extension loss, complex trig overflow/underflow | |
| SparseCore Lowering | interleaved float-to-int packing errors, sublane bit-reversal omissions, cross-core DMA synchronization and alignment faults | |
| Custom Kernel Lowering | Vector transpose op duplication, layout inference attribute mismatches, bitcast register-chunk corruption |
Software Vulnerability Detection
Software vulnerability detection identifies insecure code, pinpoints weaknesses, and helps developers validate or prevent exploitable bugs.
See all 149 papers