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

I'm training an AlphaZero-style agent (Gumbel MuZero via DeepMind's ), 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:

1. 
   The buffers: [B, N+1, A] grows linearly with N.


2. 
   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 (.