The Alibi of a Correct Answer

The shipping note for nx_vulkan 0.3.0, which exists because the previous release was an inference backend published as a training one. Every test passed throughout. That is not a mitigating detail; it is the entire subject.

There is a species of software failure that no amount of testing will catch, and it is distinguished by the awkward property of returning the correct answer. It does not crash. It does not drift in the sixth decimal. It does not produce a stack trace, a warning, a log line, or a single one of the small courtesies by which a program normally admits that something has gone wrong. It computes precisely what you asked for, to the last bit, and then hands it back with an expression of perfect innocence. For three months this project shipped one, and the suite was green the entire time.

The claim on the front page of the repository was that autograd came free — write your model, call Nx.Defn.grad, and the gradients run on the GPU along with everything else. The gradients did not run on the GPU. They ran, in their considerable entirety, in interpreted Elixir on the CPU, at a cost the reader can judge for themselves: a LeNet-style training step took 20.9 seconds. It now takes 84 milliseconds. Anybody who believed the README received something like a two-hundred-and-fifty fold less than what it promised, and received it in the form of numbers that were, to repeat the point that makes this interesting, entirely correct.

The safety net and the blind spot are the same object

The mechanism deserves to be stated plainly, because it is not a bug in the ordinary sense and it will not be the last of its kind. Every operation in this backend has the same shape: attempt the GPU; if the operation is unsupported, transfer the tensors to Nx.BinaryBackend — the reference implementation, written in plain Elixir, slow and impeccable — compute there, and transfer the result back. This is a good design. It is why the backend has never returned a wrong number in its life. It is also, and by the identical mechanism, why nothing could tell us that it had stopped using the GPU.

Consider what a test can actually assert. It runs an operation and compares the result against a reference. The reference is BinaryBackend. The fallback is BinaryBackend. A test that fires when the GPU path is silently abandoned would have to distinguish a value from itself. Our entire verification apparatus was, in the strictest sense, asking the accused to confirm his own alibi — and the alibi was true. Every time.

An ecosystem that agrees not to look

I would like to report that this was a local failure of diligence, and in part it was. But when we went to check what the wider Nx ecosystem verifies about a third-party backend, the findings were not flattering to anyone, this project very much included.

