10. GPU training

The idea

Everything expensive in histogram GBT is a data-parallel reduction or scatter over rows: bin the raw values, sum gradients into cells, scan cells for the best split, route rows to children. A GPU has ~10× the memory bandwidth of a CPU socket, and bandwidth is exactly what these loops are starved for. What a GPU does not have is cheap access to the sequential, branchy part: deciding which nodes to split, bookkeeping the tree, enforcing constraints.

So the design question is not "port the algorithm to CUDA" but "draw the boundary": the device owns every per-row loop, the host owns every per-node decision, and the two exchange the smallest possible messages. bonsai has a CUDA grower for each of its three growth strategies, cuda_depthwise, cuda_leafwise, and cuda_levelwise, and each draws that boundary in its own way. depthwise and levelwise share a boundary called the level transaction (decision 53): the CPU and CUDA engines implement the same transactions, so a grower cannot tell which backend it is running on. leafwise draws an analogous boundary, one node per round instead of one level, through its own LeafStep seam (decisions 95 to 98). This chapter walks through the level transaction first, since it covers two of the three CUDA growers, then the leaf step. The problems that make the device side interesting, atomics and precision, keeping data resident, and knowing what to move next, apply to all three.

The math

Histogram accumulation is a race by construction. Thousands of threads add (g_i, h_i) into 255 shared cells, so the adds must be atomic, and the order they land in is whatever the scheduler produced. Floating-point addition is not associative, so a float cell carries an order-dependent error:

\[ \Big(\sum_i g_i\Big)_{\text{GPU}} \;=\; \Big(\sum_i g_i\Big)_{\text{exact}} + \varepsilon, \qquad |\varepsilon| \lesssim n\,u\,\max|g_i| \]

with \(u\) the unit roundoff, and two runs of one fit differ in the last ulps. Integer addition is associative, so the cells are int64 fixed point. Once per tree the fill reads \(M = n \max_i |g_i|\) and picks the power-of-two scale \(s = 2^{61 - \lfloor \log_2 M \rfloor}\); each visit adds \(\mathrm{round}(g_i s)\), and any subset of the tree's rows sums to under \(2^{62} + n/2\), inside int64 whatever order the adds land in. A power-of-two scale makes the multiply exact, so the one rounding per row is the integer conversion, below \(\max|g| \cdot n / 2^{62}\). Two consequences to design for: a device fit is bit-identical to itself on one device and build, and it matches the CPU fit, which accumulates float cells, to tolerance, not bit-exactly.

What the scheme costs. Less than the float fill it replaced. A 64-bit shared-memory atomic lowers to a compare-and-swap loop, so the shared stage adds each cell as two native 32-bit atomics with the carry taken from the first; the tile's shared footprint doubles from 4 to 8 bytes a cell, and the fill block grows to 512 threads to keep the SM full. The same-pod A/B in decision 124 prices the whole scheme at a 7.6 to 8.6% faster fit on an L40S. The float-per-chunk plus double-merge scheme it replaced bounded the float error term at 32k rows a chunk; the integer sum has no such term, so the chunk count follows occupancy alone. Split scoring happens in double on cells dequantised once, by one multiply, at the start of the scan.

The subtraction trick survives intact (chapter 2): children of a split partition the parent's rows, so the device builds the smaller child's histogram and derives the larger by a cell-wise integer subtract kernel, exact, on resident buffers, without the host ever seeing a histogram.

In bonsai: the level transaction

The entire backend is one CUDA C++ translation unit, src/cuda/histogram_engine.cu plus its kernel header src/cuda/detail/kernels.cuh, compiled by the project's own clang (-x cuda), not nvcc, same C++23, same libc++ as the rest of the build. Builds without CUDA link a stub that throws; bonsai::cuda_available() is the runtime predicate and the cuda_* growers are registered everywhere, so a GPU-trained model predicts fine on a CPU-only binary.

Follow one fit through the transactions:

• Ingest: cuda_ingest (decision 54): raw feature values stream to the device in ~64MB chunks and a kernel bins them with a lower_bound that reproduces the host transform exactly: same cuts, same comparisons, bit-identical bin ids. The product, an IngestPlane, rides on the Dataset as an opaque receipt; host binned columns are never materialized unless a host consumer asks. At 16M×100 this replaced 4.6s of host binning plus a 1.6GB upload with ~0.5–0.9s of transfer+kernel (host-dependent).
• begin_tree: the per-tree gradient upload, interleaved into (g,h) pairs on device.
• open_level (find): level histograms live in a slot-indexed device buffer that ping-pongs between parent and child levels. The find kernel gives each (node, feature) pair one warp: a shuffle-based prefix scan over the cells, per-lane gain scoring, and a shuffle argmax that reduces to the best split with a fixed tie-break. Only the per-node decisions (feature, bin, gain, child sums, a few hundred bytes) cross back to the host.
• apply_level (partition): route/count/scan/scatter kernels move each split node's row segment into stable left/right children entirely on device; the host receives two integers per split (the child row counts). Stability matters: it is what keeps the device partition semantically identical to the host one.
• Histogram build: the shared-memory chunked kernel for large children, a direct-to-global kernel for small ones (a full shared-memory pipeline per 40-row node is all overhead), then the subtract kernel for every larger sibling.
• end_tree: the epilogue (decision 53): the host sends the finished tree's node-value table (a few KB); a kernel maps every row's resident leaf assignment to its training value; values and leaf ids come home in two bulk copies. Before this transaction existed, the host looped over 16M rows per tree.

