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

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:

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:

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:

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.