DEV Community: Yehor Cherednichenko

Half of my FP8 convolutions were silently running in f32 – a one-line XLA fix

Yehor Cherednichenko — Thu, 17 Sep 2026 17:40:01 +0000

A stray `return` made XLA:GPU skip FP8 rewriting for most of a module

I'm training a search-based RL agent (Gumbel MuZero via DeepMind's mctx) with a small convolutional network, in JAX on a single RTX 5070 (Blackwell, sm_120). The self-play actor is the hot loop: one policy move is a root network evaluation plus 16 MCTS simulations, each of which evaluates the same network again inside a lax.while_loop, for a batch of 256 environments.

Blackwell has FP8 tensor cores and cuDNN 9 has FP8 convolution "graph" kernels, so the obvious experiment was to run the actor's network in float8_e4m3fn. The forward pass got faster right away:

batch	bf16 forward	fp8 forward
256	2.74 ms	1.98 ms
1024	9.65 ms	7.15 ms

But the compiled self-play step was not using FP8 for all of its convolutions, and finding out why turned into a one-line fix in XLA itself (PR). This post is the story.

Counting custom calls

JAX makes it easy to look at what XLA actually compiled:

txt = jax.jit(policy_move).lower(params, obs, ...).compile().as_text()

On the GPU backend an FP8 convolution that cuDNN will run natively shows up as a custom call with target __cudnn$convForwardGraph (the "graph" API, with scaling and clamping fused in). A plain convolution is __cudnn$convForward. So the check is a grep:

what is compiled	`convForwardGraph`	`convForward`
one network forward	20	1
two forwards on the same params, same computation	40	2
root forward + the same forward inside a `while_loop` (the real policy move)	20	21

The network has 20 FP8 convolutions (the single legacy call is the first layer, which stays in bf16 by design). Two copies in one computation: 40 graph convolutions, as expected. Root forward plus loop body: only one copy got the FP8 rewrite. The other 20 convolutions were compiled as legacy convolutions on operands converted back to f32.

The wrong hypothesis first

My first guess was operand sharing: the FP8 weights and their scales are parameters used both by the root forward and, through the loop carry, by the body — maybe the rewriter refuses to touch a convolution whose operand has more than one user. Row 2 of the table already argued against it (two forwards on the same parameters both get rewritten), and wrapping the parameters in optimization_barrier for one or both copies changed nothing. The difference between rows 2 and 3 is not the operands. It is that the two copies live in different HLO computations: the entry computation and the while body.

Reading the pass

The rewrite is done by CudnnFusedConvRewriter in xla/backends/gpu/transforms/cudnn_fused_conv_rewriter.cc. Its RunImpl loops over the non-fusion computations of the module and, for each of them, first runs F8GraphConv (the FP8 → graph-conv rewrite), then a dozen bias/activation fusions for non-FP8 convolutions:

for (HloComputation* comp :
     module->MakeNonfusionComputations(execution_threads)) {
  bool changed = false;
  if (!compute_capability_.IsRocm() && !compute_capability_.IsOneAPI()) {
    auto* cc = compute_capability_.cuda_compute_capability();
    ABSL_ASSIGN_OR_RETURN(changed,
                     F8GraphConv(comp, *cc, dnn_version_, toolkit_version_));
    if (changed) {
      return changed;      // <-- exits the pass, not the loop
    }
  }
  ABSL_ASSIGN_OR_RETURN(changed, FuseRemoveConvertInConv(comp));
  any_changed |= changed;
  // ... more fusions, all `any_changed |= changed`
}

Every other step in that loop accumulates into any_changed. The FP8 step alone returns. So the first computation in which an FP8 convolution gets rewritten ends the whole pass, and every FP8 convolution in every later computation is never visited. Computations are visited callees first, which is why in my case the loop body won and the entry computation lost.

Downstream nothing complains: a later fallback pass sees FP8 convolutions that were not turned into graph convolutions, converts their operands to f32 and lowers them as ordinary cuDNN convolutions. The program is correct. It is just quietly slower, and the only symptom is the custom-call count.

The fix is the obvious one:

    if (changed) {
      any_changed |= changed;
      continue;
    }

continue rather than falling through keeps the previous behaviour for the rewritten computation (the remaining fusions only match the legacy conv custom calls, so they have nothing to do there) and lets the loop reach the other computations.

Checking it end to end

XLA lives inside the jax-cuda12-pjrt plugin, so verifying the fix meant rebuilding that wheel from the JAX 0.10.2 sources against a local XLA tree with the patch (about an hour on 28 cores with Bazel). With the rebuilt plugin the policy move compiles to 40 graph convolutions out of 40, and a three-copy variant to 60 out of 60.

