Skip to main content

Module cpu_pool

Module cpu_pool 

Source
Expand description

A persistent CPU worker pool parked on a spin-then-park barrier.

§Why this exists

Decode opens a parallel region per weight matrix, per layer, per token – roughly seven per layer. Rayon answers each of those with a fork-join: a job is pushed onto the global injector, sleeping workers are woken through a mutex and a condvar, and the caller then waits on a latch. That is a fixed per-region cost, so the smaller the model the larger its share: measured at ~75% of decode wall time at 135M parameters and ~9% at 8B (issue #27).

llama.cpp does not fork. ggml_threadpool starts N workers once and parks them on a spin barrier; a graph node is published by bumping a counter the workers are already watching, and each worker pulls chunks off one shared atomic until the work is gone. Waking a spinning thread costs a cache-line transfer instead of a futex.

This module is that shape, in safe-by-construction Rust:

  • CpuPool::new spawns the workers once and keeps them.
  • CpuPool::run publishes n_tasks and a type-erased closure, bumps the epoch, then participates in draining the task counter alongside the workers.
  • Workers spin for [spin_window] and then park on a condvar, so an idle ferrox-server does not burn a core per worker. That bound is the whole reason this is not a plain spin barrier.

§What makes it sound

The one dangerous thing here is that workers dereference a pointer to a closure the submitter owns. Two rules keep that from being a use-after-free, and both are enforced in CpuPool::run:

  1. A region ends only when every worker has checked out. active is set to the worker count before the epoch bump and decremented by each worker after its last touch of the job; run does not return until it reads zero. So the closure outlives every use.
  2. Only one region at a time. submit is a mutex, and run tries it rather than blocking: a second thread that arrives while a region is in flight is told false and falls back to rayon rather than queueing behind it.

Re-entrancy is the third hazard – a task closure that itself opens a region would deadlock against rule 2 – so CpuPool::run runs nested regions inline on the calling thread.

§What this module is NOT

It is not a general work-stealing runtime. There is no task graph, no nested parallelism, no join. It runs one flat 0..n_tasks loop at a time, because that is the entire shape of a quantized matvec and the shape llama.cpp’s threadpool has.

Structs§

CpuPool
A persistent pool of parked workers.

Functions§

in_region
Whether the calling thread is currently running pool tasks.