The One-Line Lie That Crashes Google's AI Compiler (a DoS in XLA)

I found a way to crash Google's AI compiler with a single malformed line of text - by handing it a description that contradicts itself. Instead of noticing, it believed both halves at once, then fell over.
Let me explain in plain English.
What is XLA?
When you run a modern AI model - in TensorFlow, JAX, or the giant models behind today's chatbots - something has to turn that model into fast code for GPUs and TPUs. That something is XLA, Google's AI compiler. To do its job it reads an internal text format called HLO - think of it as the compiler's notebook. Engineers and tools pass HLO around as text all the time: to reproduce bugs, run tests, and feed the compiler work.
And reading text from somewhere else is exactly where software tends to trip.
The bug: two numbers that were allowed to disagree
Big models don't fit on one chip, so XLA splits the work across many devices. The HLO text describes this using two separate pieces: a shape (say “2 by 2” = 4 devices) and a list of which device goes where.
These two have to agree on one thing - the total device count. Shape says 4, the list must describe 4.
The parser never checked. So I wrote a line that lies to itself: a shape that says 4, paired with a list that says 16. Four isn't sixteen - but the parser shrugged and built the broken layout anyway.
The contradiction, in numbers
Here's the actual line I fed it:
sharding={devices=[2,2]<=[16]}
Read it like this. The [2,2] is the device grid - 2 × 2 = 4 cells. The <=[16] part says “fill those cells from an iota of 16,” and an iota is just the counting sequence 0, 1, 2, …, 15 - that's 16 device IDs.
So the compiler is asked to place 16 numbers into 4 slots. The one rule that must always hold is simple:
product of the grid = number of device IDs
2 × 2 16
4 ≠ 16 ✗
“Four equals sixteen” is false - and that single broken equation is the whole bug. The entire fix is just checking it before trusting the layout:
if product(grid_dims) ≠ product(iota_dims): reject
Put simply: it's a seating chart with 4 chairs and 16 name cards. On paper it looks fine - until people actually try to sit down.
The crash
Nothing broke right away. The contradiction just sat there, baked into the compiler's state like a typo in a blueprint nobody noticed. Then, much later - when XLA tried to use that layout - the numbers didn't add up, and it did the only thing it could: it aborted. An unrecoverable crash, thrown far from the actual bad line, which makes it miserable to trace back.
How bad is it, really?
Let me be precise. This is a denial-of-service - a crash. Feed the compiler a malformed line and the process dies. It is not remote code execution, and it does not steal data. Nobody's “hacking the AI” with this.
But a crash still matters, because HLO text travels - shared in bug reports, pasted into tools, fed into automated compile services. Anywhere a program accepts HLO it didn't write, one poisoned line can take it down. And “trusting input that contradicts itself” is the same root mistake that, elsewhere, leads to far worse bugs. The harmless version is a gift: you get to close the door before something worse walks through.
The fix, and the takeaway
The fix is boringly simple, which is the point: before building the layout, compare the two totals. If they don't match, stop and return a clear error instead of crashing three rooms later. One check, plus tests. It's now merged into both TensorFlow and XLA.
One contradiction. Two numbers that didn't match. A compiler behind a huge slice of modern AI, brought down by a line of text that argued with itself.
The lesson underneath is one of the biggest in software: never trust input that hasn't proven it makes sense.
References
- Origin fix: openxla/xla PR #46606
- Merged into TensorFlow: tensorflow/tensorflow PR #125087
- Merged into XLA: openxla/xla PR #47060
Comments
Post a Comment