The self-play numbers for one policy move at batch 256 and 16 simulations:

actor	move time
bf16	54 ms
fp8, entry copy demoted to f32 (stock plugin)	39 ms, −25 %
fp8 everywhere (patched plugin)	−27 %

The fix itself is worth only two percentage points in my workload, and the reason is instructive: the demoted copy was the root evaluation, which runs once per move, while the 16 evaluations inside the loop were the copy that got the rewrite. Had the order been the other way round, the loop would have run at bf16-or-worse speed and FP8 would have looked like a loss. Which copy loses is an accident of computation order, not of your code.

A 20-line repro with no network at all, the same FP8 convolution once in the entry computation and once in a while body:

import jax, jax.numpy as jnp
from jax import lax

def conv(x, w):
    return lax.conv_general_dilated(x, w, (1, 1), "SAME",
                                    dimension_numbers=("NHWC", "HWIO", "NHWC"),
                                    preferred_element_type=jnp.float32)

@jax.jit
def f(x, w):
    y0 = conv(x, w)                                            # entry computation
    def body(c):
        i, y = c
        return i + 1, y + conv(y.astype(jnp.float8_e4m3fn), w) # while body
    return lax.while_loop(lambda c: c[0] < 4, body, (0, y0))[1]

x = jnp.ones((1, 6, 6, 128), jnp.float8_e4m3fn)
w = jnp.ones((3, 3, 128, 128), jnp.float8_e4m3fn)
hlo = f.lower(x, w).compile().as_text()
print(hlo.count('custom_call_target="__cudnn$convForwardGraph"'),
      hlo.count('custom_call_target="__cudnn$convForward"'))

Stock jax-cuda12-pjrt 0.10.2 prints 1 1. The patched plugin prints 2 0.

The PR adds a regression test to cudnn_fused_conv_rewriter_test.cc: a module whose two conditional branches each contain an FP8 convolution; both branches must come out as graph convolutions. Without the fix the test fails on the second branch, which still carries the legacy __cudnn$convForward target.

Who is affected

Any XLA:GPU program that has FP8 convolutions in more than one HLO computation: a while/scan loop next to code outside it, both branches of a cond, a jax.checkpoint-ed block next to an un-checkpointed one. FP8 matmuls are not affected (the GEMM rewriter loops over computations correctly), and neither are TPU or ROCm (the FP8 graph-conv path is CUDA-only). The early return is in the XLA shipped with jaxlib 0.10.2 and is still on main as of this writing; the pass's FP8 tests only ever had single-computation modules, so nothing caught it.

To check your own program: count custom_call_target="__cudnn$convForward" (with the closing quote, so that it does not match ...Graph") in the compiled HLO. If it is not zero for a network you quantized to FP8, look at which computation those calls live in.

Until the fix lands, the workarounds are unattractive but simple: keep all FP8 convolutions in a single computation (for example, run the root evaluation as iteration zero of the loop), or accept that one copy runs in f32 and choose which one by keeping the other in bf16 explicitly.

Two things I got wrong along the way

1. I blamed the data flow, not the pass. Operand sharing, aliasing, optimization_barrier — all plausible, all wrong, and each one cost an hour of compile-and-count. The pass source was 1,700 lines and the answer was on line 1,700. Reading the pass that owns the custom call you are missing is a better first step than theorizing about the graph.

2. FP8 self-play was a win for the actor and a loss for the agent. The 25–27 % faster policy move is real, but in my training runs the FP8 actor learned worse: from a cold start it failed to take off, and after switching to FP8 once the policy had ignited, the final score on my evaluation set dropped from 0.86 to 0.79 for the same budget. That is a separate story about what MCTS does with slightly noisier value estimates, but it is why I tell it here: the compiler fix is worth having, and it does not make FP8 a free lunch for search-based RL.

Takeaways

The compiled HLO is the ground truth for "is my model running in FP8". .compile().as_text() plus counting custom-call targets takes a minute.
When a compiler pass produces the right result for one part of a program and silently not for another, look at how the pass iterates before looking at what the parts contain.
A fallback path that makes a program correct but slow is the hardest kind of bug to notice. Count kernels, not just check outputs.

The fix is open as openxla/xla#49007.

Accidentally quadratic: buffer copies made MCTS in DeepMind's mctx 3 slower

Yehor Cherednichenko — Wed, 22 Jul 2026 12:14:20 +0000

I'm training an AlphaZero-style agent (Gumbel MuZero via DeepMind's mctx) to lay out working factory modules for Factorio: the network places machines, belts and inserters on a grid, and the reward comes from an exact throughput verifier. Everything runs in JAX on a single RTX 5070: 128 environments in one batch, an action space of A = 1729 (3 entity types × 144 cells × 4 rotations + "done"), and a small 474k-parameter conv net in bf16.

