warpgbm 0.1.16__tar.gz → 0.1.18__tar.gz

{warpgbm-0.1.16/warpgbm.egg-info → warpgbm-0.1.18}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: warpgbm
-Version: 0.1.16
+Version: 0.1.18
 Summary: A fast GPU-accelerated Gradient Boosted Decision Tree library with PyTorch + CUDA
 License:                     GNU GENERAL PUBLIC LICENSE
                                Version 3, 29 June 2007
@@ -706,7 +706,7 @@ WarpGBM is a high-performance, GPU-accelerated Gradient Boosted Decision Tree (G
 
 ## Performance Note
 
-In our initial tests on an NVIDIA 3090 (local) and A100 (Google Colab Pro), WarpGBM achieves **14x to 20x faster training times** compared to LightGBM using default configurations. It also consumes **significantly less RAM and CPU**. These early results hint at more thorough benchmarking to come.
+In our initial tests on an NVIDIA 3090 (local) and A100 (Google Colab Pro), WarpGBM achieves **14x to 20x faster training times** compared to LightGBM's CPU version and **2x faster** on the GPU version using default configurations. Speed also outperforms XGBoost and CatBoost on regression problems. It also consumes **significantly less RAM and CPU**. These early results hint at more thorough benchmarking to come.
 
 ---
 

{warpgbm-0.1.16 → warpgbm-0.1.18}/README.md

@@ -18,7 +18,7 @@ WarpGBM is a high-performance, GPU-accelerated Gradient Boosted Decision Tree (G
 
 ## Performance Note
 
-In our initial tests on an NVIDIA 3090 (local) and A100 (Google Colab Pro), WarpGBM achieves **14x to 20x faster training times** compared to LightGBM using default configurations. It also consumes **significantly less RAM and CPU**. These early results hint at more thorough benchmarking to come.
+In our initial tests on an NVIDIA 3090 (local) and A100 (Google Colab Pro), WarpGBM achieves **14x to 20x faster training times** compared to LightGBM's CPU version and **2x faster** on the GPU version using default configurations. Speed also outperforms XGBoost and CatBoost on regression problems. It also consumes **significantly less RAM and CPU**. These early results hint at more thorough benchmarking to come.
 
 ---
 

{warpgbm-0.1.16 → warpgbm-0.1.18}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "warpgbm"
-version = "0.1.16"
+version = "0.1.18"
 description = "A fast GPU-accelerated Gradient Boosted Decision Tree library with PyTorch + CUDA"
 readme = "README.md"
 requires-python = ">=3.8"

{warpgbm-0.1.16 → warpgbm-0.1.18}/setup.py

@@ -22,6 +22,7 @@ def get_extensions():
                 sources=[
                     "warpgbm/cuda/histogram_kernel.cu",
                     "warpgbm/cuda/best_split_kernel.cu",
+                    "warpgbm/cuda/binner.cu",
                     "warpgbm/cuda/node_kernel.cpp",
                 ]
             )

warpgbm-0.1.18/tests/test_fit_predict_corr.py

@@ -0,0 +1,46 @@
+import numpy as np
+from warpgbm import WarpGBM
+from sklearn.datasets import make_regression
+
+import numpy as np
+import time
+from warpgbm import WarpGBM
+from sklearn.datasets import make_regression
+
+def test_fit_predictpytee_correlation():
+    np.random.seed(42)
+    N = 100_000
+    F = 1000
+    X, y = make_regression(n_samples=N, n_features=F, noise=0.1, random_state=42)
+    era = np.zeros(N, dtype=np.int32)
+    corrs = []
+
+    for hist_type in ['hist1', 'hist2', 'hist3']:
+        print(f"\nTesting histogram method: {hist_type}")
+
+        model = WarpGBM(
+            max_depth=10,
+            num_bins=10,
+            n_estimators=10,
+            learning_rate=1,
+            verbosity=False,
+            histogram_computer=hist_type,
+            threads_per_block=128,
+            rows_per_thread=4
+        )
+
+        start_fit = time.time()
+        model.fit(X, y, era_id=era)
+        fit_time = time.time() - start_fit
+        print(f"  Fit time:     {fit_time:.3f} seconds")
+
+        start_pred = time.time()
+        preds = model.predict(X)
+        pred_time = time.time() - start_pred
+        print(f"  Predict time: {pred_time:.3f} seconds")
+
+        corr = np.corrcoef(preds, y)[0, 1]
+        print(f"  Correlation:  {corr:.4f}")
+        corrs.append(corr)
+
+    assert (np.array(corrs) > 0.95).all(), f"In-sample correlation too low: {corrs}"

warpgbm-0.1.18/version.txt

@@ -0,0 +1 @@
+0.1.18

{warpgbm-0.1.16 → warpgbm-0.1.18}/warpgbm/core.py

@@ -12,40 +12,6 @@ histogram_kernels = {
     'hist3': node_kernel.compute_histogram3
 }
 
-@torch.jit.script
-def jit_find_best_split(
-    G: Tensor, H: Tensor, 
-    lambda_l2: float, 
-    lambda_l1: float,  # unused placeholder for now
-    min_split_gain: float, 
-    min_child_weight: float
-) -> Tuple[int, int]:
-    F, B = G.size()
-    Bm1 = B - 1
-    eps = 0
-
-    GH = torch.stack([G, H], dim=0).cumsum(dim=2)  # [2, F, B]
-    GL, HL_raw = GH[0, :, :-1], GH[1, :, :-1]      # [F, B-1]
-    GP, HP = GH[0, :, -1:], GH[1, :, -1:]          # [F, 1]
-    H_R_raw = HP - HL_raw
-
-    # Validity mask using raw child hessians
-    valid = (HL_raw >= min_child_weight) & (H_R_raw >= min_child_weight)
-
-    # Closed-form gain
-    HL, HP = HL_raw + lambda_l2, HP + lambda_l2
-    num = (HP * GL - HL * GP).pow(2)
-    denom = HP * HL * (HP - HL) + eps
-    gain = torch.where(valid & (num / denom >= min_split_gain), num / denom, torch.full_like(num, -float("inf")))
-
-    gain_flat = gain.view(-1)
-    best_idx = torch.argmax(gain_flat)
-
-    if gain_flat[best_idx].item() == float('-inf'):
-        return -1, -1
-
-    return best_idx // Bm1, best_idx % Bm1
-
 class WarpGBM(BaseEstimator, RegressorMixin):
     def __init__(
         self,
@@ -80,12 +46,8 @@ class WarpGBM(BaseEstimator, RegressorMixin):
         self.Y_gpu = None
         self.num_features = None
         self.num_samples = None
-        self.out_feature = torch.zeros(1, device=self.device, dtype=torch.int32)
-        self.out_bin = torch.zeros(1, device=self.device, dtype=torch.int32)
         self.min_child_weight = min_child_weight
         self.min_split_gain = min_split_gain
-        self.best_gain = torch.tensor([-float('inf')], dtype=torch.float32, device=self.device)
-        self.best_feature = torch.tensor([-1], dtype=torch.int32, device=self.device)
         self.best_bin = torch.tensor([-1], dtype=torch.int32, device=self.device)
         self.compute_histogram = histogram_kernels[histogram_computer]
         self.threads_per_block = threads_per_block
@@ -102,45 +64,49 @@ class WarpGBM(BaseEstimator, RegressorMixin):
         self.root_node_indices = torch.arange(self.num_samples, device=self.device)
         self.base_prediction = self.Y_gpu.mean().item()
         self.gradients += self.base_prediction
-        self.split_gains = torch.zeros((self.num_features, self.num_bins - 1), device=self.device)
-        self.forest = self.grow_forest()
+        self.best_gains = torch.zeros(self.num_features, device=self.device)
+        self.best_bins = torch.zeros(self.num_features, device=self.device, dtype=torch.int32)
+        with torch.no_grad():
+            self.forest = self.grow_forest()
         return self
-
-    def compute_quantile_bins(self, X, num_bins):
-        quantiles = torch.linspace(0, 1, num_bins + 1)[1:-1]  # exclude 0% and 100%
-        bin_edges = torch.quantile(X, quantiles, dim=0)       # shape: [B-1, F]
-        return bin_edges.T  # shape: [F, B-1]
     
     def preprocess_gpu_data(self, X_np, Y_np, era_id_np):
-        self.num_samples, self.num_features = X_np.shape
-        Y_gpu = torch.from_numpy(Y_np).type(torch.float32).to(self.device)
-        era_id_gpu = torch.from_numpy(era_id_np).type(torch.int32).to(self.device)
-        is_integer_type = np.issubdtype(X_np.dtype, np.integer)
-        if is_integer_type:
-            max_vals = X_np.max(axis=0)
-            if np.all(max_vals < self.num_bins):
-                print("Detected pre-binned integer input — skipping quantile binning.")
-                bin_indices = torch.from_numpy(X_np).to(self.device).contiguous().to(torch.int8)
-    
-                # We'll store None or an empty tensor in self.bin_edges
-                # to indicate that we skip binning at predict-time
-                bin_edges = torch.arange(1, self.num_bins, dtype=torch.float32).repeat(self.num_features, 1)
-                bin_edges = bin_edges.to(self.device)
-                unique_eras, era_indices = torch.unique(era_id_gpu, return_inverse=True)
-                return bin_indices, era_indices, bin_edges, unique_eras, Y_gpu
-            else:
-                print("Integer input detected, but values exceed num_bins — falling back to quantile binning.")
-    
-        print("Performing quantile binning on CPU...")
-        X_cpu = torch.from_numpy(X_np).type(torch.float32)  # CPU tensor
-        bin_edges_cpu = self.compute_quantile_bins(X_cpu, self.num_bins).type(torch.float32).contiguous()
-        bin_indices_cpu = torch.empty((self.num_samples, self.num_features), dtype=torch.int8)
-        for f in range(self.num_features):
-            bin_indices_cpu[:, f] = torch.bucketize(X_cpu[:, f], bin_edges_cpu[f], right=False).type(torch.int8)
-        bin_indices = bin_indices_cpu.to(self.device).contiguous()
-        bin_edges = bin_edges_cpu.to(self.device)
-        unique_eras, era_indices = torch.unique(era_id_gpu, return_inverse=True)
-        return bin_indices, era_indices, bin_edges, unique_eras, Y_gpu
+        with torch.no_grad():
+            self.num_samples, self.num_features = X_np.shape
+            Y_gpu = torch.from_numpy(Y_np).type(torch.float32).to(self.device)
+            era_id_gpu = torch.from_numpy(era_id_np).type(torch.int32).to(self.device)
+            is_integer_type = np.issubdtype(X_np.dtype, np.integer)
+            if is_integer_type:
+                max_vals = X_np.max(axis=0)
+                if np.all(max_vals < self.num_bins):
+                    print("Detected pre-binned integer input — skipping quantile binning.")
+                    bin_indices = torch.from_numpy(X_np).to(self.device).contiguous().to(torch.int8)
+        
+                    # We'll store None or an empty tensor in self.bin_edges
+                    # to indicate that we skip binning at predict-time
+                    bin_edges = torch.arange(1, self.num_bins, dtype=torch.float32).repeat(self.num_features, 1)
+                    bin_edges = bin_edges.to(self.device)
+                    unique_eras, era_indices = torch.unique(era_id_gpu, return_inverse=True)
+                    return bin_indices, era_indices, bin_edges, unique_eras, Y_gpu
+                else:
+                    print("Integer input detected, but values exceed num_bins — falling back to quantile binning.")
+        
+            bin_indices = torch.empty((self.num_samples, self.num_features), dtype=torch.int8, device='cuda')
+            bin_edges = torch.empty((self.num_features, self.num_bins - 1), dtype=torch.float32, device='cuda')
+
+            X_np = torch.from_numpy(X_np).to(torch.float32).pin_memory()
+
+            for f in range(self.num_features):
+                X_f = X_np[:, f].to('cuda', non_blocking=True)
+                quantiles = torch.linspace(0, 1, self.num_bins + 1, device='cuda', dtype=X_f.dtype)[1:-1]
+                bin_edges_f = torch.quantile(X_f, quantiles, dim=0).contiguous()  # shape: [B-1] for 1D input
+                bin_indices_f = bin_indices[:, f].contiguous()  # view into output
+                node_kernel.custom_cuda_binner(X_f, bin_edges_f, bin_indices_f)
+                bin_indices[:,f] = bin_indices_f
+                bin_edges[f,:] = bin_edges_f
+
+            unique_eras, era_indices = torch.unique(era_id_gpu, return_inverse=True)
+            return bin_indices, era_indices, bin_edges, unique_eras, Y_gpu
 
     def compute_histograms(self, bin_indices_sub, gradients):
         grad_hist = torch.zeros((self.num_features, self.num_bins), device=self.device, dtype=torch.float32)
@@ -158,15 +124,24 @@ class WarpGBM(BaseEstimator, RegressorMixin):
         return grad_hist, hess_hist
 
     def find_best_split(self, gradient_histogram, hessian_histogram):
-        f,b = jit_find_best_split(
+        node_kernel.compute_split(
             gradient_histogram,
             hessian_histogram,
-            self.L2_reg,
-            self.L1_reg,
             self.min_split_gain,
             self.min_child_weight,
+            self.L2_reg,
+            self.best_gains,
+            self.best_bins,
+            self.threads_per_block
         )
-        return (f, b)
+
+        if torch.all(self.best_bins == -1):
+            return -1, -1  # No valid split found
+
+        f = torch.argmax(self.best_gains).item()
+        b = self.best_bins[f].item()
+
+        return f, b
     
     def grow_tree(self, gradient_histogram, hessian_histogram, node_indices, depth):
         if depth == self.max_depth:
@@ -226,10 +201,10 @@ class WarpGBM(BaseEstimator, RegressorMixin):
                 depth=0
             )
             forest[i] = tree
-            # loss = ((self.Y_gpu - self.gradients) ** 2).mean().item()
-            # self.training_loss.append(loss)
-            # print(f"🌲 Tree {i+1}/{self.n_estimators} - MSE: {loss:.6f}")
-    
+        # loss = ((self.Y_gpu - self.gradients) ** 2).mean().item()
+        # self.training_loss.append(loss)
+        # print(f"🌲 Tree {i+1}/{self.n_estimators} - MSE: {loss:.6f}")
+
         print("Finished training forest.")
         return forest
 
@@ -239,103 +214,104 @@ class WarpGBM(BaseEstimator, RegressorMixin):
         We assume `flatten_forest_to_tensors` has produced self.flat_forest with
         "features", "thresholds", "leaf_values", all shaped [n_trees, max_nodes].
         """
-        # 1) Convert X_np -> bin_indices
-        is_integer_type = np.issubdtype(X_np.dtype, np.integer)
-        if is_integer_type:
-            max_vals = X_np.max(axis=0)
-            if np.all(max_vals < self.num_bins):
-                bin_indices = X_np.astype(np.int8)
+        with torch.no_grad():
+            # 1) Convert X_np -> bin_indices
+            is_integer_type = np.issubdtype(X_np.dtype, np.integer)
+            if is_integer_type:
+                max_vals = X_np.max(axis=0)
+                if np.all(max_vals < self.num_bins):
+                    bin_indices = X_np.astype(np.int8)
+                else:
+                    raise ValueError("Pre-binned integers must be < num_bins")
             else:
-                raise ValueError("Pre-binned integers must be < num_bins")
-        else:
-            X_cpu = torch.from_numpy(X_np).type(torch.float32)
-            bin_indices = torch.empty((X_np.shape[0], X_np.shape[1]), dtype=torch.int8)
-            bin_edges_cpu = self.bin_edges.to('cpu')
-            for f in range(self.num_features):
-                bin_indices[:, f] = torch.bucketize(X_cpu[:, f], bin_edges_cpu[f], right=False).type(torch.int8)
-            bin_indices = bin_indices.numpy()
-
-        # 2) Ensure we have a padded representation
-        self.flat_forest = self.flatten_forest_to_tensors(self.forest)
-
-        features_t   = self.flat_forest["features"]      # [n_trees, max_nodes], int16
-        thresholds_t = self.flat_forest["thresholds"]    # [n_trees, max_nodes], int16
-        values_t     = self.flat_forest["leaf_values"]    # [n_trees, max_nodes], float32
-        max_nodes    = self.flat_forest["max_nodes"]
-
-        n_trees = features_t.shape[0]
-        N       = bin_indices.shape[0]
-        out     = np.zeros(N, dtype=np.float32)
-
-        # 3) Process rows in chunks
-        for start in tqdm(range(0, N, chunk_size)):
-            end = min(start + chunk_size, N)
-            chunk_np  = bin_indices[start:end]  # shape [chunk_size, F]
-            chunk_gpu = torch.from_numpy(chunk_np).to(self.device)  # [chunk_size, F], int8
-
-            # Accumulate raw (unscaled) leaf sums
-            chunk_preds = torch.zeros((end - start,), dtype=torch.float32, device=self.device)
-
-            # node_idx[i] tracks the current node index in the padded tree for row i
-            node_idx = torch.zeros((end - start,), dtype=torch.int32, device=self.device)
-
-            # 'active' is a boolean mask over [0..(end-start-1)], indicating which rows haven't reached a leaf
-            active = torch.ones((end - start,), dtype=torch.bool, device=self.device)
-
-            for t in range(n_trees):
-                # Reset for each tree (each tree is independent)
-                node_idx.fill_(0)
-                active.fill_(True)
-
-                tree_features = features_t[t]     # shape [max_nodes], int16
-                tree_thresh   = thresholds_t[t]    # shape [max_nodes], int16
-                tree_values   = values_t[t]          # shape [max_nodes], float32
-
-                # Up to self.max_depth+1 layers
-                for _level in range(self.max_depth + 1):
-                    active_idx = active.nonzero(as_tuple=True)[0]
-                    if active_idx.numel() == 0:
-                        break  # all rows are done in this tree
-
-                    current_node_idx = node_idx[active_idx]
-                    f    = tree_features[current_node_idx]    # shape [#active], int16
-                    thr  = tree_thresh[current_node_idx]       # shape [#active], int16
-                    vals = tree_values[current_node_idx]       # shape [#active], float32
-
-                    mask_no_node = (f == -2)
-                    mask_leaf    = (f == -1)
-
-                    # If leaf, add leaf value and mark inactive.
-                    if mask_leaf.any():
-                        leaf_rows = active_idx[mask_leaf]
-                        chunk_preds[leaf_rows] += vals[mask_leaf]
-                        active[leaf_rows] = False
-
-                    # If no node, mark inactive.
-                    if mask_no_node.any():
-                        no_node_rows = active_idx[mask_no_node]
-                        active[no_node_rows] = False
-
-                    # For internal nodes, perform bin comparison.
-                    mask_internal = (~mask_leaf & ~mask_no_node)
-                    if mask_internal.any():
-                        internal_rows = active_idx[mask_internal]
-                        act_f   = f[mask_internal].long()
-                        act_thr = thr[mask_internal]
-                        binvals = chunk_gpu[internal_rows, act_f]
-                        go_left = (binvals <= act_thr)
-                        new_left_idx  = current_node_idx[mask_internal] * 2 + 1
-                        new_right_idx = current_node_idx[mask_internal] * 2 + 2
-                        node_idx[internal_rows[go_left]]  = new_left_idx[go_left]
-                        node_idx[internal_rows[~go_left]] = new_right_idx[~go_left]
-                # end per-tree layer loop
-            # end for each tree
-
-            out[start:end] = (
-                self.base_prediction + self.learning_rate * chunk_preds
-            ).cpu().numpy()
-
-        return out
+                X_cpu = torch.from_numpy(X_np).type(torch.float32)
+                bin_indices = torch.empty((X_np.shape[0], X_np.shape[1]), dtype=torch.int8)
+                bin_edges_cpu = self.bin_edges.to('cpu')
+                for f in range(self.num_features):
+                    bin_indices[:, f] = torch.bucketize(X_cpu[:, f], bin_edges_cpu[f], right=False).type(torch.int8)
+                bin_indices = bin_indices.numpy()
+
+            # 2) Ensure we have a padded representation
+            self.flat_forest = self.flatten_forest_to_tensors(self.forest)
+
+            features_t   = self.flat_forest["features"]      # [n_trees, max_nodes], int16
+            thresholds_t = self.flat_forest["thresholds"]    # [n_trees, max_nodes], int16
+            values_t     = self.flat_forest["leaf_values"]    # [n_trees, max_nodes], float32
+            max_nodes    = self.flat_forest["max_nodes"]
+
+            n_trees = features_t.shape[0]
+            N       = bin_indices.shape[0]
+            out     = np.zeros(N, dtype=np.float32)
+
+            # 3) Process rows in chunks
+            for start in tqdm(range(0, N, chunk_size)):
+                end = min(start + chunk_size, N)
+                chunk_np  = bin_indices[start:end]  # shape [chunk_size, F]
+                chunk_gpu = torch.from_numpy(chunk_np).to(self.device)  # [chunk_size, F], int8
+
+                # Accumulate raw (unscaled) leaf sums
+                chunk_preds = torch.zeros((end - start,), dtype=torch.float32, device=self.device)
+
+                # node_idx[i] tracks the current node index in the padded tree for row i
+                node_idx = torch.zeros((end - start,), dtype=torch.int32, device=self.device)
+
+                # 'active' is a boolean mask over [0..(end-start-1)], indicating which rows haven't reached a leaf
+                active = torch.ones((end - start,), dtype=torch.bool, device=self.device)
+
+                for t in range(n_trees):
+                    # Reset for each tree (each tree is independent)
+                    node_idx.fill_(0)
+                    active.fill_(True)
+
+                    tree_features = features_t[t]     # shape [max_nodes], int16
+                    tree_thresh   = thresholds_t[t]    # shape [max_nodes], int16
+                    tree_values   = values_t[t]          # shape [max_nodes], float32
+
+                    # Up to self.max_depth+1 layers
+                    for _level in range(self.max_depth + 1):
+                        active_idx = active.nonzero(as_tuple=True)[0]
+                        if active_idx.numel() == 0:
+                            break  # all rows are done in this tree
+
+                        current_node_idx = node_idx[active_idx]
+                        f    = tree_features[current_node_idx]    # shape [#active], int16
+                        thr  = tree_thresh[current_node_idx]       # shape [#active], int16
+                        vals = tree_values[current_node_idx]       # shape [#active], float32
+
+                        mask_no_node = (f == -2)
+                        mask_leaf    = (f == -1)
+
+                        # If leaf, add leaf value and mark inactive.
+                        if mask_leaf.any():
+                            leaf_rows = active_idx[mask_leaf]
+                            chunk_preds[leaf_rows] += vals[mask_leaf]
+                            active[leaf_rows] = False
+
+                        # If no node, mark inactive.
+                        if mask_no_node.any():
+                            no_node_rows = active_idx[mask_no_node]
+                            active[no_node_rows] = False
+
+                        # For internal nodes, perform bin comparison.
+                        mask_internal = (~mask_leaf & ~mask_no_node)
+                        if mask_internal.any():
+                            internal_rows = active_idx[mask_internal]
+                            act_f   = f[mask_internal].long()
+                            act_thr = thr[mask_internal]
+                            binvals = chunk_gpu[internal_rows, act_f]
+                            go_left = (binvals <= act_thr)
+                            new_left_idx  = current_node_idx[mask_internal] * 2 + 1
+                            new_right_idx = current_node_idx[mask_internal] * 2 + 2
+                            node_idx[internal_rows[go_left]]  = new_left_idx[go_left]
+                            node_idx[internal_rows[~go_left]] = new_right_idx[~go_left]
+                    # end per-tree layer loop
+                # end for each tree
+
+                out[start:end] = (
+                    self.base_prediction + self.learning_rate * chunk_preds
+                ).cpu().numpy()
+
+            return out
 
     def flatten_forest_to_tensors(self, forest):
         """

warpgbm-0.1.18/warpgbm/cuda/best_split_kernel.cu

@@ -0,0 +1,79 @@
+#include <torch/extension.h>
+#include <cuda.h>
+#include <cuda_runtime.h>
+
+__global__ void best_split_kernel_global_only(
+    const float *__restrict__ G, // [F x B]
+    const float *__restrict__ H, // [F x B]
+    int F,
+    int B,
+    float min_split_gain,
+    float min_child_samples,
+    float eps,
+    float *__restrict__ best_gains, // [F]
+    int *__restrict__ best_bins     // [F]
+)
+{
+    int f = blockIdx.x * blockDim.x + threadIdx.x;
+    if (f >= F)
+        return;
+
+    float G_total = 0.0f, H_total = 0.0f;
+    for (int b = 0; b < B; ++b)
+    {
+        G_total += G[f * B + b];
+        H_total += H[f * B + b];
+    }
+
+    float G_L = 0.0f, H_L = 0.0f;
+    float best_gain = min_split_gain;
+    int best_bin = -1;
+
+    for (int b = 0; b < B - 1; ++b)
+    {
+        G_L += G[f * B + b];
+        H_L += H[f * B + b];
+        float G_R = G_total - G_L;
+        float H_R = H_total - H_L;
+
+        if (H_L >= min_child_samples && H_R >= min_child_samples)
+        {
+            float gain = (G_L * G_L) / (H_L + eps) + (G_R * G_R) / (H_R + eps);
+            if (gain > best_gain)
+            {
+                best_gain = gain;
+                best_bin = b;
+            }
+        }
+    }
+
+    best_gains[f] = best_gain;
+    best_bins[f] = best_bin;
+}
+
+void launch_best_split_kernel_cuda(
+    const at::Tensor &G, // [F x B]
+    const at::Tensor &H, // [F x B]
+    float min_split_gain,
+    float min_child_samples,
+    float eps,
+    at::Tensor &best_gains, // [F], float32
+    at::Tensor &best_bins,  // [F], int32
+    int threads)
+{
+    int F = G.size(0);
+    int B = G.size(1);
+
+    int blocks = (F + threads - 1) / threads;
+
+    best_split_kernel_global_only<<<blocks, threads>>>(
+        G.data_ptr<float>(),
+        H.data_ptr<float>(),
+        F,
+        B,
+        min_split_gain,
+        min_child_samples,
+        eps,
+        best_gains.data_ptr<float>(),
+        best_bins.data_ptr<int>());
+}

warpgbm-0.1.18/warpgbm/cuda/binner.cu

@@ -0,0 +1,52 @@
+#include <torch/extension.h>
+#include <cuda.h>
+#include <cuda_runtime.h>
+
+__global__ void bin_column_kernel(
+    const float *__restrict__ X,         // [N]
+    const float *__restrict__ bin_edges, // [B - 1]
+    int8_t *__restrict__ bin_indices,    // [N]
+    int N,
+    int B_minus1)
+{
+    int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    if (idx >= N)
+        return;
+
+    float val = X[idx];
+    int bin = 0;
+
+    // Linear scan over edges: bin_edges is sorted
+    while (bin < B_minus1 && val >= bin_edges[bin])
+    {
+        ++bin;
+    }
+
+    bin_indices[idx] = static_cast<int8_t>(bin);
+}
+
+// C++ launcher for calling from Python
+void launch_bin_column_kernel(
+    at::Tensor X,          // [N]
+    at::Tensor bin_edges,  // [B - 1]
+    at::Tensor bin_indices // [N]
+)
+{
+    const int N = X.size(0);
+    const int B = bin_edges.size(0);
+
+    const int threads = 256;
+    const int blocks = (N + threads - 1) / threads;
+
+    bin_column_kernel<<<blocks, threads>>>(
+        X.data_ptr<float>(),
+        bin_edges.data_ptr<float>(),
+        bin_indices.data_ptr<int8_t>(),
+        N,
+        B);
+
+    // Optional: sync and error check
+    cudaError_t err = cudaGetLastError();
+    if (err != cudaSuccess)
+        printf("CUDA error: %s\n", cudaGetErrorString(err));
+}

{warpgbm-0.1.16 → warpgbm-0.1.18}/warpgbm/cuda/node_kernel.cpp

@@ -21,15 +21,14 @@ void launch_histogram_kernel_cuda_2(
     int rows_per_thread = 1);
 
 void launch_best_split_kernel_cuda(
-    const at::Tensor &G,
-    const at::Tensor &H,
-    int F,
-    int B,
+    const at::Tensor &G, // [F x B]
+    const at::Tensor &H, // [F x B]
     float min_split_gain,
     float min_child_samples,
     float eps,
-    at::Tensor &out_feature,
-    at::Tensor &out_bin);
+    at::Tensor &best_gains, // [F], float32
+    at::Tensor &best_bins,
+    int threads);
 
 void launch_histogram_kernel_cuda_configurable(
     const at::Tensor &bin_indices,
@@ -40,6 +39,11 @@ void launch_histogram_kernel_cuda_configurable(
     int threads_per_block = 256,
     int rows_per_thread = 1);
 
+void launch_bin_column_kernel(
+    at::Tensor X,
+    at::Tensor bin_edges,
+    at::Tensor bin_indices);
+
 // Bindings
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
 {
@@ -47,4 +51,5 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
     m.def("compute_histogram2", &launch_histogram_kernel_cuda_2, "Histogram (CUDA) 2");
     m.def("compute_histogram3", &launch_histogram_kernel_cuda_configurable, "Histogram Feature Shared Mem");
     m.def("compute_split", &launch_best_split_kernel_cuda, "Best Split (CUDA)");
+    m.def("custom_cuda_binner", &launch_bin_column_kernel, "Custom CUDA binning kernel");
 }

{warpgbm-0.1.16 → warpgbm-0.1.18/warpgbm.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: warpgbm
-Version: 0.1.16
+Version: 0.1.18
 Summary: A fast GPU-accelerated Gradient Boosted Decision Tree library with PyTorch + CUDA
 License:                     GNU GENERAL PUBLIC LICENSE
                                Version 3, 29 June 2007
@@ -706,7 +706,7 @@ WarpGBM is a high-performance, GPU-accelerated Gradient Boosted Decision Tree (G
 
 ## Performance Note
 
-In our initial tests on an NVIDIA 3090 (local) and A100 (Google Colab Pro), WarpGBM achieves **14x to 20x faster training times** compared to LightGBM using default configurations. It also consumes **significantly less RAM and CPU**. These early results hint at more thorough benchmarking to come.
+In our initial tests on an NVIDIA 3090 (local) and A100 (Google Colab Pro), WarpGBM achieves **14x to 20x faster training times** compared to LightGBM's CPU version and **2x faster** on the GPU version using default configurations. Speed also outperforms XGBoost and CatBoost on regression problems. It also consumes **significantly less RAM and CPU**. These early results hint at more thorough benchmarking to come.
 
 ---
 

{warpgbm-0.1.16 → warpgbm-0.1.18}/warpgbm.egg-info/SOURCES.txt

@@ -15,5 +15,6 @@ warpgbm.egg-info/requires.txt
 warpgbm.egg-info/top_level.txt
 warpgbm/cuda/__init__.py
 warpgbm/cuda/best_split_kernel.cu
+warpgbm/cuda/binner.cu
 warpgbm/cuda/histogram_kernel.cu
 warpgbm/cuda/node_kernel.cpp

warpgbm-0.1.16/tests/test_fit_predict_corr.py

@@ -1,66 +0,0 @@
-import numpy as np
-from warpgbm import WarpGBM
-from sklearn.datasets import make_regression
-
-def test_fit_predict_correlation():
-    np.random.seed(42)
-    N = 1_000_000
-    F = 100
-    X, y = make_regression(n_samples=N, n_features=F, noise=0.1, random_state=42)
-    era = np.zeros(N, dtype=np.int32)
-    corrs = []
-
-    model = WarpGBM(
-        max_depth = 10,
-        num_bins = 10,
-        n_estimators = 10,
-        learning_rate = 1,
-        verbosity=False,
-        histogram_computer='hist1',
-        threads_per_block=32,
-        rows_per_thread=4
-    )
-
-    model.fit(X, y, era_id=era)
-    preds = model.predict(X)
-
-    # Pearson correlation in-sample
-    corr = np.corrcoef(preds, y)[0, 1]
-    corrs.append(corr)
-
-    model = WarpGBM(
-        max_depth = 10,
-        num_bins = 10,
-        n_estimators = 10,
-        learning_rate = 1,
-        verbosity=False,
-        histogram_computer='hist2',
-        threads_per_block=32,
-        rows_per_thread=4
-    )
-
-    model.fit(X, y, era_id=era)
-    preds = model.predict(X)
-
-    # Pearson correlation in-sample
-    corr = np.corrcoef(preds, y)[0, 1]
-    corrs.append(corr)
-
-    model = WarpGBM(
-        max_depth = 10,
-        num_bins = 10,
-        n_estimators = 10,
-        learning_rate = 1,
-        verbosity=False,
-        histogram_computer='hist3',
-        threads_per_block=32,
-        rows_per_thread=4
-    )
-
-    model.fit(X, y, era_id=era)
-    preds = model.predict(X)
-
-    # Pearson correlation in-sample
-    corr = np.corrcoef(preds, y)[0, 1]
-    corrs.append(corr)
-    assert ( np.array(corrs) > 0.95 ).all(), f"In-sample correlation too low: {corr:.4f}"

warpgbm-0.1.16/version.txt

@@ -1 +0,0 @@
-0.1.16

warpgbm-0.1.16/warpgbm/cuda/best_split_kernel.cu

@@ -1,112 +0,0 @@
-#include <torch/extension.h>
-#include <cuda.h>
-#include <cuda_runtime.h>
-
-__global__ void best_split_kernel(
-    const float *__restrict__ G, // [F x B]
-    const float *__restrict__ H, // [F x B]
-    int F,
-    int B,
-    float min_split_gain,
-    float min_child_samples,
-    float eps,
-    int *out_feature,
-    int *out_bin,
-    void *shared_mem)
-{
-    int f = blockIdx.x * blockDim.x + threadIdx.x;
-    if (f >= F)
-        return;
-
-    // Cast shared memory
-    extern __shared__ char smem[];
-    float *gains = reinterpret_cast<float *>(smem);
-    int *features = reinterpret_cast<int *>(&gains[blockDim.x]);
-    int *bins = reinterpret_cast<int *>(&features[blockDim.x]);
-
-    // Calculate total G and H for this feature
-    float G_total = 0.0f, H_total = 0.0f;
-    for (int b = 0; b < B; ++b)
-    {
-        G_total += G[f * B + b];
-        H_total += H[f * B + b];
-    }
-
-    float G_L = 0.0f, H_L = 0.0f;
-    float best_gain = min_split_gain;
-    int best_bin = -1;
-
-    for (int b = 0; b < B - 1; ++b)
-    {
-        G_L += G[f * B + b];
-        H_L += H[f * B + b];
-        float G_R = G_total - G_L;
-        float H_R = H_total - H_L;
-
-        if (H_L > min_child_samples && H_R > min_child_samples)
-        {
-            float gain = (G_L * G_L) / (H_L + eps) + (G_R * G_R) / (H_R + eps);
-            if (gain > best_gain)
-            {
-                best_gain = gain;
-                best_bin = b;
-            }
-        }
-    }
-
-    gains[threadIdx.x] = best_gain;
-    features[threadIdx.x] = f;
-    bins[threadIdx.x] = best_bin;
-    __syncthreads();
-
-    // Thread 0 in each block finds best among its block
-    if (threadIdx.x == 0)
-    {
-        float block_best_gain = min_split_gain;
-        int block_best_feature = -1;
-        int block_best_bin = -1;
-        for (int i = 0; i < blockDim.x && blockIdx.x * blockDim.x + i < F; ++i)
-        {
-            if (gains[i] > block_best_gain)
-            {
-                block_best_gain = gains[i];
-                block_best_feature = features[i];
-                block_best_bin = bins[i];
-            }
-        }
-
-        // Write to global outputs
-        *out_feature = block_best_feature;
-        *out_bin = block_best_bin;
-    }
-}
-
-void launch_best_split_kernel_cuda(
-    const at::Tensor &G,
-    const at::Tensor &H,
-    int F,
-    int B,
-    float min_split_gain,
-    float min_child_samples,
-    float eps,
-    at::Tensor &out_feature,
-    at::Tensor &out_bin)
-{
-    int threads = 256;
-    int blocks = (F + threads - 1) / threads;
-
-    size_t shared_mem_bytes = threads * (sizeof(float) + 2 * sizeof(int));
-
-    best_split_kernel<<<blocks, threads, shared_mem_bytes>>>(
-        G.data_ptr<float>(),
-        H.data_ptr<float>(),
-        F,
-        B,
-        min_split_gain,
-        min_child_samples,
-        eps,
-        out_feature.data_ptr<int>(),
-        out_bin.data_ptr<int>(),
-        nullptr // shared memory pointer not needed; just launch size
-    );
-}