hessboost 0.2.2__tar.gz → 0.2.4__tar.gz

{hessboost-0.2.2 → hessboost-0.2.4}/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "hessboost"
-version = "0.2.2"
+version = "0.2.4"
 edition = "2024"
 rust-version = "1.93"
 license = "Apache-2.0"
@@ -23,6 +23,12 @@ include = [
 ]
 readme = "README.md"
 
+# docs.rs builds on Linux, where the `wgpu` backend compiles (its drivers
+# load at run time) and `metal` cannot (no Apple SDK), so the wgpu module's
+# real docs render there and Metal's stand-in does.
+[package.metadata.docs.rs]
+features = ["wgpu"]
+
 [dependencies]
 rayon = "1.12"
 serde = { version = "1.0", features = ["derive"] }
@@ -32,11 +38,31 @@ serde = { version = "1.0", features = ["derive"] }
 serde_json = { version = "1.0", features = ["float_roundtrip"] }
 zstd = { version = "0.14", default-features = false }
 
+# `wgpu`'s own backends: Vulkan (Linux, Windows, Android), Metal (macOS),
+# DirectX 12 (Windows); GL is left out (no 64-bit integers). `bytemuck`
+# casts plain-data slices to the bytes the GPU buffers take, so the backend
+# has no `unsafe`; `parking_lot` is the lock the backend's pools use. Both
+# are already in wgpu's dependency tree.
+wgpu = { version = "30", default-features = false, features = [
+    "std",
+    "parking_lot",
+    "vulkan",
+    "metal",
+    "dx12",
+    "wgsl",
+], optional = true }
+bytemuck = { version = "1.22", features = ["derive"], optional = true }
+parking_lot = { version = "0.12", optional = true }
+
 [features]
 # Native Metal acceleration for histogram training and prediction on
 # macOS (Apple Silicon and Intel Macs with a Metal GPU). Opt-in and
 # off by default; see `src/backend/`.
 metal = ["dep:objc2", "dep:objc2-metal", "dep:objc2-foundation"]
+# Portable GPU acceleration through wgpu (Vulkan, Metal, DirectX 12) for
+# histogram training (`device = wgpu`) and prediction
+# (`BoostedModel::to_wgpu`). Opt-in and off by default; see `src/backend/`.
+wgpu = ["dep:wgpu", "dep:bytemuck", "dep:parking_lot"]
 
 [target.'cfg(target_os = "macos")'.dependencies]
 objc2 = { version = "0.6", default-features = false, features = ["std"], optional = true }

