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

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:

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.

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
    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:

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):

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.