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