GPUTreeIG: optional CUDA execution

TreeIG on the CPU is already fast enough for most applications and remains the default. When attribution speed matters and an NVIDIA GPU is available, GPUTreeIG can be materially faster: the T4 comparisons below show roughly 9–20× speedups on the reported workloads. The gain depends on the problem, and small workloads can still favor the CPU.

The treeig package also provides GPUTreeIG, an optional CUDA implementation for large, repeated prediction-attribution workloads with a fixed model and baseline distribution. It uses the same attribution semantics as TreeIG, up to floating-point accumulation differences. It is not a separate package.

On a host with an NVIDIA CUDA GPU, install the optional CUDA support:

pip install "treeig[cuda]"

This installs numba-cuda; a compatible NVIDIA driver and CUDA toolkit are also required. Follow the Numba CUDA installation instructions for your platform. CUDA is loaded only when constructing GPUTreeIG; ordinary TreeIG use requires no CUDA installation. Apple Metal and AMD ROCm are not supported by this backend. If CUDA is unavailable, GPUTreeIG raises an error; choose TreeIG for CPU execution.

import treeig as tig

gpu_ig = tig.GPUTreeIG(
    model,
    baseline=background.rows,
    baseline_weights=background.weights,
)
phi = gpu_ig.attribute(large_batch, batch_size=1000)

Construction uploads the model, baseline distribution, and weights. Reuse the instance across calls to retain that state and reuse device buffers. Each call transfers input rows, initializes output buffers, and retrieves results. Let the instance go out of scope or use del gpu_ig when it is no longer needed; device memory may remain in the CUDA allocator cache.

GPUTreeIG supports attribute() and explain() for prediction attribution. The baseline distribution, weights, and target are fixed at construction. Observation batching is supported; baseline batching, per-baseline output, loss attribution, traces, diagnostics, and warmup() are not. Models requiring stack or leaf-segment scratch widths above 1,024 entries are rejected.

T4 versus CPU benchmarks

The following recorded Google Colab Tesla T4 comparisons illustrate workloads where GPU execution helped. They are examples, not a promise of speedup on other hardware or problems. CPU TreeIG remains the default.

Here n is the number of observations and B is the number of weighted baseline rows. The persistent explainer retains its model and baseline state on the GPU and reuses its work buffers. Timings are medians of three warmed calls; model fitting, parsing, baseline preparation, and JIT compilation are outside the timed region. GPU calls include observation transfers and result retrieval.

Shallow gradient boosting

The recorded persistent run used a 100-tree, depth-3 boosting ensemble:

Observations (n)	Baselines (B)	CPU (ms)	T4 (ms)	CPU time / T4 time
10	100	26.00	2.69	9.68×
100	10	24.95	2.71	9.22×
100	100	252.24	15.22	16.58×
1,000	10	259.69	15.22	17.06×
1,000	100	2,415.75	137.41	17.58×

Maximum absolute CPU/GPU attribution difference was below 2.0e-14 in the recorded persistent comparison.

Depth-6 random forest

A later run used the current, smaller per-thread scratch buffers. The recorded comparisons for the 100-tree forest were:

Observations (n)	Baselines (B)	CPU (ms)	T4 (ms)	CPU time / T4 time
10	100	≈41.78	3.74	11.17×
100	10	≈41.72	3.53	11.82×
100	100	≈385.16	22.63	17.02×
1,000	10	≈445.14	23.64	18.83×
1,000	100	≈4,092.65	207.96	19.68×

CPU times marked ≈ are reconstructed as T4 time × recorded speedup. The notes retain these two rounded quantities rather than raw CPU timings for this later run, so the reconstructed CPU times are approximate. Across the scratch-specialization runs, which also included depth-8 and depth-10 forests, the maximum absolute CPU/GPU difference was 2.62e-14.

Small workloads can favor the CPU because GPU launch and transfer costs dominate. Compare warmed calls on your own model and baseline distribution when deciding whether to use the optional backend. CUDA simulator checks validate semantics; simulator timings are not GPU performance evidence.

The full CUDA benchmark notes preserve the earlier stateless and persistent measurements and the scratch-buffer comparison. The benchmark script is available in the repository. These are historical measurements, not new runs performed for this documentation update.