The community standard for validating a backend is to run doctest Nx with your backend installed as the default. We do this; we pass 851 of them on three separate GPUs. Nx’s doctests contain zero gradient examples. Not few — zero. Upstream Nx does possess a proper gradient suite: a 6,031-line grad_test.exs with 293 tests, and a central-differences helper, check_grads!, already used thirty-three times within the project. Both live under test/, and the Hex package ships lib/** only. A backend author who runs mix deps.get receives no deps/nx/test directory at all. The best gradient-verification tool in the ecosystem exists, works, and is sealed behind the packaging wall for no reason anyone has articulated.

Nor is there much comfort in the neighbours. Torchx has no gradient tests. EXLA has roughly four, all of them corner cases. The reason this has cost them nothing is instructive: neither has a blanket silent fallback. EXLA touches BinaryBackend in three specific places; Torchx simply raises. They do not have this bug class because they declined the convenience that produces it. We took the convenience, and we did not pay for it until we did.

It is worth adding, since it bears on any future proposal for a shared conformance kit: a kit that checked only values would not have caught this either. You cannot find this defect by comparing outputs. You can only find it by asking a different question, which is where the work actually started.

Count what you fear, not what you hope

The instrument that broke the case open is embarrassingly modest. It is a counter. Nx.Vulkan.Fallback tallies, per process, every operation that leaves the device, attributing each to the callback that abandoned it. The assertion it enables is not assert_all_close but assert Fallback.count_total(fun) == 0: not is this right, but did it happen where I said it would. The first instance of the bug took a day of measurement and source-reading to isolate. The counter found the next six in an afternoon apiece.

What it found was a pattern, and the pattern is the finding worth carrying away. Every GPU fast path in the backend was guarded by a predicate describing the shapes a forward pass produces. Nx.Defn.Grad emits shapes that no human writes by hand — transposed, permuted, dilated, contracted along the other axis — and each of those predicates looked at the gradient, failed to recognise it, and politely declined. Eight instances:

Six of the eight were narrow gates rather than missing capability. The shaders could already do the work and were being refused it. Only reverse and broadcast needed genuinely new kernels, and both were written from the index-remap skeleton that already existed for transpose_nd. There is a certain grim comedy in a backend that turned a three-minute MNIST epoch into an eleven-hour one — 1,875 steps at batch 32, do the multiplication — because eight if statements were describing the wrong half of the calculus.

The second sub-pattern is the one I would put on a poster. Nx materialises integer literals as {:s, 32}. The 0 in max(x, 0). A mean’s divisor. select’s zeros. Pooling’s init_value. Four unrelated operations, one cause: a four-byte integer constant meeting a gate that demanded an exact float dtype, and dragging an entire {32,16,14,14} tensor to the CPU behind it. Any gate in a numerical codebase that demands an exact dtype match is a gate that will eventually be wrong. Coerce, then check.

The counter lies too, in one specific direction

Honesty about the instrument: the count is a lower bound, and it is not monotone. It only sees operations that reach this backend. Once a fallback strands a tensor on BinaryBackend, Nx dispatches everything downstream of it there as well, and none of that is recorded. Which means fixing one operation reveals fallbacks that were always happening and were never visible. This occurred four separate times — window_scatter_max, select and reduce all appeared because something upstream had been fixed. A rising number can be evidence that you are winning. Read the composition, never the total.

The stopwatch would have stopped after the second fix

Here is the result I did not expect and would not have predicted. Performance is not linear in the number of fallbacks. Removing them from eleven to three barely moved the clock at all. Going from three to one moved everything:

Cost is dominated by the largest tensor that leaves the device, not by the number of operations that leave. Ten cheap fallbacks on a {32,10} logit tensor cost less than one on {32,16,14,14}, because the host leg is pure-Elixir and scales with elements. Several of the intermediate fixes produced no wall-clock improvement whatsoever and were entirely correct and entirely necessary.

The methodological point is worth more than the measurement. A process driven by the stopwatch — make a change, time it, keep it if the number moves — would have abandoned this work after the second fix, concluded that fallbacks were cheap, and left twenty seconds a step on the table. The census kept insisting that work was moving on-device while the clock insisted nothing had happened. The census was right and the clock was right; only the inference from the clock was wrong. Measure the thing you are actually changing, and give it more credence than the thing you wish would change.

Then the fix turned out not to be the fix

With the backward pass on-device, we raced the whole thing against EXLA on CUDA, using the Axon MNIST model, one training step at batch 32. This project has a fusion compiler — Nx.Vulkan.Compiler, an Nx.Defn compiler that traces a whole graph and compiles it to an on-device stage schedule. It represents the single largest body of work in the repository. It was the obvious candidate to close the gap.

Fusion is 0.76×. It is twenty-four percent slower than doing nothing clever at all, and bit-identical while being so. On a convolutional graph it is 0.98×, which is to say neutral. The compiler splits its stages at dot boundaries, and a graph that is almost entirely dot offers its tracing, scheduling and boundary buffers nothing whatever to amortise against.

I want to dwell on this rather than hurry past it, because the temptation to hurry past it is precisely the disease. The most useful thing this race produced was not a speedup. It was the demolition of a plan. Anyone reading the eager-versus-fused rows and concluding “we need more fusion” would have spent a quarter building the wrong thing, and would have been encouraged in it by every instinct a compiler education instils. An optimisation that removes work from the shaders cannot possibly explain a deficit that removing work from the shaders makes worse. The gap therefore had to be per-dispatch cost. Which meant the answer was not more compiler. It was fewer submits.

One command buffer, one fence, 1.45–1.71×

Every operation used to build its own command buffer, submit it to the queue, and block on wait_idle(). Dispatches are now recorded into a pending queue and submitted as one command buffer with one fence wait, flushed automatically at every host boundary — a download, an upload, a concat, an FFT — and at a size cap. Correctness never depends on anyone remembering to flush.

This is unglamorous and it is safe, for two reasons worth naming. Vulkano’s AutoCommandBufferBuilder tracks resource usage while recording and inserts the pipeline barriers between commands when the buffer is built, so a read-after-write between two batched dispatches is synchronised without hand-rolled barriers. And the only route by which a value reaches the host is a download, which flushes first. The one genuine obstacle was that the builder deliberately does not implement Send — a command buffer may not migrate threads mid-recording, and consecutive NIF calls land on whichever dirty scheduler happens to be free — so the pending queue holds closures, replayed into a builder created on the flushing thread.

Raced across the whole fleet: an RTX 3060 Ti on Linux, a GT 650M and a GT 750M on FreeBSD, cards spanning 2012 to 2021 and two operating systems.

The loss is 2.6447360515594482 in every cell of that table — every arm, every cap, all three hosts, two architectures, two operating systems. Batching changes when work is submitted and never what is computed, and this is how you demonstrate that rather than assert it.

It also wins on all three, which could not be assumed and is the reason the race was run. This project has two optimisations that reverse sign across hardware: register-blocked GEMM wins on Ampere and regresses on both Keplers, and the many-slot fused reduction wins 4.4× on Kepler and collapses to 0.44× on Ampere. Batching needs no device-class gate. It ships on by default.

And now the part that stings. The evidence for this predated the work by three months. A planning document from May recorded a measurement of 1.13 ms of fence wait against 138 µs of submit per dispatch — the wait dominating by a factor of eight, which is precisely what batching amortises. The finding sat in a Markdown file, correct and unread, until a race against a competitor independently pointed at the same place. Writing something down is not the same as knowing it.

Two occasions on which we were wrong in our own favour

A release note that only recounted the discovery of other people’s mistakes would be worth very little, so here are ours, and specifically the ones that ran in the flattering direction, because those are the ones to be most suspicious of.

We accused EXLA of something it did not do. A committed benchmark document stated that EXLA “failed to compile” convolutions. This was false. It was written from a single failing run without isolating the cause, and it happened to make this project look better at a competitor’s expense. A seventeen-variant matrix narrowed the real failure to two stacked convolutions with stride 2 and channels: :first, in the gradient only; any single relaxation compiles, and Axon’s default layout avoids the case entirely. EXLA compiles and trains convolutional models perfectly well. The correction is in the changelog, because a correction that lives only in a commit message is not a correction.

And our own benchmark lied to us twice. The first version of the batching benchmark neglected to call Nx.global_default_backend/1, so every tensor that defn materialises internally — gradient constants, the scalar inside max(h, 0.0) — landed on BinaryBackend and dragged the graph to the CPU. It measured 6.8 seconds per step, which is exactly the BinaryBackend row of the race table, a coincidence that should have been recognised immediately and was not.

The second version built its inputs with Nx.sin, which is not a supported unary op and therefore falls back. Its result lands on BinaryBackend, and everything downstream then computes there without being recorded, because the counter only sees operations that reach this backend. A 32×784×128 matmul underneath it spent 1,039 ms on the CPU while the fallback counter cheerfully reported an empty map. The same matmul on the GPU is 1.41 ms. This is the counter’s documented lower-bound caveat, arriving with a price tag attached, in the very tooling written to prevent it. The benchmark now refuses to report a timing unless every input and every gradient asserts residency on the GPU backend, alongside the existing check that the loss is a number.

Two smaller things, for the record

Ties have to be constructed, never hoped for. window_scatter_max assigns the gradient to the last maximum in row-major order, so the shader requires >= and not >. With > it is correct on random data and wrong wherever values repeat — and the output of a relu is full of exact ties at zero. Random floats will never generate that case for you. Build the ties.

Occasionally the reference is the broken one. Nx.BinaryBackend.window_scatter_max/5 round-trips f64 through f32 (2.4715269558223154 becomes 2.471526861190796). For f64 pooling gradients this backend is now more accurate than the thing it is tested against, which makes agreement with the reference the wrong assertion. The test instead requires that each value be an exact element of the source tensor — a stronger claim, and one that declines to inherit the reference’s defect.

What 0.3.0 actually is

A CNN training step now performs exactly one host fallback: pow in f64, which GLSL.std.450 does not provide, and which should remain a fallback rather than quietly boundary-cast through f32 and trade real precision for a tidier table. The gradient suite, the fallback counter and the full test suite — 851 doctests, 415 tests, no failures — run green on all three GPUs at every batching cap, including a cap small enough to force a flush mid-graph on nearly every operation, because a missed barrier would present as nondeterministic wrong numbers rather than a crash and one green run proves nothing about it.

The remaining distance to EXLA on a small dense model is roughly twelve-fold after batching, and we now know what it consists of, which is per-dispatch cost and GEMM quality rather than anything a compiler pass will fix. That is a smaller and more honest claim than the one this project was making three months ago, and it has the advantage of being supported by measurements taken on hardware we can name.

If there is a general lesson, it is not the tidy one about writing more tests. It is that a test suite is a set of questions, and a suite composed entirely of “is this value right?” will return a clean bill of health to a program that has quietly stopped doing the thing you built it to do. Ours did, for a full release, with a perfectly straight face. The correct answer is not evidence of correct behaviour. It is merely an alibi, and the first duty of anyone verifying anything is to notice when the alibi is the only thing on offer.

0.3.0 is on Hex. If you are on 0.2.0 and training on a GPU, upgrade — your results were right, and they cost you about two hundred and fifty times what they should have. The audit that this post condenses is docs/BACKWARD_PASS_AUDIT.md; the fleet data behind the batching table is bench_results/BATCHED_DISPATCH.md; the EXLA race, including the retracted conv claim, is bench_results/MNIST_EXLA_RACE.md; and what the Nx ecosystem does and does not verify about a third-party backend is docs/BACKEND_VERIFICATION_GAP.md. The fusion compiler whose regression sent us looking elsewhere was introduced in Compute It Twice.