The host control plane between transactions is src/level_step.hpp and the growers, the same plan_level/commit_children logic the CPU path uses, because it is the CPU path's logic. When the device declines (a feature's bins exceed the shared-memory budget), the same grower runs the same tree on the host engine mid-fit.

What this buys is measured on the perf standings, which own the digits and are refreshed with each release. At the gpu-tall scenario every bonsai device plane finishes the fit ahead of the reference GPU trainer it is paired with, in about a third of the peak host memory. Much of that margin is ingest rather than boosting: bonsai bins on the device, where the references sketch on the host and ship the result over the bus. The path from 3× slower to ahead is chapter 11.

In bonsai: the leaf step

cuda_leafwise grows the leafwise grower's best-first strategy fully on device: histograms live in a per-tree slot pool instead of the level plane's ping-pong buffers, and the gain heap that picks which leaf to expand next stays on the host, the same structure the CPU leafwise grower already uses. The design record, including the slot-pool layout, is decisions 95 to 98 in the decision log; the adopt/refuse reading of LightGBM's learner is a performance engineering case study.

It is what device="cuda" runs by default: the Python estimators default to grower="leafwise", and the device mapping is a prefix (leafwise becomes cuda_leafwise) that preserves the growth strategy rather than substituting one. Pass grower="depthwise" (CPU) or grower="cuda_depthwise" (GPU) explicitly to select the fastest plane at matched knobs instead.

Leaf-wise and depthwise agree structurally at a capped depth with a full leaf budget: the tree that comes out is the same tree either way. The two strategies only differentiate at uncapped depth or under a small leaf budget, where best-first order decides which nodes get split at all.

Where the leaf plane stands against LightGBM's CUDA leaf-wise is on the perf standings: bonsai trains faster at every published GPU scenario, on less host and device memory, and LightGBM scores the better test r2 on those same cells. Uncapped depth is the cell to read carefully, because best-first there must find a split for every leaf it creates and so pays about twice the rounds; the last measurement of it still put bonsai ahead (decision 100), which carries the reasoning.

The device-resident objective (MSE, LogLoss, or Poisson; no DART; all-rows or Bernoulli sampling) arms for cuda_leafwise the same way it arms for cuda_depthwise and cuda_levelwise. There is nothing to configure: an eligible fit keeps labels and scores on the device for the whole fit, and BONSAI_HOST_OBJECTIVE=1 forces the host path.

Try it

No CUDA device on your machine? The RunPod runbook gets you a validated GPU session for well under a dollar.

# The device suite: SKIPs without a GPU, exercises real kernels with one.
./build-cuda/tests/bonsai_tests "[cuda]"

# Any fit, with the device profile lines:
BONSAI_CUDA_PROFILE=1 BONSAI_GROW_PROFILE=1 \
  bonsai fit --config configs/california_housing.toml \
  --set dispatch.grower_name=cuda_depthwise

Read the cuda-profile line first: upload vs gpu tells you where the fit's time went. A fit whose max_bin pushes a feature's histogram past the device's shared-memory ceiling does not appear here at all, because it does not run: the fit stops with an error naming the limit, and the remedy is a lower max_bin or device="cpu". Then cuda-upload-decomp splits every transfer by transaction: this is the line every optimization in chapter 11 was priced against.

# CPU vs GPU on the same data — expect tolerance-level agreement, not bit-equality:
uv run scripts/compare.py --config configs/california_housing.toml \
    --growers depthwise,cuda_depthwise

Gotchas & war stories

• The allocator can cost more than the kernels. On GeForce drivers, the default CUDA mempool returns freed memory to the OS at every sync, and the resulting alloc/free churn synchronizes the whole process: a measured 11–14 seconds per fit on an RTX 5090 before decision 48 pinned the pool's release threshold. If a GPU fit is inexplicably slow, suspect memory management before kernels.
• Some rented GPUs are just broken. A whole class of hosts showed ~300µs per device synchronization (healthy: ~4µs), invariant to every software knob, an ASPM/IRQ-level defect. Multiply by ~25k syncs per fit and the machine is unusable. The bench harness front-loads a 30-second sync probe and rejects hosts above 50µs (decision 48).
• A demoted split is not the same thing on both planes. demote_empty_splits un-splits a node whose child would be empty. On the host that is bookkeeping; on the device the rows were already scattered and stamped, so demoting orphaned their leaf assignments, one mega-leaf per tree at the rightmost spine, quality collapsing only at 16M rows where double-precision sums stopped masking it (PR #29). Control-plane operations must know which plane's invariants they touch.
• "First use is single-threaded" is an assumption, not a law. The lazy host materialization of device-binned columns was first guarded by nothing, on the precedent of another lazy member, but its first consumer runs inside a parallel loop, and the race corrupted the heap intermittently (one clean validation pass meant nothing). It is now a call_once (PR #37). On the same theme: while the cells were float, two GPU fits of identical inputs differed in atomic order, and the leafwise parity test flaked one run in five at the 1e-4 bound as the tail of that spread; a distribution's tail is not a tolerance's job to absorb, so the cells went fixed point (decision 124).
• Byte-identity is available where it matters. Device binning is bit-identical to host binning (same cuts, same lower_bound, integer outputs, no atomics), which is what lets a device-binned dataset train the same model as a host-binned one and lets the test suite say so exactly. A device fit is bit-identical to itself run to run (integer cells), so the suite compares model bytes, not predictions within a bound.