Batched MCTS inference

Deep dive

How it works

Parts:

• tree — nodes with prior, visits N, value sum W, in-flight count (virtual loss)
• collector — PUCT descent that applies virtual loss and stops at an unexpanded leaf
• batch builder — packs the collected leaves into one contiguous input tensor
• evaluator thread — owns the GPU; runs the network on whole batches
• scatter step — expands leaves with priors, backs up values, removes virtual loss

Operations: collect, batch, evaluate, scatter, overlap (steps below).

1. Collect. For each tree, descend with PUCT where every in-flight visit counts as an extra visit with value −1 (virtual loss): Q' = (W − v·n_vl) / (N + n_vl), N' = N + n_vl. Increment n_vl along the path. Stop at an unexpanded leaf.
2. Terminal leaves are backed up immediately; collisions (leaf already in flight) are dropped (or counted as extra visits to back up later, Lc0's "collision visits").
3. Batch. Encode every collected leaf into a contiguous input tensor; send (inputs, count, reply channel) to the evaluator.
4. Evaluate. The evaluator thread runs the net on the whole batch (one kernel sequence), returns policy priors and values.
5. Scatter. Expand each leaf with its priors, back up its value along its path (sign flips per ply), and remove the virtual loss.
6. Overlap. Keep 2+ batches in flight: while the GPU runs batch k, the CPU collects batch k+1 (double buffering), or run several search threads feeding one queue.

 games/trees ──collect leaves (virtual loss)──► batch builder ──► queue ──► GPU thread
      ▲                                                                       │
      └──────── expand + backup (remove virtual loss) ◄── results ◄───────────┘

Code

use std::{sync::mpsc, thread};

/// Game interface: value is from the side to move's point of view.
pub trait Game: Clone {
    fn legal_moves(&self) -> Vec<u16>;
    fn play(&mut self, m: u16);
    fn terminal_value(&self) -> Option<f32>; // Some(v) if the game is over
    fn encode(&self, out: &mut Vec<f32>);   // network input planes
}

struct Node { prior: f32, n: u32, w: f32, vloss: u32, children: Vec<(u16, u32)>, expanded: bool }

pub struct Tree<G: Game> { root_state: G, nodes: Vec<Node> }

const C_PUCT: f32 = 1.5;
const VLOSS: f32 = 1.0; // each in-flight visit counts as one loss

impl<G: Game> Tree<G> {
    pub fn new(state: G) -> Self {
        let root = Node { prior: 1.0, n: 0, w: 0.0, vloss: 0, children: vec![], expanded: false };
        Tree { root_state: state, nodes: vec![root] }
    }

    /// PUCT with virtual loss: in-flight visits raise N and lower Q, so the next
    /// descent in the same batch prefers other branches.
    fn select_child(&self, id: u32) -> (u16, u32) {
        let p = &self.nodes[id as usize];
        let sqrt_n = ((p.n + p.vloss) as f32).sqrt();
        let score = |&(_, c): &(u16, u32)| {
            let ch = &self.nodes[c as usize];
            let n = (ch.n + ch.vloss) as f32;
            // child w is from the child's side to move: negate for the parent
            let q = if n > 0.0 { (-ch.w - VLOSS * ch.vloss as f32) / n } else { 0.0 };
            q + C_PUCT * ch.prior * sqrt_n / (1.0 + n)
        };
        *p.children.iter().max_by(|a, b| score(a).total_cmp(&score(b))).unwrap()
    }

    /// Descend to a leaf, applying virtual loss on the path. Returns the path and the
    /// leaf state, or None if the leaf was terminal (backed up immediately) or a collision.
    pub fn collect_leaf(&mut self) -> Option<(Vec<u32>, G)> {
        let mut state = self.root_state.clone();
        let mut path = vec![0u32];
        let mut id = 0u32;
        while self.nodes[id as usize].expanded {
            if self.nodes[id as usize].children.is_empty() { break; }
            let (m, c) = self.select_child(id);
            state.play(m);
            id = c;
            path.push(id);
        }
        if let Some(v) = state.terminal_value() { self.backup(&path, v, false); return None; }
        if self.nodes[id as usize].vloss > 0 { return None; } // already in flight: collision
        for &n in &path { self.nodes[n as usize].vloss += 1; }
        Some((path, state))
    }

    /// Expand the leaf with network priors and back up its value, removing virtual loss.
    pub fn apply_eval(&mut self, path: &[u32], state: &G, priors: &[f32], value: f32) {
        let leaf = *path.last().unwrap();
        let moves = state.legal_moves();
        for (i, &m) in moves.iter().enumerate() {
            let c = self.nodes.len() as u32;
            self.nodes.push(Node { prior: priors[i], n: 0, w: 0.0, vloss: 0, children: vec![], expanded: false });
            self.nodes[leaf as usize].children.push((m, c));
        }
        self.nodes[leaf as usize].expanded = true;
        self.backup(path, value, true);
    }

    fn backup(&mut self, path: &[u32], mut v: f32, had_vloss: bool) {
        for &n in path.iter().rev() {
            let node = &mut self.nodes[n as usize];
            node.n += 1;
            node.w += v;
            if had_vloss { node.vloss -= 1; }
            v = -v; // parent sees the opposite sign
        }
    }

    pub fn best_move(&self) -> u16 {
        let r = &self.nodes[0];
        r.children.iter().max_by_key(|&&(_, c)| self.nodes[c as usize].n).unwrap().0
    }
}

/// A batch request: flat input planes for `count` positions.
pub struct Batch { pub inputs: Vec<f32>, pub count: usize, pub reply: mpsc::Sender<Vec<(Vec<f32>, f32)>> }

/// One search step over many games: gather up to `per_tree` leaves from each tree,
/// send one batch to the evaluator thread, scatter the results back.
pub fn step<G: Game>(trees: &mut [Tree<G>], per_tree: usize, eval_tx: &mpsc::Sender<Batch>) {
    let mut pending: Vec<(usize, Vec<u32>, G)> = Vec::new();
    let mut inputs = Vec::new();
    for (t, tree) in trees.iter_mut().enumerate() {
        for _ in 0..per_tree {
            if let Some((path, state)) = tree.collect_leaf() {
                state.encode(&mut inputs);
                pending.push((t, path, state));
            }
        }
    }
    if pending.is_empty() { return; }
    let (tx, rx) = mpsc::channel();
    eval_tx.send(Batch { inputs, count: pending.len(), reply: tx }).unwrap();
    let results = rx.recv().unwrap(); // in a real engine: overlap with collecting the next batch
    for ((t, path, state), (priors, value)) in pending.into_iter().zip(results) {
        trees[t].apply_eval(&path, &state, &priors, value);
    }
}

/// Evaluator thread: owns the network (GPU context). Here a stand-in closure.
pub fn spawn_evaluator(net: impl Fn(&[f32], usize) -> Vec<(Vec<f32>, f32)> + Send + 'static)
    -> mpsc::Sender<Batch> {
    let (tx, rx) = mpsc::channel::<Batch>();
    thread::spawn(move || {
        for b in rx { let out = net(&b.inputs, b.count); let _ = b.reply.send(out); }
    });
    tx
}

Toy test (64 Nim games, about 1,600 visits per root): 8 leaves per tree per batch found the winning move in 46 of 64 games, 1 leaf per tree in 55 of 64. That is the virtual-loss tax that parallel games avoid.

Complexity and constants

• Host↔device copies: inputs are small (e.g. 112 planes × 8×8 for Lc0-style chess input ≈ 7 KB per position as f32, much less as bit-packed planes); use pinned memory and async copies.
• Memory: tree nodes order of 32-64 bytes each plus children arrays; a node per evaluated leaf.
• Selection: O(depth × branching) per leaf, CPU-bound; for chess-sized trees order of µs per leaf, so one CPU core keeps up with a few thousand evals/s and more cores or games are needed for large GPUs.
• Network: one call per batch; throughput ≈ B / latency(B). Measure latency(B) for your net: it is often nearly flat from B = 1 to order of 64-256, then grows linearly once the GPU is saturated.

Efficiency

• Cost: per leaf, selection is order of µs on one CPU core and the network call dominates; throughput ≈ B / latency(B). AlphaGo Zero used mini-batches of 8 at 1,600 simulations per move (Silver et al. 2017). Bound: GPU latency when the batch is small, CPU selection and host-to-device copies when it is large.
• Levers:
  • batch across independent searches — no virtual-loss bias, GPU filled — many parallel self-play games or problems
  • virtual loss inside one tree — fills a batch from a single search, at a cost in visit quality (toy test above) — one big analysis tree
  • evaluation cache keyed by hash — skips repeated states — transpositions, repeated openings
  • double buffering (two batches in flight) — hides collect time behind GPU time — GPU idle otherwise
  • padded, precompiled batch sizes (CUDA graphs, TensorRT) — lower launch overhead — fixed-shape engines
• Naive → standard → state of the art:
  • one evaluation per call · latency-bound, GPU mostly idle
  • one tree, batch of 8 with virtual loss · AlphaGo Zero setting
  • many parallel searches, async evaluator, cache, fixed batch shapes · GPU-bound; JAX libraries such as mctx batch the search itself on the accelerator

Implementation notes

• Virtual loss size: 1-3 losses per in-flight visit. Larger values spread a batch more but bias it more; tune with the batch size.
• Many games, one leaf each: with G games ≥ batch size, take one leaf per game and skip virtual loss entirely; the search is then identical to sequential MCTS. KataGo and Lc0 self-play run many games per process this way.
• Eval cache: hash positions (Zobrist) and cache (policy, value); transpositions and repeated openings in self-play hit often. Check the cache before adding a leaf to the batch.
• Batch assembly: write inputs straight into a preallocated pinned buffer at index i; keep a parallel array of (tree, path) for scattering. No per-leaf allocation.
• Async queue: an MPSC channel or a bounded MPMC queue feeds one evaluator per GPU; it waits for "batch full or timeout of order of 1 ms".
• Fixed batch shapes: inference engines (TensorRT, ONNX Runtime with CUDA graphs) are fastest with a few precompiled batch sizes; pad partial batches up to the next size.
• Lock-free trees with threads: if several threads search one tree, apply virtual loss with atomic adds on N and vloss; the code above is single-threaded per tree, which is simpler and enough when parallelism comes from games.

Pitfalls

• Forgetting to remove virtual loss on collisions, terminal leaves or dropped leaves: the path is permanently penalised and the tree stops visiting a good move.
• Value sign errors across plies when backing up a batch out of order: always back up from the leaf's side-to-move perspective and flip per level.
• Batch built from one tree with a huge virtual loss early in the search: all leaves spread over bad moves; ramp batch size with root visits or cap leaves per tree.
• Dirichlet noise and temperature applied per batch instead of per root: self-play becomes deterministic across games sharing a batch.
• Synchronous design (collect, wait, scatter) leaves both CPU and GPU idle half the time; keep at least two batches in flight.

Variants and alternatives

In production

• Board games: AlphaGo/AlphaGo Zero/AlphaZero (virtual loss, TPU batches), Leela Chess Zero (minibatches per search thread, collision handling, multiplexing backends), KataGo (many parallel self-play games and search threads sharing one batched GPU evaluator).
• Chemistry: retrosynthesis planning with MCTS guided by neural policy networks (Segler et al. 2018).
• Algorithm discovery: AlphaTensor (faster matrix multiplication) and AlphaDev (faster sorting routines) search with MCTS over a learned network.
• Video compression: MuZero-based rate control for VP9 (Mandhane et al. 2022).
• LLM reasoning: RAP runs MCTS over reasoning steps with the language model as world model and evaluator (Hao et al. 2023).

Implementations

Sources

• David Silver et al., Mastering the game of Go with deep neural networks and tree search, Nature, 2016. paper
• David Silver et al., Mastering the game of Go without human knowledge, Nature, 2017. paper
• Guillaume M. J.-B. Chaslot, Mark H. M. Winands, H. Jaap van den Herik, Parallel Monte-Carlo Tree Search, Computers and Games, 2008. paper
• Richard B. Segal, On the Scalability of Parallel UCT, Computers and Games, 2010.
• David J. Wu, Accelerating Self-Play Learning in Go, arXiv:1902.10565, 2019. paper · code
• Anji Liu, Jianshu Chen, Mingze Yu, Yu Zhai, Xuewen Zhou, Ji Liu, Watch the Unobserved: A Simple Approach to Parallelizing Monte Carlo Tree Search, ICLR 2020.
• Marc Segler, Mike Preuss, Mark P. Waller, Planning chemical syntheses with deep neural networks and symbolic AI, Nature 555, 2018. paper
• Amol Mandhane et al., MuZero with Self-competition for Rate Control in VP9 Video Compression, arXiv:2202.06626, 2022. paper
• Shibo Hao et al., Reasoning with Language Model is Planning with World Model, arXiv:2305.14992, 2023. paper
• Fawzi et al., Discovering faster matrix multiplication algorithms with reinforcement learning, Nature 610, 2022. paper
• Mankowitz et al., Faster sorting algorithms discovered using deep reinforcement learning, Nature 618, 2023. paper