Implementation:Rapidsai Cuml Treelite Integration

Knowledge Sources	Rapidsai_Cuml
Domains	Machine_Learning, Forest_Inference
Last Updated	2026-02-08 12:00 GMT

Overview

Provides integration between Treelite tree models and cuML's forest traversal framework, enabling conversion of Treelite model structures into cuML's traversal node and forest abstractions for GPU-accelerated inference.

Description

The treelite.hpp header bridges the Treelite library's tree representation with cuML's generic forest traversal framework. It defines specialized traversal types and utility functions that operate on Treelite model objects:

treelite_traversal_node: A concrete implementation of cuML's traversal_node interface that wraps a Treelite tree node. It provides methods for:

Leaf detection (is_leaf()).
Child navigation (hot_child(), distant_child()) with correct handling of comparison operators.
Split properties: feature index, threshold, inclusivity, categorical detection, and category lists.
Output retrieval for leaf nodes (scalar or vector).
Internal logic to determine which child is "hot" (default path) based on the Treelite comparison operator and categorical split direction.

treelite_traversal_forest: A concrete implementation of traversal_forest that wraps a treelite::ModelPreset. It constructs root node identifiers for all trees and implements get_node() to create traversal nodes on demand.

Utility Functions (operating on treelite::Model):

tree_for_each: Applies a lambda to each tree in the model.
tree_transform: Transforms each tree using a lambda, writing results to an output iterator.
tree_accumulate: Accumulates a value across all trees using a lambda.
node_for_each: Traverses all nodes in the forest in a specified order (depth-first, breadth-first, or layered), applying a lambda.
node_transform: Transforms each node during traversal, writing results to an output iterator.
node_accumulate: Accumulates a value across all nodes during traversal.

All node traversal functions are templated on forest_order to support different traversal strategies.

Usage

Use this integration when loading Treelite models into cuML's Forest Inference Library (FIL) for GPU-accelerated prediction. The traversal abstractions enable generic algorithms to operate on Treelite models without coupling to Treelite-specific types. This is the primary pathway for importing trained tree models (from XGBoost, LightGBM, etc.) into the cuML inference pipeline.

Code Reference

Source Location

Repository: Rapidsai_Cuml
File: cpp/include/cuml/forest/integrations/treelite.hpp

Signature

namespace ML {
namespace forest {

using TREELITE_NODE_ID_T = int;

template <typename tl_threshold_t, typename tl_output_t>
struct treelite_traversal_node : public traversal_node<TREELITE_NODE_ID_T> {
  treelite_traversal_node(treelite::Tree<tl_threshold_t, tl_output_t> const& tl_tree,
                          id_type node_id);
  bool is_leaf() const override;
  id_type hot_child() const override;
  id_type distant_child() const override;
  auto default_distant() const;
  auto get_feature() const;
  auto is_inclusive() const;
  auto is_categorical() const;
  auto get_categories() const;
  auto threshold() const;
  auto max_num_categories() const;
  auto get_output() const;
  auto get_treelite_id() const;
};

template <typename tl_threshold_t, typename tl_output_t>
struct treelite_traversal_forest
  : public traversal_forest<treelite_traversal_node<tl_threshold_t, tl_output_t>> {
  treelite_traversal_forest(
    treelite::ModelPreset<tl_threshold_t, tl_output_t> const& tl_model);
  node_type get_node(tree_id_type tree_id, node_id_type node_id) const override;
};

template <typename lambda_t>
void tree_for_each(treelite::Model const& tl_model, lambda_t&& lambda);

template <typename iter_t, typename lambda_t>
void tree_transform(treelite::Model const& tl_model, iter_t out_iter, lambda_t&& lambda);

template <typename T, typename lambda_t>
auto tree_accumulate(treelite::Model const& tl_model, T init, lambda_t&& lambda);

template <forest_order order, typename lambda_t>
void node_for_each(treelite::Model const& tl_model, lambda_t&& lambda);

template <forest_order order, typename iter_t, typename lambda_t>
void node_transform(treelite::Model const& tl_model, iter_t output_iter, lambda_t&& lambda);

template <forest_order order, typename T, typename lambda_t>
auto node_accumulate(treelite::Model const& tl_model, T init, lambda_t&& lambda);

} // namespace forest
} // namespace ML

Import

#include <cuml/forest/integrations/treelite.hpp>

I/O Contract

Inputs

Name	Type	Required	Description
tl_model	treelite::ModelPreset<tl_threshold_t, tl_output_t> const&	Yes	Treelite model containing the tree ensemble
tl_tree	treelite::Tree<tl_threshold_t, tl_output_t> const&	Yes	Individual Treelite tree (for node construction)
node_id	TREELITE_NODE_ID_T (int)	Yes	Node identifier within a tree
lambda	lambda_t	Yes	Callable for traversal operations
order	forest_order	Yes (template)	Traversal order: depth_first, breadth_first, layered_children_segregated, or layered_children_together

Outputs

Name	Type	Description
treelite_traversal_node methods	various	Node properties: is_leaf, children, feature, threshold, output, categories
tree_for_each	void	Applies lambda to each tree (side effects only)
tree_transform	void	Writes transformed results to output iterator
tree_accumulate	T	Accumulated value across all trees
node_for_each	void	Applies lambda to each node in traversal order
node_transform	void	Writes transformed node results to output iterator
node_accumulate	T	Accumulated value across all nodes

Usage Examples

#include <cuml/forest/integrations/treelite.hpp>

// Assume tl_model is a loaded treelite::Model
treelite::Model const& tl_model = /* loaded model */;

// Count total nodes in the forest using depth-first traversal
auto total_nodes = ML::forest::node_accumulate<ML::forest::forest_order::depth_first>(
    tl_model,
    std::size_t{0},
    [](auto acc, auto&& tree_id, auto&& node, auto&& depth, auto&& parent_index) {
        return acc + 1;
    });

// Iterate over all trees
ML::forest::tree_for_each(tl_model, [](auto&& tree) {
    // Process each tree
});

// Find maximum depth across all trees
auto max_depth = ML::forest::node_accumulate<ML::forest::forest_order::depth_first>(
    tl_model,
    std::size_t{0},
    [](auto acc, auto&& tree_id, auto&& node, auto&& depth, auto&& parent_index) {
        return std::max(acc, depth);
    });

Related Pages

Environment:Rapidsai_Cuml_CUDA_GPU

Page Connections

Double-click a node to navigate. Hold to expand connections.

Principle

Implementation

Heuristic

Environment