While benchmarking training configurations I hit this:

MCTS simulations per move	training throughput	XLA compile time
16	143 episodes/s	5 s
32	47 episodes/s	13 s
64	9 episodes/s	60 s

Doubling the simulation budget should roughly double the cost — each simulation is one network call plus some tree bookkeeping. Instead, 16→32 costs ×3 and 32→64 costs ×5. Something in the search was superlinear, and this post is the story of finding it in the compiled HLO and fixing it by rewriting one ~80-line function (PR #116), with bitwise-identical search results.

Ruling out the network

First, components in isolation (batch 128): one network evaluation takes ~0.8 ms, and the rest of recurrent_fn (environment step + observation + legal-action mask) adds almost nothing on top — the whole function is also ~0.8 ms. So at 64 simulations the network accounts for roughly 50 ms per move. But a full policy step at 64 simulations costs 362 ms. Hundreds of milliseconds were going somewhere else.

To localize them I benchmarked three variants of the same policy step:

full — production setup;
no-net — network replaced by constant logits, real environment;
tree-only — no network and no environment: recurrent_fn returns the embedding unchanged. Nothing left but mctx's own tree machinery.

sims	full	no-net	tree-only
8	10.7 ms	3.7 ms	3.8 ms
16	20.9 ms	8.3 ms	8.3 ms
32	64.7 ms	34.0 ms	33.7 ms
64	362.1 ms	124.7 ms	125.0 ms

The pure tree machinery is superlinear all by itself. Per simulation it costs 0.47 → 0.52 → 1.05 → 1.95 ms as the budget goes 8 → 16 → 32 → 64: the per-simulation cost doubles every time the number of simulations doubles. That's O(N²) machinery.

What grows with N inside a simulation

mctx stores the search tree as struct-of-arrays: six [B, N+1, A] buffers (children_index, children_prior_logits, children_values, children_visits, children_rewards, children_discounts). With my shapes at 64 simulations that's six 57.5 MB buffers, ~345 MB of tree state. Each simulation runs three phases: simulate (walk down the tree in a while_loop), expand (evaluate the leaf), backward (walk back up, updating value statistics).

backward is the interesting one. It carries the entire tree as the lax.while_loop state, and each iteration both reads from children_values/children_visits (gathers for the parent update) and scatters single elements back into the same buffers.

Reading the compiled HLO

JAX makes it easy to look at what XLA actually compiled:

txt = jax.jit(policy_step).lower(params, key, states).compile().as_text()

A grep for copy ops on the tree-shaped buffers turned up this in the backward loop body (this dump is from a 16-simulation compile, so the middle dimension is N+1 = 17):

%copy.1 = f32[128,17,1729]{2,1,0} copy(%get-tuple-element.512)
%copy   = s32[128,17,1729]{2,1,0} copy(%get-tuple-element.513)

Those are children_values (f32) and children_visits (s32) — copied in full, on every iteration of the backward loop, i.e. on every step of the leaf-to-root walk. XLA:GPU cannot prove the scatter can safely alias the input buffer (the same tensor is gathered from and scattered into within one loop iteration), so it materializes a defensive copy of each buffer, each iteration.

Now the quadratic makes sense, because two things grow with the simulation budget N:

The buffers: [B, N+1, A] grows linearly with N.
The walk length: with Gumbel sequential halving, the champion root action receives ~N/4 visits, and the deterministic interior selection extends a chain below it. Measured max tree depth (from search_tree.parents): ~1 at 16 sims, ~4 at 32, up to 15 at 64 with a rained network. And since backward is vmapped, every simulation pays for the deepest batch element out of 128.

Copies of size O(N·A) times a depth that grows with N — there's the superlinear term. It also explains the compile-time blowup (5 → 13 → 60 s): XLA's buffer assignment and scheduling passes work over ever-larger buffers.

The fix

The backward pass doesn't need the whole tree in its loop carry. The patch (PR #116) walks the leaf-to-root path with only O(num_nodes) state — the path indices, actions, and the updated per-node statistics — and then applies each statistic with a single scatter per array after the loop:

def body_fun(state):
  index = state.index
  parent = tree.parents[index]
  action = tree.action_from_parent[index]
  # ... same value recursion as before, into small path buffers:
  return _BackwardState(
      step=step + 1, index=parent, leaf_value=leaf_value, ...
      path_parents=state.path_parents.at[step].set(parent),
      path_node_values=state.path_node_values.at[step].set(parent_value),
      ...)

end = jax.lax.while_loop(cond_fun, body_fun, initial_state)
return tree.replace(
    children_values=tree.children_values.at[
        end.path_parents, end.path_actions].set(
            end.path_children_values, mode="drop"),
    children_visits=...,   # same single-scatter treatment for the other
    node_values=..., node_visits=...)  # three updated arrays

Unused path slots point at an out-of-range index and are dropped by the
scatters. The value recursion is arithmetically identical and runs in the
same order, so this is not an approximation:

same PRNG key → bitwise-identical action and action_weights from gumbel_muzero_policy (checked at 16 and 32 simulations, batch 128);
the full mctx test suite passes, including the golden-tree comparisons that directly validate backward's outputs;
in my trainer, fixed-seed A/B runs produce identical metric traces — only the throughput changes.

Results

Policy step, interleaved A/B in a single process, trained checkpoint:

sims	before	after	speedup
16	29.3 ms	27.0 ms	×1.08
32	72.3 ms	48.5 ms	×1.49
64	394.7 ms	120.8 ms	×3.27

Per-simulation cost after the fix is flat — 1.5–1.9 ms across 16/32/64
simulations — instead of growing with the budget.

End-to-end training (full loop with SGD and reward verification, 4096
episodes per config):

sims	before	after	speedup	compile
16	138.6 eps/s	152.0 eps/s	×1.10	5.9 → 5.4 s
32	46.6 eps/s	73.6 eps/s	×1.58	12.7 → 8.7 s
64	9.1 eps/s	29.8 eps/s	×3.27	62.9 → 19.4 s

Until the PR lands you can monkeypatch — the function is module-level and
looked up at trace time:

import mctx._src.search as msearch
msearch.backward = backward_fast  # install before the first jit trace

Four things I got wrong along the way

1. I first blamed the action space. With A = 1729 the "obvious" fix is
Sampled-MuZero-style top-K action candidates (search over K = 128 candidates instead of the full space). I prototyped it — and on top of the backward fix it gained nothing (117 vs 121 ms at 64 sims). The full-buffer copies were the A-scaled term; once they're gone, per-node work over 1729 actions is cheap. Know your actual bottleneck before reaching for the textbook fix.

2. max_depth looked like the lever, and wasn't. Capping tree depth
attacks the walk length, but only trims the tail: 394.7 → 337.4 ms alone,
120.8 → 111.2 ms on top of the fix.

3. Cross-process GPU benchmarks kept lying to me. My RTX 5070 idles at
495 MHz and boosts to 2932 MHz under load. The same compiled graph measured 34 ms in one process and 57 ms in another, depending on clock ramp-up — a ×1.7 phantom difference that sent me chasing ghosts (an early version of this very fix looked slower than baseline for exactly this reason). All the A/B numbers above come from interleaved measurements inside a single process (variant A, B, C, A, B, C…, medians over rounds).

4. Benchmark with a trained network, not just a fresh one. A trained
(peaked) policy builds much deeper trees than a randomly initialized one: at 64 sims the deepest tree in the batch reached 15 levels vs 10, and the
per-example maximum averaged 11.6 vs 5.6. Deeper trees mean more backward
iterations, so the superlinear term grows as the agent improves — the
same-process checkpoint-vs-random gap at 64 sims was 394.7 vs 341.8 ms
pre-fix. Fittingly, my own first estimate of this effect was a scary ×2.2,
measured across two processes — i.e. contaminated by exactly the clock-drift artifact from the previous point. Fresh-net benchmarks still measure the cheapest point of your training run.

Takeaways

If a JAX while_loop carries big buffers and both reads and writes them in one iteration, check the HLO for defensive copies. .compile().as_text() plus grep is a ten-minute test that can be worth a ×3 speedup.
Decompose before optimizing: the net/no-net/tree-only split localized the problem in three benchmark runs.
Caveat: measured on GPU (CUDA). On TPU the buffer assignment may already handle this; the patch should be neutral there.

The fix is open as google-deepmind/mctx#116.

DEV Community: Yehor Cherednichenko

Half of my FP8 convolutions were silently running in f32 – a one-line XLA fix

A stray return made XLA:GPU skip FP8 rewriting for most of a module

Counting custom calls

The wrong hypothesis first

Reading the pass

Checking it end to end

Who is affected

Two things I got wrong along the way

Takeaways

Accidentally quadratic: buffer copies made MCTS in DeepMind's mctx 3 slower

Ruling out the network

What grows with N inside a simulation

Reading the compiled HLO

The fix

Results

Four things I got wrong along the way

Takeaways

A stray `return` made XLA:GPU skip FP8 rewriting for most of a module