{hessboost-0.2.2 → hessboost-0.2.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: hessboost
-Version: 0.2.2
+Version: 0.2.4
 Classifier: Development Status :: 4 - Beta
 Classifier: Intended Audience :: Developers
 Classifier: Intended Audience :: Science/Research
@@ -51,6 +51,8 @@ UBJSON.
   model at any `nthread`.
 - **Typed** (`py.typed`, complete type information), with the GIL released
   while training and predicting, and free-threaded CPython supported.
+  Native work runs on the extension's own thread pool, which a child
+  forked with `os.fork()` rebuilds instead of hanging.
 - **Modern modeling (opt-in).** Conformal prediction intervals,
   confidence intervals for the regression function (Boulevard boosting),
   distributional boosting (a predictive distribution per row), LightGBM/CatBoost
@@ -67,10 +69,11 @@ Extras: `hessboost[pandas]`, `hessboost[polars]`, and
 `hessboost[scikit-learn]`. numpy is the only required dependency.
 
 Prebuilt wheels are published for Linux x86_64 and aarch64 (glibc
-manylinux and musl/Alpine musllinux), macOS arm64 (with Metal support:
-`device="metal"` training histograms and `Booster.to_gpu()` GPU batch
-prediction), and Windows x86_64: one `abi3` wheel per platform for
-CPython 3.11 and newer, plus a wheel for free-threaded CPython 3.14t.
+manylinux and musl/Alpine musllinux), macOS arm64, and Windows x86_64: one
+`abi3` wheel per platform for CPython 3.11 and newer, plus a wheel for
+free-threaded CPython 3.14t. Every wheel trains and predicts on a GPU
+through wgpu (Vulkan, Metal, DirectX 12), and the macOS wheels also through
+native Metal (see [GPU training and prediction](#gpu-training-and-prediction)).
 Elsewhere the installer builds from the source distribution, which needs
 Rust 1.93 or newer and a C compiler (for libzstd).
 
@@ -114,16 +117,69 @@ leaves = booster.predict(X[800:], pred_leaf=True)  # (rows, trees) int32, all tr
 print(booster.best_iteration, booster.get_score(importance_type="gain"))
 ```
 
+Values and margins of a plain 2-D numpy array are predicted from the array
+itself, without building a `DMatrix` (bit-identical to the `DMatrix` path).
+For serving one row at a time, `predict_row` skips the matrix altogether;
+`out=` reuses a result array, and `transform_margin(s)` applies the
+objective's transform to margins computed elsewhere (bit for bit what
+`predict` reports):
+
+```python
+row = booster.predict_row(X[0])  # == booster.predict(X[:1])[0], 1-D float32
+out = np.empty(1, dtype=np.float32)
+booster.predict_row(X[1], output_margin=True, out=out)  # writes into out
+probability = booster.transform_margin(float(out[0]))
+probabilities = booster.transform_margins(margins)  # == booster.predict(X[800:])
+```
+
+`booster.model_info()` returns the model's structure as numpy arrays, enough
+to walk the trees and recompute the margins without parsing a model file: a
+`ModelInfo` with the layout, base margins, per-tree weights and outputs, and
+the `gblinear` and model-shrinkage records, and per tree a `TreeInfo` of node
+arrays (children, split features, thresholds, categories, leaf values, covers,
+gains, linear leaves).
+
 ### Input data
 
 `DMatrix` takes numpy arrays of any numeric dtype and memory layout (a
 C-contiguous `float32` array is used without a copy), pandas and polars
-DataFrames, scipy sparse matrices, and anything `numpy.asarray` accepts.
-NaN, and a frame's null, is missing (or pass `missing=`); ranking data
-takes `group=` sizes or `qid=`, and `survival:aft` takes
-`label_lower_bound=`/`label_upper_bound=`.
+DataFrames, polars LazyFrames, scipy sparse matrices, and anything
+`numpy.asarray` accepts. NaN, and a frame's null, is missing (or pass
+`missing=`); ranking data takes `group=` sizes or `qid=`, and
+`survival:aft` takes `label_lower_bound=`/`label_upper_bound=`.
 `Booster.predict` accepts the same inputs directly.
 
+With a frame, the per-row metadata can name its columns instead of
+arriving as separate arrays: `label`, `weight`, `base_margin`, `qid`,
+`label_lower_bound` and `label_upper_bound` (`label` and `base_margin`
+also a list of names, for a label matrix or per-output margins). The named
+columns leave the features:
+
+```python
+import polars as pl
+
+dtrain = hessboost.DMatrix(
+    pl.scan_parquet("train.parquet").filter(pl.col("split") == "train").drop("split"),
+    label="price",
+    weight="exposure",
+)
+```
+
+A `LazyFrame` is collected once, by polars' default engine, which since
+polars 2.0 is the streaming engine (spilling to disk when the frame
+outgrows memory). Because that engine keeps no row order after a join or
+`group_by`, a `LazyFrame` takes labels and the other per-row metadata by
+column name only, so that they come out of the same `collect` as the
+features; an array alongside it (or `group` sizes) is refused. Predictions
+on a `LazyFrame` align with its collected rows, so to attach them to a
+frame, collect it yourself and predict on the DataFrame. A polars frame
+converts through a `select` of `Float32` expressions that polars evaluates
+in parallel, in row blocks of 64 MiB of output; `Decimal` and all-null
+columns are numeric, and a column of another dtype (`Datetime`, `String`,
+polars 2.0's `Extension` for Arrow extension types the frame was read with)
+is refused with the conversion that would make it numeric. polars 1.x and
+2.x are both supported (`polars>=1.0`).
+
 ### Training controls
 
 `train` supports XGBoost's everyday arguments: `evals`, `evals_result`,
@@ -134,7 +190,18 @@ and `callbacks`. Ctrl-C stops training at the end of the current round
 and raises `KeyboardInterrupt`. `hessboost.cv` cross-validates over
 shuffled folds, explicit folds, or a scikit-learn splitter; on ranking data
 each fold must hold whole query groups (e.g. `GroupKFold` over the query
-ids). `hessboost.train_with_budget(params, dtrain, budget)` trains with one
+ids). `cv(xgb_model=...)` continues a model in every fold, and
+`cv(refit=True)` also retrains on every row for the chosen round count (the
+best one under early stopping), returning a `CvRefit` with the `history`
+and that `booster`:
+
+```python
+refit = hessboost.cv(params, dtrain, 500, early_stopping_rounds=20, refit=True)
+refit.history["test-rmse-mean"], refit.num_boost_round
+predictions = refit.booster.predict(X_test)
+```
+
+`hessboost.train_with_budget(params, dtrain, budget)` trains with one
 fitting budget in place of `eta`, tree limits, and a round count
 (PerpetualBooster's algorithm).
 
@@ -356,6 +423,13 @@ A Boulevard EBM (`{"booster": "ebm", "ebm_boulevard": True}`) additionally
 supports confidence bands on its shapes via
 `hessboost.inference.EbmInference.term_bands`.
 
+EBMs take eval sets like any booster: `train(..., evals=[(dvalid, "valid")],
+early_stopping_rounds=5)` records the history and `best_score` (classic
+EBMs stop early on it), and `cv` (with `refit=True` too) cross-validates
+them. `cv` refuses `ebm_early_stopping_rounds`, which would stop each fold's
+stages at different rounds; its `early_stopping_rounds` stops on the fold
+means instead.
+
 ### Boulevard inference
 
 Boulevard boosting (`{"booster": "boulevard"}`) samples trees with dropout
@@ -435,20 +509,39 @@ filled = forest.impute(X_with_nans, n_imputations=5)  # (5, rows, columns)
 `forest_diffusion()`; its `training` mappings are XGBoost parameters, and
 models save and load like `DiffusionModel`s.
 
-### GPU prediction
+### GPU training and prediction
 
-On macOS, `Booster.to_gpu()` lays the model out for batch prediction on
-Metal: the forest uploads once, and each call predicts bit-identically to
-`Booster.predict` (values or raw margins), faster from roughly a few
-thousand row-trees upward:
+Two GPU backends reproduce the CPU's results bit for bit: native Metal
+(`"metal"`, macOS) and wgpu (`"wgpu"`: Vulkan on Linux and Windows, Metal
+on macOS, DirectX 12 on Windows). `Booster.to_gpu()` lays a model out for
+batch prediction (on Metal on macOS and on wgpu elsewhere;
+`to_gpu("wgpu")` asks for wgpu): the forest uploads once, and each call
+predicts bit-identically to `Booster.predict` (values or raw margins).
+`GpuModel.available()` says whether the device can predict here:
 
 ```python
 from hessboost import GpuModel
 
-gpu = booster.to_gpu()
-probabilities = gpu.predict(X_test)
+if GpuModel.available():
+    gpu = booster.to_gpu()
+    probabilities = gpu.predict(X_test)
 ```
 
+Training with `device="metal"` or `device="wgpu"` builds the larger nodes'
+histograms on the GPU and gives the CPU's model bit for bit:
+
+```python
+booster = hessboost.train({"device": "wgpu", "max_depth": 6}, dtrain, 100)
+```
+
+Metal prediction is faster than the CPU from roughly a few thousand
+row-trees upward. wgpu needs an adapter with 64-bit shader integers:
+desktop Vulkan drivers, Apple GPUs, or DirectX 12 with
+`dxcompiler.dll` where Windows finds DLLs. It uses a software adapter such
+as Mesa's lavapipe (correct, but slower than the CPU) only when there is no
+other; `GpuModel.device_name("wgpu")` names the adapter it picked, and the
+`WGPU_ADAPTER_NAME` environment variable picks one by name.
+
 ### Compact models
 
 `Booster.to_compact()` packs the trees default prediction uses into
@@ -487,7 +580,10 @@ result = hessboost.cv({"max_depth": 4}, dtrain, 100, folds=splits)
 with CatBoost-style ordered target means: a training row's encoding never
 sees its own label. `label=` supplies the target for a multi-target matrix
 or a class's 0/1 indicator. In `cv`, `target_stats=` fits the encoder on
-each fold's training rows only, so no held-out label reaches an encoding:
+each fold's training rows only, so no held-out label reaches an encoding
+(`target_stats_label=` supplies its target, split with the folds); with
+`refit=True`, `CvRefit.target_encoder` is the encoder fitted on every row,
+which encodes new data for the refit booster:
 
 ```python
 from hessboost.target_stats import OrderedTargetEncoder
@@ -497,8 +593,15 @@ dtrain_encoded, stats = encoder.fit_transform(dtrain, ["city"])
 booster = hessboost.train({"max_depth": 4}, dtrain_encoded, 100)
 predictions = booster.predict(stats.transform(X_test))
 result = hessboost.cv({"max_depth": 4}, dtrain, 100, target_stats=["city"], target_encoder=encoder)
+refit = hessboost.cv({"max_depth": 4}, dtrain, 100, target_stats=["city"], refit=True)
+predictions = refit.booster.predict(refit.target_encoder.transform(X_test))
 ```
 
+`stats.save(path)` / `FittedTargetEncoder.load(path)` (and
+`to_bytes`/`from_bytes`) store the Rust crate's serde JSON, which keeps
+column indices and category codes but not feature names or frame
+categories; pickle the statistics to keep them.
+
 ### Extra training options
 
 Every hessboost training option (`path_smooth`, `extra_trees`,
@@ -531,21 +634,26 @@ LightGBM's `rank_xendcg` stream.
   `objective`, which `params` must then not set, and `num_class` is the
   custom objective's output count (XGBoost's custom-softmax convention;
   default: one per label column).
+  Unlike XGBoost, which accepts both, hessboost refuses `objective` alongside
+  `obj`: drop `objective` from ported `params`; the callback needs no change.
 - `cv` returns a dict of numpy arrays (`test-<metric>-mean`/`-std`) with
-  held-out metrics only; there is no `stratified` or `as_pandas`.
+  held-out metrics only (or, with `refit=True`, a `CvRefit` holding it and
+  the retrained booster); there is no `stratified` or `as_pandas`.
 - `predict` defaults to the iterations through `best_iteration` (XGBoost's
   scikit-learn behavior); pass `iteration_range=(0, 0)` for all. SHAP and
   leaf ranges start at iteration 0; `pred_leaf` defaults to every iteration
   instead, and returns `int32`.
-- On macOS, `Booster.to_gpu()` lays the model out for GPU batch prediction
-  on Metal (`GpuModel.predict`, bit-identical to `Booster.predict`; see
-  whether a device exists with `GpuModel.available()`).
+- GPUs: `device="metal"` (macOS) or `device="wgpu"` trains on one instead of
+  `device="cuda"`, and `Booster.to_gpu()` lays the model out for GPU batch
+  prediction (`GpuModel.predict`, bit-identical to `Booster.predict`)
+  instead of predicting on the training device.
 - Model files do not store feature names or categories (pickles do).
 - Not available: `DMatrix` from files or `QuantileDMatrix`, `inplace_predict`
-  (`predict` takes arrays directly), `Booster.get_dump`/`trees_to_dataframe`
-  /`dump_model`, attributes (`set_attr`), plotting, distributed (Dask/Spark)
-  and CUDA training, `approx_contribs`, and `strict_shape`. The macOS wheels
-  support `device="metal"` (GPU histograms while training).
+  (`predict` takes arrays directly, and `predict_row` single rows),
+  `Booster.get_dump`/`trees_to_dataframe`/`dump_model` (`model_info()`
+  returns the trees as arrays), attributes (`set_attr`), plotting,
+  distributed (Dask/Spark) and CUDA training, `approx_contribs`, and
+  `strict_shape`.
 
 ## Development
 

{hessboost-0.2.2 → hessboost-0.2.4}/README.md

@@ -128,15 +128,18 @@ runnable programs live in [`examples/`](examples)
 | `online_update` | adding and deleting training rows in place, and exact unlearning |
 | `pfn_boost` | boosting from a pretrained model's logits |
 | `metal` | CPU vs GPU prediction (macOS, `--features metal`) |
+| `wgpu` | GPU training and prediction through wgpu, bit-identical to the CPU (`--features wgpu`; runs on Mesa's lavapipe without a GPU) |
 
 ## Python
 
 [`python/`](python) holds the Python package (`pip install hessboost`):
 `DMatrix`, `train`, `cv`, and `Booster` (taking XGBoost's parameter names),
-scikit-learn estimators, pandas and polars categorical input, and the conformal,
-distributional, tree-diffusion, ForestFlow, in-place update, ordered target
-statistics, budget training, and compact model extras
-(on macOS, `Booster.to_gpu()` batch-predicts on the Metal GPU):
+scikit-learn estimators, pandas and polars (1.x and 2.x; DataFrames and
+LazyFrames, with labels taken from the frame's columns by name) categorical
+input, and the conformal, distributional, tree-diffusion, ForestFlow,
+in-place update, ordered target statistics, budget training, and compact
+model extras, with GPU training and batch prediction (`device="wgpu"` on
+every platform, `device="metal"` on macOS; `Booster.to_gpu()`):
 
 ```python
 import hessboost
@@ -157,8 +160,12 @@ See [`python/README.md`](python/README.md).
   binary/multiclass classification, ranking (LambdaMART, XE-NDCG), count, and survival (Cox, AFT),
   plus typed objective and metric APIs and custom loss hooks.
 - **Validation & workflow:** Cross-validation (including purged and forward time-series folds,
-  whole-query ranking folds, and per-fold target statistics), early stopping, feature importance,
+  whole-query ranking folds, per-fold target statistics, continuing a model, and refitting on all
+  rows at the chosen round count), early stopping, feature importance,
   SHAP values and interactions, model slicing, and iteration ranges.
+- **Prediction & inspection:** Batch, in-place (borrowed rows), and allocation-free single-row
+  prediction, bit-identical in every batch; read-only model metadata (tree weights and nodes,
+  `gblinear` weights, model-shrinkage records).
 - **Interchange:** Native binary and JSON formats, XGBoost JSON/UBJSON import/export, LightGBM model import, and models embedded in the binary at compile time.
 - **Modern modeling (opt-in):**
   - [Conformal intervals](https://docs.rs/hessboost/latest/hessboost/conformal/): Finite-sample coverage guarantees.
@@ -171,6 +178,7 @@ See [`python/README.md`](python/README.md).
   - [Compact models](https://docs.rs/hessboost/latest/hessboost/model/compact/): Bit-packed model format with identical margins.
   - [Budget training](https://docs.rs/hessboost/latest/hessboost/training/budget/): Training controlled by one budget value, based on PerpetualBooster.
   - [Metal GPU](https://docs.rs/hessboost/latest/hessboost/backend/metal/): Apple Silicon GPU prediction and training (`--features metal`).
+  - [wgpu GPU](https://docs.rs/hessboost/latest/hessboost/backend/wgpu/): Vulkan, Metal, and DirectX 12 GPU prediction and training through wgpu (`--features wgpu`), bit-identical to the CPU.
 
 ## Caveats
 
@@ -181,7 +189,7 @@ See [`python/README.md`](python/README.md).
 
 - Distributed and external-memory training.
 - CLI and C bindings.
-- GPU training outside macOS (a `wgpu` backend is planned).
+- CUDA. GPU training runs through Metal (macOS) or wgpu (Vulkan, Metal, DirectX 12).
 - A few XGBoost options exist at one setting only, and a few metrics are
   missing; the [API docs](https://docs.rs/hessboost/latest/hessboost/#not-implemented)
   list them.

{hessboost-0.2.2 → hessboost-0.2.4}/benches/training.rs

@@ -8,8 +8,8 @@ use criterion::{
     BenchmarkGroup, BenchmarkId, Criterion, Throughput, criterion_group, criterion_main,
 };
 use hessboost::config::{
-    BoosterKind, Dart, ExtraTrees, GrowPolicy, LinearTree, Monotone, MultiStrategy, QuantizedGrad,
-    TrainingParamsBuilder,
+    BoosterKind, Dart, Ebm, ExtraTrees, GrowPolicy, LinearTree, Monotone, MultiStrategy,
+    QuantizedGrad, TrainingParamsBuilder,
 };
 use hessboost::data::FeatureType;
 use hessboost::internals::{
@@ -1208,36 +1208,68 @@ fn bench_data_prep(c: &mut Criterion) {
     group.finish();
 }
 
-/// Metal GPU benches (`cargo bench --features metal` on a Mac with a Metal
-/// device): histogram construction, end-to-end training, and batch
-/// prediction, each against its CPU counterpart on identical data. The GPU
-/// results are bit-identical to the single-threaded CPU's, so the benches
-/// compare speed only.
-#[cfg(all(target_os = "macos", feature = "metal"))]
-fn bench_metal(c: &mut Criterion) {
-    use hessboost::backend::metal;
-    use hessboost::backend::metal::MetalHistBackend;
-    use hessboost::config::Device;
-
-    if let Some(reason) = metal::unavailable_reason() {
-        eprintln!("skipping metal benches: {reason}");
-        return;
+/// FAST pair ranking (`booster = ebm`): one round of 20 main effects, then
+/// FAST over all 190 feature pairs (`interactions_1`) or no pairs at all
+/// (`interactions_0`); the difference is FAST.
+fn bench_ebm_fast(c: &mut Criterion) {
+    let mut group = c.benchmark_group("ebm_fast_x20_1round");
+    group.sample_size(10);
+    for (label, rows) in [("10k", 10_000), ("100k", 100_000)] {
+        let data = make_data(rows, 20);
+        for interactions in [0, 1] {
+            let params = TrainingParams::builder()
+                .booster(BoosterKind::Ebm(
+                    Ebm::builder().interactions(interactions).build().unwrap(),
+                ))
+                .grow_policy(GrowPolicy::LossGuide)
+                .max_leaves(3)
+                .build()
+                .unwrap();
+            group.bench_function(format!("{label}_interactions_{interactions}"), |b| {
+                b.iter(|| black_box(train(&params, &data, 1).unwrap()));
+            });
+        }
     }
+    group.finish();
+}
+
+/// One GPU backend's entry points for [`bench_gpu`].
+#[cfg(any(all(target_os = "macos", feature = "metal"), feature = "wgpu"))]
+struct GpuBench<H, G> {
+    /// Group-name prefix and bench id (`metal`, `wgpu`).
+    name: &'static str,
+    /// The `device` that trains on the backend.
+    device: hessboost::config::Device,
+    /// The backend's histogram builder for an index.
+    hist_backend: fn(&GHistIndex) -> Result<H>,
+    /// The model laid out for the backend's prediction.
+    to_gpu: fn(&BoostedModel) -> Result<G>,
+    /// The GPU model's margins of every iteration.
+    predict_margin: fn(&G, &DMatrix) -> Result<hessboost::model::Predictions>,
+}
+
+/// A GPU backend's benches: histogram construction, end-to-end training,
+/// and batch prediction, each against its CPU counterpart on identical
+/// data. The GPU results are bit-identical to the single-threaded CPU's, so
+/// the benches compare speed only.
+#[cfg(any(all(target_os = "macos", feature = "metal"), feature = "wgpu"))]
+fn bench_gpu<H: HistogramBackend, G>(c: &mut Criterion, gpu: &GpuBench<H, G>) {
+    let name = gpu.name;
     // Histogram construction at the sizes where the GPU pays off.
     {
-        let mut group = c.benchmark_group("metal_histogram_build");
+        let mut group = c.benchmark_group(format!("{name}_histogram_build"));
         group.sample_size(10);
         for &n in &[100_000usize, 1_000_000] {
             let (ghist, gpair, rows) = histogram_case(&make_data(n, 30));
-            let gpu = MetalHistBackend::new(&ghist).unwrap();
+            let backend = (gpu.hist_backend)(&ghist).unwrap();
             let mut cpu_out = zeroed(ghist.total_bins());
             let mut gpu_out = zeroed(ghist.total_bins());
             group.throughput(Throughput::Elements(n as u64));
             group.bench_with_input(BenchmarkId::new("cpu", n), &n, |b, _| {
                 b.iter(|| CpuBackend.build(&ghist, &rows, &gpair, &mut cpu_out));
             });
-            group.bench_with_input(BenchmarkId::new("metal", n), &n, |b, _| {
-                b.iter(|| gpu.build(&ghist, &rows, &gpair, &mut gpu_out));
+            group.bench_with_input(BenchmarkId::new(name, n), &n, |b, _| {
+                b.iter(|| backend.build(&ghist, &rows, &gpair, &mut gpu_out));
             });
         }
         group.finish();
@@ -1245,9 +1277,9 @@ fn bench_metal(c: &mut Criterion) {
     // End-to-end training: identical parameters, CPU against GPU histograms.
     {
         let data = make_data(200_000, 30);
-        let mut group = c.benchmark_group("metal_train_200k_x30_50rounds_depth8");
+        let mut group = c.benchmark_group(format!("{name}_train_200k_x30_50rounds_depth8"));
         group.sample_size(10);
-        for (name, device) in [("cpu", Device::Cpu), ("metal", Device::Metal)] {
+        for (id, device) in [("cpu", hessboost::config::Device::Cpu), (name, gpu.device)] {
             let params = TrainingParams::builder()
                 .objective(Objective::SquaredError(RegLoss::default()))
                 .tree_method(TreeMethod::Hist)
@@ -1256,7 +1288,7 @@ fn bench_metal(c: &mut Criterion) {
                 .device(device)
                 .build()
                 .unwrap();
-            group.bench_function(name, |b| {
+            group.bench_function(id, |b| {
                 b.iter(|| black_box(train(&params, &data, 50).unwrap()));
             });
         }
@@ -1266,9 +1298,9 @@ fn bench_metal(c: &mut Criterion) {
     {
         let model_data = make_data(100_000, 30);
         let model = trained_model(&model_data, 100);
-        let gpu = model.to_gpu().unwrap();
+        let gpu_model = (gpu.to_gpu)(&model).unwrap();
         let data = make_data(500_000, 30);
-        let mut group = c.benchmark_group("metal_predict_500k_x30_100trees_depth6");
+        let mut group = c.benchmark_group(format!("{name}_predict_500k_x30_100trees_depth6"));
         group.sample_size(10);
         group.throughput(Throughput::Elements(data.n_rows() as u64));
         group.bench_function("cpu", |b| {
@@ -1278,23 +1310,73 @@ fn bench_metal(c: &mut Criterion) {
                     .unwrap()
             });
         });
-        group.bench_function("metal", |b| {
-            b.iter(|| {
-                gpu.predict_margin(black_box(&data), Iterations::Best)
-                    .unwrap()
-            });
+        group.bench_function(name, |b| {
+            b.iter(|| (gpu.predict_margin)(&gpu_model, black_box(&data)).unwrap());
         });
         group.finish();
     }
 }
 
+/// Metal GPU benches (`cargo bench --features metal` on a Mac with a Metal
+/// device); see [`bench_gpu`].
 #[cfg(all(target_os = "macos", feature = "metal"))]
-fn bench_metal_registered(c: &mut Criterion) {
-    bench_metal(c);
+fn bench_metal(c: &mut Criterion) {
+    use hessboost::backend::metal;
+
+    if let Some(reason) = metal::unavailable_reason() {
+        eprintln!("skipping metal benches: {reason}");
+        return;
+    }
+    bench_gpu(
+        c,
+        &GpuBench {
+            name: "metal",
+            device: hessboost::config::Device::Metal,
+            hist_backend: metal::MetalHistBackend::new,
+            to_gpu: BoostedModel::to_gpu,
+            predict_margin: |gpu, data| gpu.predict_margin(data, Iterations::Best),
+        },
+    );
 }
 
 #[cfg(not(all(target_os = "macos", feature = "metal")))]
-fn bench_metal_registered(_c: &mut Criterion) {}
+fn bench_metal(_c: &mut Criterion) {}
+
+/// wgpu GPU benches (`cargo bench --features wgpu` on a machine with a
+/// Vulkan, Metal, or DirectX 12 adapter); see [`bench_gpu`]. On a software
+/// adapter (lavapipe, WARP) the GPU side measures the CPU emulating one,
+/// which says nothing about a GPU.
+#[cfg(feature = "wgpu")]
+fn bench_wgpu(c: &mut Criterion) {
+    use hessboost::backend::wgpu;
+
+    if let Some(reason) = wgpu::unavailable_reason() {
+        eprintln!("skipping wgpu benches: {reason}");
+        return;
+    }
+    eprintln!(
+        "wgpu adapter: {}{}",
+        wgpu::device_name().unwrap_or_default(),
+        if wgpu::is_software_adapter() == Some(true) {
+            " (software renderer)"
+        } else {
+            ""
+        }
+    );
+    bench_gpu(
+        c,
+        &GpuBench {
+            name: "wgpu",
+            device: hessboost::config::Device::Wgpu,
+            hist_backend: wgpu::WgpuHistBackend::new,
+            to_gpu: BoostedModel::to_wgpu,
+            predict_margin: |gpu, data| gpu.predict_margin(data, Iterations::Best),
+        },
+    );
+}
+
+#[cfg(not(feature = "wgpu"))]
+fn bench_wgpu(_c: &mut Criterion) {}
 
 criterion_group!(
     benches,
@@ -1315,6 +1397,8 @@ criterion_group!(
     bench_predict_csr,
     bench_model_io,
     bench_data_prep,
-    bench_metal_registered
+    bench_ebm_fast,
+    bench_metal,
+    bench_wgpu
 );
 criterion_main!(benches);