From af890c6d0337f52a49e6e0f460e40d16824ae070 Mon Sep 17 00:00:00 2001
From: Hyunsu Cho <phcho@nvidia.com>
Date: Fri, 24 Apr 2026 15:58:34 -0700
Subject: [PATCH 01/17] Validate Treelite input to avoid overflow

---
 cpp/include/cuml/fil/detail/bitset.hpp             | 6 +++---
 cpp/include/cuml/fil/detail/raft_proto/ceildiv.hpp | 4 ++--
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/cpp/include/cuml/fil/detail/bitset.hpp b/cpp/include/cuml/fil/detail/bitset.hpp
index 6981c29d1a..8c1d12bb05 100644
--- a/cpp/include/cuml/fil/detail/bitset.hpp
+++ b/cpp/include/cuml/fil/detail/bitset.hpp
@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 #pragma once
@@ -39,12 +39,12 @@ struct bitset {
   // Standard bit-wise mutators and accessor
   HOST DEVICE auto& set(index_type index)
   {
-    data_[bin_from_index(index)] |= mask_in_bin(index);
+    if (index < num_bits_) { data_[bin_from_index(index)] |= mask_in_bin(index); }
     return *this;
   }
   HOST DEVICE auto& clear(index_type index)
   {
-    data_[bin_from_index(index)] &= ~mask_in_bin(index);
+    if (index < num_bits_) { data_[bin_from_index(index)] &= ~mask_in_bin(index); }
     return *this;
   }
   HOST DEVICE auto test(index_type index) const
diff --git a/cpp/include/cuml/fil/detail/raft_proto/ceildiv.hpp b/cpp/include/cuml/fil/detail/raft_proto/ceildiv.hpp
index 61887f6ffb..92a1eaf819 100644
--- a/cpp/include/cuml/fil/detail/raft_proto/ceildiv.hpp
+++ b/cpp/include/cuml/fil/detail/raft_proto/ceildiv.hpp
@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 #pragma once
@@ -9,6 +9,6 @@ namespace raft_proto {
 template <typename T, typename U>
 HOST DEVICE auto constexpr ceildiv(T dividend, U divisor)
 {
-  return (dividend + divisor - T{1}) / divisor;
+  return dividend / divisor + (dividend % divisor != 0);
 }
 }  // namespace raft_proto

From e1c460642437df66348911476903cba5d5eccd5b Mon Sep 17 00:00:00 2001
From: Hyunsu Cho <phcho@nvidia.com>
Date: Fri, 24 Apr 2026 17:09:54 -0700
Subject: [PATCH 02/17] Add validation to add_categorical_node

---
 .../fil/detail/decision_forest_builder.hpp    | 37 ++++++++++++++-----
 1 file changed, 27 insertions(+), 10 deletions(-)

diff --git a/cpp/include/cuml/fil/detail/decision_forest_builder.hpp b/cpp/include/cuml/fil/detail/decision_forest_builder.hpp
index 0fa8ac029c..01b12313c4 100644
--- a/cpp/include/cuml/fil/detail/decision_forest_builder.hpp
+++ b/cpp/include/cuml/fil/detail/decision_forest_builder.hpp
@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 #pragma once
@@ -57,20 +57,37 @@ struct decision_forest_builder {
     typename node_type::metadata_storage_type feature = typename node_type::metadata_storage_type{},
     typename node_type::offset_type offset            = typename node_type::offset_type{})
   {
-    auto constexpr const bin_width = index_type(sizeof(typename node_type::index_type) * 8);
-    auto node_value                = typename node_type::index_type{};
-    auto set_storage               = &node_value;
-    auto max_node_categories =
-      (vec_begin != vec_end) ? *std::max_element(vec_begin, vec_end) + 1 : 1;
+    auto constexpr const bin_width =
+      typename node_type::index_type{sizeof(typename node_type::index_type) * 8};
+    auto node_value  = typename node_type::index_type{};
+    auto set_storage = &node_value;
+
+    // Ensure that all category indices are non-negative.
+    auto const is_negative = [](typename iter_t::value_type x) { return x < 0; };
+    if (std::any_of(vec_begin, vec_end, is_negative)) {
+      throw model_builder_error("Category index must be non-negative");
+    }
+
+    // Ensure that (max_cat + 1) can be represented as node_type::index_type
+    // to prevent integer overflow.
+    auto max_cat                 = (vec_begin != vec_end) ? *std::max_element(vec_begin, vec_end)
+                                                          : typename iter_t::value_type{0};
+    auto const max_representable = std::numeric_limits<typename node_type::index_type>::max();
+    if (max_cat == std::numeric_limits<typename iter_t::value_type>::max() ||
+        max_cat >= max_representable) {
+      throw model_builder_error(
+        "Category index exceeds maximum representable value for this model's index type");
+    }
+    auto max_cat_plus_one =
+      static_cast<typename node_type::index_type>(max_cat) + typename node_type::index_type{1};
     if (max_num_categories_ > bin_width) {
-      // TODO(wphicks): Check for overflow here
       node_value         = categorical_storage_.size();
-      auto bins_required = raft_proto::ceildiv(max_node_categories, bin_width);
-      categorical_storage_.push_back(max_node_categories);
+      auto bins_required = raft_proto::ceildiv(max_cat_plus_one, bin_width);
+      categorical_storage_.push_back(max_cat_plus_one);
       categorical_storage_.resize(categorical_storage_.size() + bins_required);
       set_storage = &(categorical_storage_[node_value + 1]);
     }
-    auto set = bitset{set_storage, max_node_categories};
+    auto set = bitset{set_storage, max_cat_plus_one};
     std::for_each(vec_begin, vec_end, [&set](auto&& cat_index) { set.set(cat_index); });
 
     add_node(

From 0dac58a7574b5d32bf1024f8ea18024d3e9985f3 Mon Sep 17 00:00:00 2001
From: Hyunsu Cho <phcho@nvidia.com>
Date: Fri, 24 Apr 2026 17:29:31 -0700
Subject: [PATCH 03/17] Incorporate feedback from CodeRabbit

---
 cpp/include/cuml/fil/detail/bitset.hpp                 |  2 ++
 .../cuml/fil/detail/decision_forest_builder.hpp        | 10 +++++-----
 2 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/cpp/include/cuml/fil/detail/bitset.hpp b/cpp/include/cuml/fil/detail/bitset.hpp
index 8c1d12bb05..7f34b2322b 100644
--- a/cpp/include/cuml/fil/detail/bitset.hpp
+++ b/cpp/include/cuml/fil/detail/bitset.hpp
@@ -22,6 +22,8 @@ struct bitset {
   using storage_type = storage_t;
   using index_type   = index_t;
 
+  static_assert(std::is_unsigned_v<index_t>, "index_t must be unsigned");
+
   auto constexpr static const bin_width = index_type(sizeof(storage_type) * 8);
 
   HOST DEVICE bitset() : data_{nullptr}, num_bits_{0} {}
diff --git a/cpp/include/cuml/fil/detail/decision_forest_builder.hpp b/cpp/include/cuml/fil/detail/decision_forest_builder.hpp
index 01b12313c4..39bdcaea27 100644
--- a/cpp/include/cuml/fil/detail/decision_forest_builder.hpp
+++ b/cpp/include/cuml/fil/detail/decision_forest_builder.hpp
@@ -18,6 +18,7 @@
 #include <algorithm>
 #include <cmath>
 #include <cstddef>
+#include <iterator>
 #include <numeric>
 #include <optional>
 #include <vector>
@@ -63,18 +64,17 @@ struct decision_forest_builder {
     auto set_storage = &node_value;
 
     // Ensure that all category indices are non-negative.
-    auto const is_negative = [](typename iter_t::value_type x) { return x < 0; };
+    using cat_t            = typename std::iterator_traits<iter_t>::value_type;
+    auto const is_negative = [](cat_t x) { return x < 0; };
     if (std::any_of(vec_begin, vec_end, is_negative)) {
       throw model_builder_error("Category index must be non-negative");
     }
 
     // Ensure that (max_cat + 1) can be represented as node_type::index_type
     // to prevent integer overflow.
-    auto max_cat                 = (vec_begin != vec_end) ? *std::max_element(vec_begin, vec_end)
-                                                          : typename iter_t::value_type{0};
+    auto max_cat = (vec_begin != vec_end) ? *std::max_element(vec_begin, vec_end) : cat_t{0};
     auto const max_representable = std::numeric_limits<typename node_type::index_type>::max();
-    if (max_cat == std::numeric_limits<typename iter_t::value_type>::max() ||
-        max_cat >= max_representable) {
+    if (max_cat >= max_representable || max_cat + 1 >= max_representable) {
       throw model_builder_error(
         "Category index exceeds maximum representable value for this model's index type");
     }

From b91ec68cb6a8f2194d59367a56a351c4fd5b0be9 Mon Sep 17 00:00:00 2001
From: Hyunsu Cho <phcho@nvidia.com>
Date: Fri, 24 Apr 2026 17:32:27 -0700
Subject: [PATCH 04/17] Elide sign check if cat_t is unsigned

---
 .../cuml/fil/detail/decision_forest_builder.hpp        | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/cpp/include/cuml/fil/detail/decision_forest_builder.hpp b/cpp/include/cuml/fil/detail/decision_forest_builder.hpp
index 39bdcaea27..88ad865a23 100644
--- a/cpp/include/cuml/fil/detail/decision_forest_builder.hpp
+++ b/cpp/include/cuml/fil/detail/decision_forest_builder.hpp
@@ -64,10 +64,12 @@ struct decision_forest_builder {
     auto set_storage = &node_value;
 
     // Ensure that all category indices are non-negative.
-    using cat_t            = typename std::iterator_traits<iter_t>::value_type;
-    auto const is_negative = [](cat_t x) { return x < 0; };
-    if (std::any_of(vec_begin, vec_end, is_negative)) {
-      throw model_builder_error("Category index must be non-negative");
+    using cat_t = typename std::iterator_traits<iter_t>::value_type;
+    if constexpr (std::is_signed_v<cat_t>) {
+      auto const is_negative = [](cat_t x) { return x < 0; };
+      if (std::any_of(vec_begin, vec_end, is_negative)) {
+        throw model_builder_error("Category index must be non-negative");
+      }
     }
 
     // Ensure that (max_cat + 1) can be represented as node_type::index_type

From 8632d9e6414987c4d8262ef123acfa9201fd64de Mon Sep 17 00:00:00 2001
From: Hyunsu Cho <phcho@nvidia.com>
Date: Mon, 27 Apr 2026 15:51:46 -0700
Subject: [PATCH 05/17] Address review comments

---
 cpp/include/cuml/fil/detail/bitset.hpp             | 2 ++
 cpp/include/cuml/fil/detail/raft_proto/ceildiv.hpp | 3 +++
 2 files changed, 5 insertions(+)

diff --git a/cpp/include/cuml/fil/detail/bitset.hpp b/cpp/include/cuml/fil/detail/bitset.hpp
index 7f34b2322b..edbe5140ed 100644
--- a/cpp/include/cuml/fil/detail/bitset.hpp
+++ b/cpp/include/cuml/fil/detail/bitset.hpp
@@ -22,6 +22,7 @@ struct bitset {
   using storage_type = storage_t;
   using index_type   = index_t;
 
+  // Ensrue that index_t is unsigned. Bound checks below rely on index_t being unsigned
   static_assert(std::is_unsigned_v<index_t>, "index_t must be unsigned");
 
   auto constexpr static const bin_width = index_type(sizeof(storage_type) * 8);
@@ -41,6 +42,7 @@ struct bitset {
   // Standard bit-wise mutators and accessor
   HOST DEVICE auto& set(index_type index)
   {
+    // Guard against OOB writes; silently ignored to preserve memory safety
     if (index < num_bits_) { data_[bin_from_index(index)] |= mask_in_bin(index); }
     return *this;
   }
diff --git a/cpp/include/cuml/fil/detail/raft_proto/ceildiv.hpp b/cpp/include/cuml/fil/detail/raft_proto/ceildiv.hpp
index 92a1eaf819..429b8af64a 100644
--- a/cpp/include/cuml/fil/detail/raft_proto/ceildiv.hpp
+++ b/cpp/include/cuml/fil/detail/raft_proto/ceildiv.hpp
@@ -5,10 +5,13 @@
 #pragma once
 #include <cuml/fil/detail/raft_proto/gpu_support.hpp>
 
+#include <type_traits>
+
 namespace raft_proto {
 template <typename T, typename U>
 HOST DEVICE auto constexpr ceildiv(T dividend, U divisor)
 {
+  static_assert(std::is_integral_v<T> && std::is_integral_v<U>, "Arguments must be integers");
   return dividend / divisor + (dividend % divisor != 0);
 }
 }  // namespace raft_proto

From 5eed8efd76788c552d9da05fd3ec106e101f34f7 Mon Sep 17 00:00:00 2001
From: Hyunsu Cho <phcho@nvidia.com>
Date: Mon, 27 Apr 2026 16:06:31 -0700
Subject: [PATCH 06/17] Update integer check in add_categorical_node()

---
 .../fil/detail/decision_forest_builder.hpp    | 32 +++++++++++--------
 cpp/include/cuml/fil/treelite_importer.hpp    |  3 +-
 2 files changed, 20 insertions(+), 15 deletions(-)

diff --git a/cpp/include/cuml/fil/detail/decision_forest_builder.hpp b/cpp/include/cuml/fil/detail/decision_forest_builder.hpp
index 88ad865a23..c0b967d3ae 100644
--- a/cpp/include/cuml/fil/detail/decision_forest_builder.hpp
+++ b/cpp/include/cuml/fil/detail/decision_forest_builder.hpp
@@ -21,6 +21,8 @@
 #include <iterator>
 #include <numeric>
 #include <optional>
+#include <string>
+#include <utility>
 #include <vector>
 
 namespace ML {
@@ -33,10 +35,11 @@ namespace detail {
 struct model_builder_error : std::exception {
   model_builder_error() : model_builder_error("Error while building model") {}
   model_builder_error(char const* msg) : msg_{msg} {}
-  virtual char const* what() const noexcept { return msg_; }
+  model_builder_error(std::string msg) : msg_{std::move(msg)} {}
+  virtual char const* what() const noexcept { return msg_.c_str(); }
 
  private:
-  char const* msg_;
+  std::string msg_;
 };
 
 /*
@@ -52,6 +55,7 @@ struct decision_forest_builder {
   void add_categorical_node(
     iter_t vec_begin,
     iter_t vec_end,
+    std::size_t tree_id,
     std::optional<int> tl_node_id                     = std::nullopt,
     std::size_t depth                                 = std::size_t{1},
     bool default_to_distant_child                     = false,
@@ -65,23 +69,23 @@ struct decision_forest_builder {
 
     // Ensure that all category indices are non-negative.
     using cat_t = typename std::iterator_traits<iter_t>::value_type;
-    if constexpr (std::is_signed_v<cat_t>) {
-      auto const is_negative = [](cat_t x) { return x < 0; };
-      if (std::any_of(vec_begin, vec_end, is_negative)) {
-        throw model_builder_error("Category index must be non-negative");
-      }
-    }
+    static_assert(std::is_unsigned_v<cat_t>, "Category value must be an unsigned integer type");
+
+    auto max_cat = (vec_begin != vec_end) ? *std::max_element(vec_begin, vec_end) : cat_t{0};
 
     // Ensure that (max_cat + 1) can be represented as node_type::index_type
     // to prevent integer overflow.
-    auto max_cat = (vec_begin != vec_end) ? *std::max_element(vec_begin, vec_end) : cat_t{0};
-    auto const max_representable = std::numeric_limits<typename node_type::index_type>::max();
-    if (max_cat >= max_representable || max_cat + 1 >= max_representable) {
+    using index_t           = typename node_type::index_type;
+    auto const max_index    = static_cast<std::uintmax_t>(std::numeric_limits<index_t>::max());
+    auto const cat_unsigned = static_cast<std::uintmax_t>(max_cat);
+    if (cat_unsigned >= max_index) {
+      auto node_id_repr =
+        tl_node_id.has_value() ? std::to_string(tl_node_id.value()) : std::string{"n/a"};
       throw model_builder_error(
-        "Category index exceeds maximum representable value for this model's index type");
+        std::string{"Tree "} + std::to_string(tree_id) + ", Node " + node_id_repr +
+        ": Category index exceeds maximum representable value for this model's index type");
     }
-    auto max_cat_plus_one =
-      static_cast<typename node_type::index_type>(max_cat) + typename node_type::index_type{1};
+    auto max_cat_plus_one = static_cast<index_t>(cat_unsigned + 1u);
     if (max_num_categories_ > bin_width) {
       node_value         = categorical_storage_.size();
       auto bins_required = raft_proto::ceildiv(max_cat_plus_one, bin_width);
diff --git a/cpp/include/cuml/fil/treelite_importer.hpp b/cpp/include/cuml/fil/treelite_importer.hpp
index be3fdb217c..24e4ee9551 100644
--- a/cpp/include/cuml/fil/treelite_importer.hpp
+++ b/cpp/include/cuml/fil/treelite_importer.hpp
@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 #pragma once
@@ -273,6 +273,7 @@ struct treelite_importer {
                 auto categories = node.get_categories();
                 builder.add_categorical_node(std::begin(categories),
                                              std::end(categories),
+                                             tree_id,
                                              node.get_treelite_id(),
                                              depth,
                                              node.default_distant(),

From bb12e5b02b4eebcdd51d5b0c82b9f177ee8d7400 Mon Sep 17 00:00:00 2001
From: Hyunsu Cho <phcho@nvidia.com>
Date: Mon, 27 Apr 2026 16:36:14 -0700
Subject: [PATCH 07/17] Add missing inline keywords

---
 .../cuml/fil/detail/degenerate_trees.hpp      |  4 +--
 cpp/include/cuml/fil/treelite_importer.hpp    | 30 ++++++++++---------
 2 files changed, 18 insertions(+), 16 deletions(-)

diff --git a/cpp/include/cuml/fil/detail/degenerate_trees.hpp b/cpp/include/cuml/fil/detail/degenerate_trees.hpp
index a0d20714f7..5dcf19d208 100644
--- a/cpp/include/cuml/fil/detail/degenerate_trees.hpp
+++ b/cpp/include/cuml/fil/detail/degenerate_trees.hpp
@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 #pragma once
@@ -17,7 +17,7 @@ namespace ML::fil::detail {
 // This function returns a modified copy of a given Treelite model if it contains
 // at least one degenerate tree (a single root node with no child).
 // If the model contains no degenerate tree, then the function returns nullptr.
-std::unique_ptr<treelite::Model> convert_degenerate_trees(treelite::Model const& tl_model)
+inline std::unique_ptr<treelite::Model> convert_degenerate_trees(treelite::Model const& tl_model)
 {
   bool contains_degenerate =
     ML::forest::tree_accumulate(tl_model, false, [](auto&& contains, auto&& tree) {
diff --git a/cpp/include/cuml/fil/treelite_importer.hpp b/cpp/include/cuml/fil/treelite_importer.hpp
index 24e4ee9551..122c0ef268 100644
--- a/cpp/include/cuml/fil/treelite_importer.hpp
+++ b/cpp/include/cuml/fil/treelite_importer.hpp
@@ -430,13 +430,14 @@ struct treelite_importer {
  * @param stream The CUDA stream to use for loading this model (can be
  * omitted for CPU).
  */
-auto import_from_treelite_model(treelite::Model const& tl_model,
-                                tree_layout layout                       = preferred_tree_layout,
-                                index_type align_bytes                   = index_type{},
-                                std::optional<bool> use_double_precision = std::nullopt,
-                                raft_proto::device_type dev_type = raft_proto::device_type::cpu,
-                                int device                       = 0,
-                                raft_proto::cuda_stream stream   = raft_proto::cuda_stream{})
+inline auto import_from_treelite_model(
+  treelite::Model const& tl_model,
+  tree_layout layout                       = preferred_tree_layout,
+  index_type align_bytes                   = index_type{},
+  std::optional<bool> use_double_precision = std::nullopt,
+  raft_proto::device_type dev_type         = raft_proto::device_type::cpu,
+  int device                               = 0,
+  raft_proto::cuda_stream stream           = raft_proto::cuda_stream{})
 {
   auto result = forest_model{};
   switch (layout) {
@@ -480,13 +481,14 @@ auto import_from_treelite_model(treelite::Model const& tl_model,
  * @param stream The CUDA stream to use for loading this model (can be
  * omitted for CPU).
  */
-auto import_from_treelite_handle(TreeliteModelHandle tl_handle,
-                                 tree_layout layout                       = preferred_tree_layout,
-                                 index_type align_bytes                   = index_type{},
-                                 std::optional<bool> use_double_precision = std::nullopt,
-                                 raft_proto::device_type dev_type = raft_proto::device_type::cpu,
-                                 int device                       = 0,
-                                 raft_proto::cuda_stream stream   = raft_proto::cuda_stream{})
+inline auto import_from_treelite_handle(
+  TreeliteModelHandle tl_handle,
+  tree_layout layout                       = preferred_tree_layout,
+  index_type align_bytes                   = index_type{},
+  std::optional<bool> use_double_precision = std::nullopt,
+  raft_proto::device_type dev_type         = raft_proto::device_type::cpu,
+  int device                               = 0,
+  raft_proto::cuda_stream stream           = raft_proto::cuda_stream{})
 {
   return import_from_treelite_model(*static_cast<treelite::Model*>(tl_handle),
                                     layout,

From b404011f14c99662a6b67923cee903f4d2f234b0 Mon Sep 17 00:00:00 2001
From: Hyunsu Cho <phcho@nvidia.com>
Date: Mon, 27 Apr 2026 17:40:10 -0700
Subject: [PATCH 08/17] Simplify category bounds check + add gtest

---
 .../fil/detail/decision_forest_builder.hpp    |  33 ++---
 cpp/tests/CMakeLists.txt                      |   5 +-
 .../fil/treelite_importer_invalid_inputs.cpp  | 115 ++++++++++++++++++
 3 files changed, 136 insertions(+), 17 deletions(-)
 create mode 100644 cpp/tests/sg/fil/treelite_importer_invalid_inputs.cpp

diff --git a/cpp/include/cuml/fil/detail/decision_forest_builder.hpp b/cpp/include/cuml/fil/detail/decision_forest_builder.hpp
index c0b967d3ae..74cc26c122 100644
--- a/cpp/include/cuml/fil/detail/decision_forest_builder.hpp
+++ b/cpp/include/cuml/fil/detail/decision_forest_builder.hpp
@@ -67,25 +67,26 @@ struct decision_forest_builder {
     auto node_value  = typename node_type::index_type{};
     auto set_storage = &node_value;
 
-    // Ensure that all category indices are non-negative.
-    using cat_t = typename std::iterator_traits<iter_t>::value_type;
-    static_assert(std::is_unsigned_v<cat_t>, "Category value must be an unsigned integer type");
+    // Check invariants for data types
+    using cat_t   = typename std::iterator_traits<iter_t>::value_type;
+    using index_t = typename node_type::index_type;
+    static_assert(std::is_same_v<cat_t, std::uint32_t>, "Category value must be uint32_t");
+    static_assert(std::is_same_v<index_t, std::uint32_t> || std::is_same_v<index_t, std::uint64_t>,
+                  "Index type in tree node must be either uint32_t or uint64_t");
 
+    // Ensure that (max_cat + 1) can be represented as index_t to prevent integer overflow.
     auto max_cat = (vec_begin != vec_end) ? *std::max_element(vec_begin, vec_end) : cat_t{0};
-
-    // Ensure that (max_cat + 1) can be represented as node_type::index_type
-    // to prevent integer overflow.
-    using index_t           = typename node_type::index_type;
-    auto const max_index    = static_cast<std::uintmax_t>(std::numeric_limits<index_t>::max());
-    auto const cat_unsigned = static_cast<std::uintmax_t>(max_cat);
-    if (cat_unsigned >= max_index) {
-      auto node_id_repr =
-        tl_node_id.has_value() ? std::to_string(tl_node_id.value()) : std::string{"n/a"};
-      throw model_builder_error(
-        std::string{"Tree "} + std::to_string(tree_id) + ", Node " + node_id_repr +
-        ": Category index exceeds maximum representable value for this model's index type");
+    if constexpr (std::is_same_v<cat_t, index_t>) {
+      if (max_cat == std::numeric_limits<index_t>::max()) {
+        auto node_id_repr =
+          tl_node_id.has_value() ? std::to_string(tl_node_id.value()) : std::string{"n/a"};
+        throw model_builder_error(std::string{"Tree "} + std::to_string(tree_id) + ", Node " +
+                                  node_id_repr + ": Category index must be at most " +
+                                  std::to_string(std::numeric_limits<index_t>::max() - 1));
+      }
     }
-    auto max_cat_plus_one = static_cast<index_t>(cat_unsigned + 1u);
+    auto max_cat_plus_one = static_cast<index_t>(max_cat) + index_t{1};
+
     if (max_num_categories_ > bin_width) {
       node_value         = categorical_storage_.size();
       auto bins_required = raft_proto::ceildiv(max_cat_plus_one, bin_width);
diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt
index 25e6109cf9..1f8844d008 100644
--- a/cpp/tests/CMakeLists.txt
+++ b/cpp/tests/CMakeLists.txt
@@ -109,7 +109,10 @@ if(all_algo OR fil_algo)
   ConfigureTest(PREFIX SG NAME DEVICE_BUFFER_TEST sg/fil/raft_proto/buffer.cu ML_INCLUDE)
   ConfigureTest(PREFIX SG NAME FOREST_TRAVERSAL_TEST sg/forest/traversal_forest.cpp ML_INCLUDE)
   ConfigureTest(PREFIX SG NAME TREELITE_TRAVERSAL_TEST sg/forest/treelite_traversal.cpp ML_INCLUDE)
-  ConfigureTest(PREFIX SG NAME TREELITE_IMPORTER_TEST sg/fil/treelite_importer.cpp ML_INCLUDE)
+  ConfigureTest(
+    PREFIX SG NAME TREELITE_IMPORTER_TEST sg/fil/treelite_importer.cpp
+                                          sg/fil/treelite_importer_invalid_inputs.cpp ML_INCLUDE
+  )
 endif()
 
 # todo: organize linear models better
diff --git a/cpp/tests/sg/fil/treelite_importer_invalid_inputs.cpp b/cpp/tests/sg/fil/treelite_importer_invalid_inputs.cpp
new file mode 100644
index 0000000000..4c7cc6f510
--- /dev/null
+++ b/cpp/tests/sg/fil/treelite_importer_invalid_inputs.cpp
@@ -0,0 +1,115 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#include <cuml/fil/detail/decision_forest_builder.hpp>
+#include <cuml/fil/tree_layout.hpp>
+#include <cuml/fil/treelite_importer.hpp>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include <treelite/enum/task_type.h>
+#include <treelite/enum/typeinfo.h>
+#include <treelite/model_builder.h>
+
+#include <cstddef>
+#include <cstdint>
+#include <limits>
+
+namespace ML {
+namespace fil {
+
+TEST(TreeliteImporter, large_category_value)
+{
+  // For tree models with 32-bit data storage,
+  // attempting to use UINT32_MAX in a categorical test node
+  // must throw an exception.
+
+  auto metadata = treelite::model_builder::Metadata{
+    1,
+    treelite::TaskType::kRegressor,
+    false,
+    1,
+    {1},
+    {1, 1},
+  };
+  auto tree_annotation = treelite::model_builder::TreeAnnotation{1, {0}, {0}};
+  auto model_builder =
+    treelite::model_builder::GetModelBuilder(treelite::TypeInfo::kFloat32,
+                                             treelite::TypeInfo::kFloat32,
+                                             metadata,
+                                             tree_annotation,
+                                             treelite::model_builder::PostProcessorFunc{"identity"},
+                                             {0.0});
+
+  model_builder->StartTree();
+  model_builder->StartNode(0);
+  model_builder->CategoricalTest(
+    0, false, {std::numeric_limits<std::uint32_t>::max()}, false, 1, 2);
+  model_builder->EndNode();
+
+  model_builder->StartNode(1);
+  model_builder->LeafScalar(1.0);
+  model_builder->EndNode();
+
+  model_builder->StartNode(2);
+  model_builder->LeafScalar(-1.0);
+  model_builder->EndNode();
+
+  model_builder->EndTree();
+
+  auto tl_model = model_builder->CommitModel();
+
+  auto expected_error_msg = std::string{"Tree 0, Node 0: Category index must be at most "} +
+                            std::to_string(std::numeric_limits<std::uint32_t>::max() - 1);
+
+  ASSERT_THAT(
+    [&]() { import_from_treelite_model(*tl_model, tree_layout::breadth_first); },
+    testing::ThrowsMessage<detail::model_builder_error>(testing::HasSubstr(expected_error_msg)));
+}
+
+TEST(TreeliteImporter, large_category_value2)
+{
+  // For tree models with 32-bit data storage,
+  // it should be possible to use (UINT32_MAX - 1) in a categorical test node
+
+  auto metadata = treelite::model_builder::Metadata{
+    1,
+    treelite::TaskType::kRegressor,
+    false,
+    1,
+    {1},
+    {1, 1},
+  };
+  auto tree_annotation = treelite::model_builder::TreeAnnotation{1, {0}, {0}};
+  auto model_builder =
+    treelite::model_builder::GetModelBuilder(treelite::TypeInfo::kFloat32,
+                                             treelite::TypeInfo::kFloat32,
+                                             metadata,
+                                             tree_annotation,
+                                             treelite::model_builder::PostProcessorFunc{"identity"},
+                                             {0.0});
+
+  model_builder->StartTree();
+  model_builder->StartNode(0);
+  model_builder->CategoricalTest(
+    0, false, {std::numeric_limits<std::uint32_t>::max() - 1}, false, 1, 2);
+  model_builder->EndNode();
+
+  model_builder->StartNode(1);
+  model_builder->LeafScalar(1.0);
+  model_builder->EndNode();
+
+  model_builder->StartNode(2);
+  model_builder->LeafScalar(-1.0);
+  model_builder->EndNode();
+
+  model_builder->EndTree();
+
+  auto tl_model = model_builder->CommitModel();
+  ASSERT_NO_THROW(import_from_treelite_model(*tl_model, tree_layout::breadth_first));
+}
+
+}  // namespace fil
+}  // namespace ML

From 99d13f69c8b15b382fee24cfbf96b14b26376181 Mon Sep 17 00:00:00 2001
From: Hyunsu Cho <phcho@nvidia.com>
Date: Mon, 27 Apr 2026 19:01:17 -0700
Subject: [PATCH 09/17] Add bound check for tree_index

---
 .../cuml/fil/detail/infer_kernel/gpu.cuh      | 95 ++++++++++---------
 1 file changed, 49 insertions(+), 46 deletions(-)

diff --git a/cpp/include/cuml/fil/detail/infer_kernel/gpu.cuh b/cpp/include/cuml/fil/detail/infer_kernel/gpu.cuh
index 95801160da..e09dd2ad7c 100644
--- a/cpp/include/cuml/fil/detail/infer_kernel/gpu.cuh
+++ b/cpp/include/cuml/fil/detail/infer_kernel/gpu.cuh
@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 #pragma once
@@ -125,56 +125,59 @@ CUML_KERNEL void __launch_bounds__(MAX_THREADS_PER_BLOCK, MIN_BLOCKS_PER_SM) inf
     // work within the loop if the task_index is below the actual task_count.
     auto const task_count_rounded_up = blockDim.x * raft_proto::ceildiv(task_count, blockDim.x);
 
-    // Infer on each tree and row
-    for (auto task_index = threadIdx.x; task_index < task_count_rounded_up;
-         task_index += blockDim.x) {
-      auto row_index = task_index % chunk_size;
-      auto real_task = task_index < task_count && row_index < rows_in_this_iteration;
-      row_index *= real_task;
-      auto tree_index  = task_index * real_task / chunk_size;
-      auto grove_index = (threadIdx.x / chunk_size) * (infer_type == infer_kind::default_kind);
-
-      auto tree_output  = std::conditional_t<has_vector_leaves,
-                                             typename node_t::index_type,
-                                             typename node_t::threshold_type>{};
-      auto leaf_node_id = index_type{};
-      if (infer_type == infer_kind::leaf_id) {
-        leaf_node_id =
-          evaluate_tree<has_vector_leaves, has_categorical_nodes, has_nonlocal_categories, true>(
-            forest, tree_index, input_data + row_index * col_count, categorical_data);
-      } else {
-        tree_output =
-          evaluate_tree<has_vector_leaves, has_categorical_nodes, has_nonlocal_categories, false>(
-            forest, tree_index, input_data + row_index * col_count, categorical_data);
-      }
+    // Ensure that tree_index doesn't overflow in the loop below
+    auto max_tree_index = (task_count - 1) / chunk_size;
+    if (max_tree_index < forest.tree_count()) {
+      // Infer on each tree and row
+      for (auto task_index = threadIdx.x; task_index < task_count_rounded_up;
+           task_index += blockDim.x) {
+        auto row_index = task_index % chunk_size;
+        auto real_task = task_index < task_count && row_index < rows_in_this_iteration;
+        row_index *= real_task;
+        auto tree_index  = task_index * real_task / chunk_size;
+        auto grove_index = (threadIdx.x / chunk_size) * (infer_type == infer_kind::default_kind);
+
+        auto tree_output  = std::conditional_t<has_vector_leaves,
+                                               typename node_t::index_type,
+                                               typename node_t::threshold_type>{};
+        auto leaf_node_id = index_type{};
+        if (infer_type == infer_kind::leaf_id) {
+          leaf_node_id =
+            evaluate_tree<has_vector_leaves, has_categorical_nodes, has_nonlocal_categories, true>(
+              forest, tree_index, input_data + row_index * col_count, categorical_data);
+        } else {
+          tree_output =
+            evaluate_tree<has_vector_leaves, has_categorical_nodes, has_nonlocal_categories, false>(
+              forest, tree_index, input_data + row_index * col_count, categorical_data);
+        }
 
-      if (infer_type == infer_kind::leaf_id) {
-        output_workspace[row_index * num_outputs * num_grove + tree_index * num_grove +
-                         grove_index] = static_cast<typename forest_t::io_type>(leaf_node_id);
-      } else {
-        if constexpr (has_vector_leaves) {
-          auto output_offset =
-            (row_index * num_outputs * num_grove +
-             tree_index * default_num_outputs * num_grove * (infer_type == infer_kind::per_tree) +
-             grove_index);
-          for (auto output_index = index_type{}; output_index < default_num_outputs;
-               ++output_index) {
-            if (real_task) {
-              output_workspace[output_offset + output_index * num_grove] +=
-                vector_output_p[tree_output * default_num_outputs + output_index];
+        if (infer_type == infer_kind::leaf_id) {
+          output_workspace[row_index * num_outputs * num_grove + tree_index * num_grove +
+                           grove_index] = static_cast<typename forest_t::io_type>(leaf_node_id);
+        } else {
+          if constexpr (has_vector_leaves) {
+            auto output_offset =
+              (row_index * num_outputs * num_grove +
+               tree_index * default_num_outputs * num_grove * (infer_type == infer_kind::per_tree) +
+               grove_index);
+            for (auto output_index = index_type{}; output_index < default_num_outputs;
+                 ++output_index) {
+              if (real_task) {
+                output_workspace[output_offset + output_index * num_grove] +=
+                  vector_output_p[tree_output * default_num_outputs + output_index];
+              }
             }
+          } else {
+            auto output_offset =
+              (row_index * num_outputs * num_grove +
+               (tree_index % default_num_outputs) * num_grove *
+                 (infer_type == infer_kind::default_kind) +
+               tree_index * num_grove * (infer_type == infer_kind::per_tree) + grove_index);
+            if (real_task) { output_workspace[output_offset] += tree_output; }
           }
-        } else {
-          auto output_offset =
-            (row_index * num_outputs * num_grove +
-             (tree_index % default_num_outputs) * num_grove *
-               (infer_type == infer_kind::default_kind) +
-             tree_index * num_grove * (infer_type == infer_kind::per_tree) + grove_index);
-          if (real_task) { output_workspace[output_offset] += tree_output; }
         }
+        __syncthreads();
       }
-
-      __syncthreads();
     }
 
     auto padded_num_groves = raft_proto::padded_size(num_grove, WARP_SIZE);

From ecfe23a3e80bb625be9ab40ae06394f34dcea53d Mon Sep 17 00:00:00 2001
From: Hyunsu Cho <phcho@nvidia.com>
Date: Tue, 28 Apr 2026 18:43:05 -0700
Subject: [PATCH 10/17] Add check on node::node(); clean up exceptions

---
 .../fil/detail/decision_forest_builder.hpp    | 49 ++++++++++---------
 cpp/include/cuml/fil/detail/node.hpp          | 40 +++++++++------
 cpp/include/cuml/fil/exceptions.hpp           |  5 +-
 cpp/include/cuml/fil/treelite_importer.hpp    | 18 ++++---
 .../fil/treelite_importer_invalid_inputs.cpp  |  7 ++-
 5 files changed, 68 insertions(+), 51 deletions(-)

diff --git a/cpp/include/cuml/fil/detail/decision_forest_builder.hpp b/cpp/include/cuml/fil/detail/decision_forest_builder.hpp
index 74cc26c122..000b0079dd 100644
--- a/cpp/include/cuml/fil/detail/decision_forest_builder.hpp
+++ b/cpp/include/cuml/fil/detail/decision_forest_builder.hpp
@@ -29,19 +29,6 @@ namespace ML {
 namespace fil {
 namespace detail {
 
-/*
- * Exception indicating that FIL model could not be built from given input
- */
-struct model_builder_error : std::exception {
-  model_builder_error() : model_builder_error("Error while building model") {}
-  model_builder_error(char const* msg) : msg_{msg} {}
-  model_builder_error(std::string msg) : msg_{std::move(msg)} {}
-  virtual char const* what() const noexcept { return msg_.c_str(); }
-
- private:
-  std::string msg_;
-};
-
 /*
  * Struct used to build FIL forests
  */
@@ -80,9 +67,9 @@ struct decision_forest_builder {
       if (max_cat == std::numeric_limits<index_t>::max()) {
         auto node_id_repr =
           tl_node_id.has_value() ? std::to_string(tl_node_id.value()) : std::string{"n/a"};
-        throw model_builder_error(std::string{"Tree "} + std::to_string(tree_id) + ", Node " +
-                                  node_id_repr + ": Category index must be at most " +
-                                  std::to_string(std::numeric_limits<index_t>::max() - 1));
+        throw model_import_error{std::string{"Tree "} + std::to_string(tree_id) + ", Node " +
+                                 node_id_repr + ": Category index must be at most " +
+                                 std::to_string(std::numeric_limits<index_t>::max() - 1)};
       }
     }
     auto max_cat_plus_one = static_cast<index_t>(max_cat) + index_t{1};
@@ -97,14 +84,23 @@ struct decision_forest_builder {
     auto set = bitset{set_storage, max_cat_plus_one};
     std::for_each(vec_begin, vec_end, [&set](auto&& cat_index) { set.set(cat_index); });
 
-    add_node(
-      node_value, tl_node_id, depth, false, default_to_distant_child, true, feature, offset, false);
+    add_node(node_value,
+             tree_id,
+             tl_node_id,
+             depth,
+             false,
+             default_to_distant_child,
+             true,
+             feature,
+             offset,
+             false);
   }
 
   /* Add a leaf node with vector output */
   template <typename iter_t>
   void add_leaf_vector_node(iter_t vec_begin,
                             iter_t vec_end,
+                            std::size_t tree_id,
                             std::optional<int> tl_node_id = std::nullopt,
                             std::size_t depth             = std::size_t{1})
   {
@@ -112,6 +108,7 @@ struct decision_forest_builder {
     std::copy(vec_begin, vec_end, std::back_inserter(vector_output_));
 
     add_node(leaf_index,
+             tree_id,
              tl_node_id,
              depth,
              true,
@@ -126,6 +123,7 @@ struct decision_forest_builder {
   template <typename value_t>
   void add_node(
     value_t val,
+    std::size_t tree_id,
     std::optional<int> tl_node_id                     = std::nullopt,
     std::size_t depth                                 = std::size_t{1},
     bool is_leaf_node                                 = true,
@@ -140,7 +138,7 @@ struct decision_forest_builder {
         if (cur_node_index_ % alignment_ != index_type{}) {
           auto padding = (alignment_ - cur_node_index_ % alignment_);
           for (auto i = index_type{}; i < padding; ++i) {
-            add_node(typename node_type::threshold_type{}, std::nullopt);
+            add_node(typename node_type::threshold_type{}, tree_id, std::nullopt);
           }
         }
       }
@@ -148,8 +146,15 @@ struct decision_forest_builder {
     }
 
     if (is_inclusive) { val = std::nextafter(val, std::numeric_limits<value_t>::infinity()); }
-    nodes_.emplace_back(
-      val, is_leaf_node, default_to_distant_child, is_categorical_node, feature, offset);
+    try {
+      nodes_.emplace_back(
+        val, is_leaf_node, default_to_distant_child, is_categorical_node, feature, offset);
+    } catch (const model_import_error& e) {
+      auto node_id_repr =
+        tl_node_id.has_value() ? std::to_string(tl_node_id.value()) : std::string{"n/a"};
+      throw model_import_error{std::string{"Tree "} + std::to_string(tree_id) + ", Node " +
+                               node_id_repr + ": " + e.what()};
+    }
     // 0 indicates the lack of ID mapping for a particular node
     node_id_mapping_.push_back(static_cast<index_type>(tl_node_id.value_or(0)));
     ++cur_node_index_;
@@ -177,7 +182,7 @@ struct decision_forest_builder {
   void set_output_size(index_type val)
   {
     if (output_size_ != index_type{1} && output_size_ != val) {
-      throw model_import_error("Inconsistent leaf vector size");
+      throw unusable_model_exception("Inconsistent leaf vector size");
     }
     output_size_ = val;
   }
diff --git a/cpp/include/cuml/fil/detail/node.hpp b/cpp/include/cuml/fil/detail/node.hpp
index 6a0a3765d4..0a5bf88e2c 100644
--- a/cpp/include/cuml/fil/detail/node.hpp
+++ b/cpp/include/cuml/fil/detail/node.hpp
@@ -1,11 +1,12 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 #pragma once
 
 #include <cuml/fil/detail/index_type.hpp>
 #include <cuml/fil/detail/raft_proto/gpu_support.hpp>
+#include <cuml/fil/exceptions.hpp>
 #include <cuml/fil/tree_layout.hpp>
 
 #include <iostream>
@@ -104,14 +105,15 @@ struct alignas(detail::get_node_alignment<threshold_t, index_t, metadata_storage
 
   // TODO(wphicks): Add custom type to ensure given child offset is at least
   // one
-#pragma GCC diagnostic push
-#pragma GCC diagnostic ignored "-Wnarrowing"
-  HOST DEVICE constexpr node(threshold_type value             = threshold_type{},
-                             bool is_leaf_node                = true,
-                             bool default_to_distant_child    = false,
-                             bool is_categorical_node         = false,
-                             metadata_storage_type feature    = metadata_storage_type{},
-                             offset_type distant_child_offset = offset_type{})
+
+  // Assumption: Node construction occurs on the host. This allows us to perform
+  // bound check on the 'feature' parameter.
+  constexpr node(threshold_type value             = threshold_type{},
+                 bool is_leaf_node                = true,
+                 bool default_to_distant_child    = false,
+                 bool is_categorical_node         = false,
+                 metadata_storage_type feature    = metadata_storage_type{},
+                 offset_type distant_child_offset = offset_type{})
     : aligned_data{
         .inner_data = {
           {.value = value},
@@ -120,12 +122,12 @@ struct alignas(detail::get_node_alignment<threshold_t, index_t, metadata_storage
   {
   }
 
-  HOST DEVICE constexpr node(index_type index,
-                             bool is_leaf_node                = true,
-                             bool default_to_distant_child    = false,
-                             bool is_categorical_node         = false,
-                             metadata_storage_type feature    = metadata_storage_type{},
-                             offset_type distant_child_offset = offset_type{})
+  constexpr node(index_type index,
+                 bool is_leaf_node                = true,
+                 bool default_to_distant_child    = false,
+                 bool is_categorical_node         = false,
+                 metadata_storage_type feature    = metadata_storage_type{},
+                 offset_type distant_child_offset = offset_type{})
     : aligned_data{
         .inner_data = {
           {.index = index},
@@ -133,7 +135,6 @@ struct alignas(detail::get_node_alignment<threshold_t, index_t, metadata_storage
           construct_metadata(is_leaf_node, default_to_distant_child, is_categorical_node, feature)}}
   {
   }
-#pragma GCC diagnostic pop
 
   /* The index of the feature for this node */
   HOST DEVICE auto constexpr feature_index() const
@@ -213,6 +214,13 @@ struct alignas(detail::get_node_alignment<threshold_t, index_t, metadata_storage
                                            bool is_categorical_node      = false,
                                            metadata_storage_type feature = metadata_storage_type{})
   {
+    // Ensure that 'feature' is not truncated.
+    static_assert(std::is_unsigned_v<metadata_storage_type>, "metadata storage must be unsigned");
+    if (feature > FEATURE_MASK) {
+      throw model_import_error{std::string{"The 'feature' value in the node must be at most "} +
+                               std::to_string(FEATURE_MASK) + "."};
+    }
+
     return metadata_storage_type(
       (is_leaf_node << LEAF_BIT) + (default_to_distant_child << DEFAULT_DISTANT_BIT) +
       (is_categorical_node << CATEGORICAL_BIT) + (feature & FEATURE_MASK));
diff --git a/cpp/include/cuml/fil/exceptions.hpp b/cpp/include/cuml/fil/exceptions.hpp
index 8a8c2a00fd..39c00a43e3 100644
--- a/cpp/include/cuml/fil/exceptions.hpp
+++ b/cpp/include/cuml/fil/exceptions.hpp
@@ -23,11 +23,12 @@ struct unusable_model_exception : std::exception {
 /** Exception indicating model import failed */
 struct model_import_error : std::exception {
   model_import_error() : model_import_error("Error while importing model") {}
+  model_import_error(std::string msg) : msg_{msg} {}
   model_import_error(char const* msg) : msg_{msg} {}
-  virtual char const* what() const noexcept { return msg_; }
+  virtual char const* what() const noexcept { return msg_.c_str(); }
 
  private:
-  char const* msg_;
+  std::string msg_;
 };
 
 /**
diff --git a/cpp/include/cuml/fil/treelite_importer.hpp b/cpp/include/cuml/fil/treelite_importer.hpp
index 122c0ef268..6fb917136f 100644
--- a/cpp/include/cuml/fil/treelite_importer.hpp
+++ b/cpp/include/cuml/fil/treelite_importer.hpp
@@ -186,7 +186,7 @@ struct treelite_importer {
       result.constant = tl_model.sigmoid_alpha;
       result.element  = element_op::sigmoid;
     } else {
-      throw model_import_error{"Unrecognized Treelite pred_transform string"};
+      throw unusable_model_exception{"Unrecognized Treelite pred_transform string"};
     }
     return result;
   }
@@ -197,7 +197,7 @@ struct treelite_importer {
     switch (tl_model.GetThresholdType()) {
       case treelite::TypeInfo::kFloat64: result = true; break;
       case treelite::TypeInfo::kFloat32: result = false; break;
-      default: throw model_import_error("Unrecognized Treelite threshold type");
+      default: throw unusable_model_exception("Unrecognized Treelite threshold type");
     }
     return result;
   }
@@ -209,7 +209,7 @@ struct treelite_importer {
       case treelite::TypeInfo::kFloat64: result = true; break;
       case treelite::TypeInfo::kFloat32: result = false; break;
       case treelite::TypeInfo::kUInt32: result = false; break;
-      default: throw model_import_error("Unrecognized Treelite threshold type");
+      default: throw unusable_model_exception("Unrecognized Treelite threshold type");
     }
     return result;
   }
@@ -221,7 +221,7 @@ struct treelite_importer {
       case treelite::TypeInfo::kFloat64: result = false; break;
       case treelite::TypeInfo::kFloat32: result = false; break;
       case treelite::TypeInfo::kUInt32: result = true; break;
-      default: throw model_import_error("Unrecognized Treelite threshold type");
+      default: throw unusable_model_exception("Unrecognized Treelite threshold type");
     }
     return result;
   }
@@ -263,10 +263,13 @@ struct treelite_importer {
               builder.set_output_size(output.size());
               if (output.size() > index_type{1}) {
                 builder.add_leaf_vector_node(
-                  std::begin(output), std::end(output), node.get_treelite_id(), depth);
+                  std::begin(output), std::end(output), tree_id, node.get_treelite_id(), depth);
               } else {
-                builder.add_node(
-                  typename forest_model_t::io_type(output[0]), node.get_treelite_id(), depth, true);
+                builder.add_node(typename forest_model_t::io_type(output[0]),
+                                 tree_id,
+                                 node.get_treelite_id(),
+                                 depth,
+                                 true);
               }
             } else {
               if (node.is_categorical()) {
@@ -281,6 +284,7 @@ struct treelite_importer {
                                              offsets[node_index]);
               } else {
                 builder.add_node(typename forest_model_t::threshold_type(node.threshold()),
+                                 tree_id,
                                  node.get_treelite_id(),
                                  depth,
                                  false,
diff --git a/cpp/tests/sg/fil/treelite_importer_invalid_inputs.cpp b/cpp/tests/sg/fil/treelite_importer_invalid_inputs.cpp
index 4c7cc6f510..efc5dfba8c 100644
--- a/cpp/tests/sg/fil/treelite_importer_invalid_inputs.cpp
+++ b/cpp/tests/sg/fil/treelite_importer_invalid_inputs.cpp
@@ -3,7 +3,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-#include <cuml/fil/detail/decision_forest_builder.hpp>
+#include <cuml/fil/exceptions.hpp>
 #include <cuml/fil/tree_layout.hpp>
 #include <cuml/fil/treelite_importer.hpp>
 
@@ -64,9 +64,8 @@ TEST(TreeliteImporter, large_category_value)
   auto expected_error_msg = std::string{"Tree 0, Node 0: Category index must be at most "} +
                             std::to_string(std::numeric_limits<std::uint32_t>::max() - 1);
 
-  ASSERT_THAT(
-    [&]() { import_from_treelite_model(*tl_model, tree_layout::breadth_first); },
-    testing::ThrowsMessage<detail::model_builder_error>(testing::HasSubstr(expected_error_msg)));
+  ASSERT_THAT([&]() { import_from_treelite_model(*tl_model, tree_layout::breadth_first); },
+              testing::ThrowsMessage<model_import_error>(testing::HasSubstr(expected_error_msg)));
 }
 
 TEST(TreeliteImporter, large_category_value2)

From 7007801bb0878ef1310a8ef24851c5ab558e80e7 Mon Sep 17 00:00:00 2001
From: Hyunsu Cho <phcho@nvidia.com>
Date: Tue, 28 Apr 2026 19:04:37 -0700
Subject: [PATCH 11/17] Add gtest coverage for node bound check

---
 .../fil/treelite_importer_invalid_inputs.cpp  | 67 +++++++++++++++++++
 1 file changed, 67 insertions(+)

diff --git a/cpp/tests/sg/fil/treelite_importer_invalid_inputs.cpp b/cpp/tests/sg/fil/treelite_importer_invalid_inputs.cpp
index efc5dfba8c..160febc030 100644
--- a/cpp/tests/sg/fil/treelite_importer_invalid_inputs.cpp
+++ b/cpp/tests/sg/fil/treelite_importer_invalid_inputs.cpp
@@ -3,6 +3,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
+#include <cuml/fil/decision_forest.hpp>
 #include <cuml/fil/exceptions.hpp>
 #include <cuml/fil/tree_layout.hpp>
 #include <cuml/fil/treelite_importer.hpp>
@@ -110,5 +111,71 @@ TEST(TreeliteImporter, large_category_value2)
   ASSERT_NO_THROW(import_from_treelite_model(*tl_model, tree_layout::breadth_first));
 }
 
+TEST(TreeliteImporter, large_feature_id)
+{
+  // Tree models with 16-bit storage for node metadata should throw
+  // an exception for feature IDs larger than 0x1FFF.
+
+  auto metadata = treelite::model_builder::Metadata{
+    9000,
+    treelite::TaskType::kRegressor,
+    false,
+    1,
+    {1},
+    {1, 1},
+  };
+  auto tree_annotation = treelite::model_builder::TreeAnnotation{1, {0}, {0}};
+  auto model_builder =
+    treelite::model_builder::GetModelBuilder(treelite::TypeInfo::kFloat32,
+                                             treelite::TypeInfo::kFloat32,
+                                             metadata,
+                                             tree_annotation,
+                                             treelite::model_builder::PostProcessorFunc{"identity"},
+                                             {0.0});
+
+  model_builder->StartTree();
+  model_builder->StartNode(0);
+  // Use a "large" feature ID here
+  model_builder->NumericalTest(8999, 0.0, false, treelite::Operator::kGT, 1, 2);
+  model_builder->EndNode();
+
+  model_builder->StartNode(1);
+  model_builder->LeafScalar(1.0);
+  model_builder->EndNode();
+
+  model_builder->StartNode(2);
+  model_builder->LeafScalar(-1.0);
+  model_builder->EndNode();
+
+  model_builder->EndTree();
+
+  auto tl_model = model_builder->CommitModel();
+
+  // Normally, treelite_importer::import() would choose the right size
+  // for the metadata storage, sufficient to hold all given feature IDs.
+  // For this example, it chooses 32-bit storage type (due to the use of feature ID 8999).
+  ASSERT_NO_THROW(import_from_treelite_model(*tl_model, tree_layout::breadth_first));
+
+  // Trick the importer to pick 16-bit storage type for metadata storage.
+  auto variant_index = get_forest_variant_index(false, 2, 1);
+  auto importer      = treelite_importer<tree_layout::breadth_first>{};
+
+  // The importer should throw an informative error message rather than silently
+  // truncating the feature ID.
+  auto expected_error_msg =
+    std::string{"Tree 0, Node 0: The 'feature' value in the node must be at most "} +
+    std::to_string(0x1FFF);
+  ASSERT_THAT(
+    [&]() {
+      importer.import_to_specific_variant<index_type{}>(variant_index,
+                                                        *tl_model,
+                                                        importer.get_num_class(*tl_model),
+                                                        importer.get_num_feature(*tl_model),
+                                                        importer.get_max_num_categories(*tl_model),
+                                                        importer.get_offsets(*tl_model));
+    },
+    testing::ThrowsMessage<model_import_error>(testing::HasSubstr(expected_error_msg)));
+}
+
 }  // namespace fil
 }  // namespace ML

From 36315bfa6077b01c2f4fb798b7080b0bc990206d Mon Sep 17 00:00:00 2001
From: Hyunsu Cho <phcho@nvidia.com>
Date: Tue, 28 Apr 2026 19:29:06 -0700
Subject: [PATCH 12/17] Safe float casting in get_decision_forest()

---
 .../fil/detail/decision_forest_builder.hpp    | 68 +++++++++++++---
 cpp/include/cuml/fil/treelite_importer.hpp    |  2 +-
 .../fil/treelite_importer_invalid_inputs.cpp  | 81 +++++++++++++++++--
 3 files changed, 134 insertions(+), 17 deletions(-)

diff --git a/cpp/include/cuml/fil/detail/decision_forest_builder.hpp b/cpp/include/cuml/fil/detail/decision_forest_builder.hpp
index 000b0079dd..5a9787f6eb 100644
--- a/cpp/include/cuml/fil/detail/decision_forest_builder.hpp
+++ b/cpp/include/cuml/fil/detail/decision_forest_builder.hpp
@@ -21,7 +21,9 @@
 #include <iterator>
 #include <numeric>
 #include <optional>
+#include <sstream>
 #include <string>
+#include <type_traits>
 #include <utility>
 #include <vector>
 
@@ -29,6 +31,45 @@ namespace ML {
 namespace fil {
 namespace detail {
 
+struct floating_point_truncation_error : std::exception {
+  floating_point_truncation_error() {}
+  floating_point_truncation_error(std::string msg) : msg_{msg} {}
+  floating_point_truncation_error(char const* msg) : msg_{msg} {}
+  virtual char const* what() const noexcept { return msg_.c_str(); }
+
+ private:
+  std::string msg_;
+};
+
+template <typename To, typename From>
+To safe_cast_floating_point(From x)
+{
+  static_assert(std::is_floating_point_v<From> && std::is_floating_point_v<To>,
+                "Source and destination types must be both floating-point types.");
+  if constexpr (sizeof(To) >= sizeof(From)) {
+    // Widening cast
+    return static_cast<To>(x);
+  } else {
+    // Narrowing cast: should be checked
+    if (!std::isfinite(x)) {
+      throw floating_point_truncation_error{"Cannot cast an INF or NaN value"};
+    }
+    auto constexpr lower_limit = static_cast<From>(std::numeric_limits<To>::lowest());
+    auto constexpr upper_limit = static_cast<From>(std::numeric_limits<To>::max());
+    if (x < lower_limit) {
+      std::ostringstream ss;
+      ss << "Input must be at least " << lower_limit << ".";
+      throw floating_point_truncation_error{ss.str()};
+    }
+    if (x > upper_limit) {
+      std::ostringstream ss;
+      ss << "Input must be at most " << upper_limit << ".";
+      throw floating_point_truncation_error{ss.str()};
+    }
+    return static_cast<To>(x);
+  }
+}
+
 /*
  * Struct used to build FIL forests
  */
@@ -165,7 +206,7 @@ struct decision_forest_builder {
   /* Set the row-wise postprocessing operation for this model */
   void set_row_postproc(row_op val) { row_postproc_ = val; }
   /* Set the value to divide by during postprocessing */
-  void set_average_factor(double val) { average_factor_ = val; }
+  void set_average_factor(float val) { average_factor_ = val; }
   /* Set the bias term, which is added to the output. The bias term
    * should have the same length as output_size. */
   void set_bias(std::vector<double> val)
@@ -211,11 +252,21 @@ struct decision_forest_builder {
                            int device                       = 0,
                            raft_proto::cuda_stream stream   = raft_proto::cuda_stream{})
   {
-    // Allow narrowing for preprocessing constants. They are stored as doubles
-    // for consistency in the builder but must be converted to the proper types
-    // for the concrete forest model.
-#pragma GCC diagnostic push
-#pragma GCC diagnostic ignored "-Wnarrowing"
+    // Invariant, so that static_cast<typename node_type::threshold_type>(average_factor_) is safe.
+    static_assert(sizeof(typename node_type::threshold_type) >= sizeof(average_factor_),
+                  "Threshold type was assumed to be big enough to hold average factor");
+
+    // Safely cast postproc_constant_ to typename node_type::threshold_type
+    auto postproc_constant_casted = typename node_type::threshold_type{};
+    try {
+      postproc_constant_casted =
+        safe_cast_floating_point<typename node_type::threshold_type>(postproc_constant_);
+      // We can't use cuda::narrow here, because it throws for imprecise conversion, i.e. casting
+      // double{3.1} to float.
+    } catch (const floating_point_truncation_error& e) {
+      throw unusable_model_exception{
+        std::string{"Found an invalid value for postprocessing constant: "} + e.what()};
+    }
     return decision_forest_t{
       raft_proto::buffer{
         raft_proto::buffer{nodes_.data(), nodes_.size()}, mem_type, device, stream},
@@ -249,8 +300,7 @@ struct decision_forest_builder {
       row_postproc_,
       element_postproc_,
       static_cast<typename node_type::threshold_type>(average_factor_),
-      static_cast<typename node_type::threshold_type>(postproc_constant_)};
-#pragma GCC diagnostic pop
+      postproc_constant_casted};
   }
 
  private:
@@ -260,7 +310,7 @@ struct decision_forest_builder {
   index_type output_size_;
   row_op row_postproc_;
   element_op element_postproc_;
-  double average_factor_;
+  float average_factor_;
   double postproc_constant_;
 
   std::vector<node_type> nodes_;
diff --git a/cpp/include/cuml/fil/treelite_importer.hpp b/cpp/include/cuml/fil/treelite_importer.hpp
index 6fb917136f..2c0f051f3e 100644
--- a/cpp/include/cuml/fil/treelite_importer.hpp
+++ b/cpp/include/cuml/fil/treelite_importer.hpp
@@ -151,7 +151,7 @@ struct treelite_importer {
     } else {
       result = 1.0;
     }
-    return result;
+    return static_cast<float>(result);
   }
 
   auto get_bias(treelite::Model const& tl_model) { return tl_model.base_scores.AsVector(); }
diff --git a/cpp/tests/sg/fil/treelite_importer_invalid_inputs.cpp b/cpp/tests/sg/fil/treelite_importer_invalid_inputs.cpp
index 160febc030..aa0bae0661 100644
--- a/cpp/tests/sg/fil/treelite_importer_invalid_inputs.cpp
+++ b/cpp/tests/sg/fil/treelite_importer_invalid_inputs.cpp
@@ -51,11 +51,11 @@ TEST(TreeliteImporter, large_category_value)
   model_builder->EndNode();
 
   model_builder->StartNode(1);
-  model_builder->LeafScalar(1.0);
+  model_builder->LeafScalar(1.0f);
   model_builder->EndNode();
 
   model_builder->StartNode(2);
-  model_builder->LeafScalar(-1.0);
+  model_builder->LeafScalar(-1.0f);
   model_builder->EndNode();
 
   model_builder->EndTree();
@@ -98,11 +98,11 @@ TEST(TreeliteImporter, large_category_value2)
   model_builder->EndNode();
 
   model_builder->StartNode(1);
-  model_builder->LeafScalar(1.0);
+  model_builder->LeafScalar(1.0f);
   model_builder->EndNode();
 
   model_builder->StartNode(2);
-  model_builder->LeafScalar(-1.0);
+  model_builder->LeafScalar(-1.0f);
   model_builder->EndNode();
 
   model_builder->EndTree();
@@ -140,11 +140,11 @@ TEST(TreeliteImporter, large_feature_id)
   model_builder->EndNode();
 
   model_builder->StartNode(1);
-  model_builder->LeafScalar(1.0);
+  model_builder->LeafScalar(1.0f);
   model_builder->EndNode();
 
   model_builder->StartNode(2);
-  model_builder->LeafScalar(-1.0);
+  model_builder->LeafScalar(-1.0f);
   model_builder->EndNode();
 
   model_builder->EndTree();
@@ -166,7 +166,7 @@ TEST(TreeliteImporter, large_feature_id)
     std::string{"Tree 0, Node 0: The 'feature' value in the node must be at most "} +
     std::to_string(0x1FFF);
   ASSERT_THAT(
-    [&]() {
+    [&] {
       importer.import_to_specific_variant<index_type{}>(variant_index,
                                                         *tl_model,
                                                         importer.get_num_class(*tl_model),
@@ -177,5 +177,72 @@ TEST(TreeliteImporter, large_feature_id)
     testing::ThrowsMessage<model_import_error>(testing::HasSubstr(expected_error_msg)));
 }
 
+TEST(TreeliteImporter, safe_cast_floating_point)
+{
+  /* Valid casts */
+  ASSERT_NO_THROW(
+    detail::safe_cast_floating_point<float>(double{3.1}));  // Some loss of precision, but o.k.
+  ASSERT_NO_THROW(detail::safe_cast_floating_point<double>(std::numeric_limits<float>::max()));
+
+  // INFs and NANs are allowed for widening cast
+  ASSERT_NO_THROW(detail::safe_cast_floating_point<float>(std::numeric_limits<float>::infinity()));
+  ASSERT_NO_THROW(
+    detail::safe_cast_floating_point<double>(std::numeric_limits<double>::infinity()));
+  ASSERT_NO_THROW(detail::safe_cast_floating_point<double>(std::numeric_limits<float>::infinity()));
+  ASSERT_NO_THROW(detail::safe_cast_floating_point<float>(std::numeric_limits<float>::quiet_NaN()));
+  ASSERT_NO_THROW(
+    detail::safe_cast_floating_point<double>(std::numeric_limits<double>::quiet_NaN()));
+  ASSERT_NO_THROW(
+    detail::safe_cast_floating_point<double>(std::numeric_limits<float>::quiet_NaN()));
+
+  // Invalid casts
+  auto inf_msg = std::string{"Cannot cast an INF or NaN value"};
+  ASSERT_THAT(
+    [] { detail::safe_cast_floating_point<float>(std::numeric_limits<double>::infinity()); },
+    testing::ThrowsMessage<detail::floating_point_truncation_error>(testing::HasSubstr(inf_msg)));
+  ASSERT_THAT(
+    [] { detail::safe_cast_floating_point<float>(std::numeric_limits<double>::quiet_NaN()); },
+    testing::ThrowsMessage<detail::floating_point_truncation_error>(testing::HasSubstr(inf_msg)));
+  ASSERT_THAT([] { detail::safe_cast_floating_point<float>(double{1e100}); },
+              testing::ThrowsMessage<detail::floating_point_truncation_error>(
+                testing::HasSubstr("Input must be at most")));
+  ASSERT_THAT([] { detail::safe_cast_floating_point<float>(double{-1e100}); },
+              testing::ThrowsMessage<detail::floating_point_truncation_error>(
+                testing::HasSubstr("Input must be at least")));
+}
+
+TEST(TreeliteImporter, invalid_postproc_constant)
+{
+  auto metadata = treelite::model_builder::Metadata{
+    1,
+    treelite::TaskType::kRegressor,
+    false,
+    1,
+    {1},
+    {1, 1},
+  };
+  auto tree_annotation = treelite::model_builder::TreeAnnotation{1, {0}, {0}};
+  auto model_builder   = treelite::model_builder::GetModelBuilder(
+    treelite::TypeInfo::kFloat32,
+    treelite::TypeInfo::kFloat32,
+    metadata,
+    tree_annotation,
+    treelite::model_builder::PostProcessorFunc{
+      "sigmoid", {{"sigmoid_alpha", std::numeric_limits<double>::quiet_NaN()}}},
+    {0.0});
+
+  model_builder->StartTree();
+  model_builder->StartNode(0);
+  model_builder->LeafScalar(0.0f);
+  model_builder->EndNode();
+  model_builder->EndTree();
+
+  auto tl_model = model_builder->CommitModel();
+
+  ASSERT_THAT([&] { import_from_treelite_model(*tl_model, tree_layout::breadth_first); },
+              testing::ThrowsMessage<unusable_model_exception>(
+                testing::HasSubstr("Found an invalid value for postprocessing constant")));
+}
+
 }  // namespace fil
 }  // namespace ML

From d20d069d490a6dddd2b057a5a504dd16dc2c5b81 Mon Sep 17 00:00:00 2001
From: Hyunsu Cho <phcho@nvidia.com>
Date: Tue, 28 Apr 2026 20:58:21 -0700
Subject: [PATCH 13/17] Use 64-bit for average factor, if possible

---
 .../fil/detail/decision_forest_builder.hpp    | 24 +++++++++++--------
 cpp/include/cuml/fil/treelite_importer.hpp    |  2 +-
 2 files changed, 15 insertions(+), 11 deletions(-)

diff --git a/cpp/include/cuml/fil/detail/decision_forest_builder.hpp b/cpp/include/cuml/fil/detail/decision_forest_builder.hpp
index 5a9787f6eb..6d3329030e 100644
--- a/cpp/include/cuml/fil/detail/decision_forest_builder.hpp
+++ b/cpp/include/cuml/fil/detail/decision_forest_builder.hpp
@@ -206,7 +206,7 @@ struct decision_forest_builder {
   /* Set the row-wise postprocessing operation for this model */
   void set_row_postproc(row_op val) { row_postproc_ = val; }
   /* Set the value to divide by during postprocessing */
-  void set_average_factor(float val) { average_factor_ = val; }
+  void set_average_factor(double val) { average_factor_ = val; }
   /* Set the bias term, which is added to the output. The bias term
    * should have the same length as output_size. */
   void set_bias(std::vector<double> val)
@@ -252,17 +252,21 @@ struct decision_forest_builder {
                            int device                       = 0,
                            raft_proto::cuda_stream stream   = raft_proto::cuda_stream{})
   {
-    // Invariant, so that static_cast<typename node_type::threshold_type>(average_factor_) is safe.
-    static_assert(sizeof(typename node_type::threshold_type) >= sizeof(average_factor_),
-                  "Threshold type was assumed to be big enough to hold average factor");
-
-    // Safely cast postproc_constant_ to typename node_type::threshold_type
+    // Safely cast average_factor_ and postproc_constant_ to node_type::threshold_type
+    auto average_factor_casted    = typename node_type::threshold_type{};
     auto postproc_constant_casted = typename node_type::threshold_type{};
     try {
-      postproc_constant_casted =
-        safe_cast_floating_point<typename node_type::threshold_type>(postproc_constant_);
+      average_factor_casted =
+        safe_cast_floating_point<typename node_type::threshold_type>(average_factor_);
       // We can't use cuda::narrow here, because it throws for imprecise conversion, i.e. casting
       // double{3.1} to float.
+    } catch (const floating_point_truncation_error& e) {
+      throw unusable_model_exception{std::string{"Found an invalid value for averaging factor: "} +
+                                     e.what()};
+    }
+    try {
+      postproc_constant_casted =
+        safe_cast_floating_point<typename node_type::threshold_type>(postproc_constant_);
     } catch (const floating_point_truncation_error& e) {
       throw unusable_model_exception{
         std::string{"Found an invalid value for postprocessing constant: "} + e.what()};
@@ -299,7 +303,7 @@ struct decision_forest_builder {
       output_size_,
       row_postproc_,
       element_postproc_,
-      static_cast<typename node_type::threshold_type>(average_factor_),
+      average_factor_casted,
       postproc_constant_casted};
   }
 
@@ -310,7 +314,7 @@ struct decision_forest_builder {
   index_type output_size_;
   row_op row_postproc_;
   element_op element_postproc_;
-  float average_factor_;
+  double average_factor_;
   double postproc_constant_;
 
   std::vector<node_type> nodes_;
diff --git a/cpp/include/cuml/fil/treelite_importer.hpp b/cpp/include/cuml/fil/treelite_importer.hpp
index 2c0f051f3e..6fb917136f 100644
--- a/cpp/include/cuml/fil/treelite_importer.hpp
+++ b/cpp/include/cuml/fil/treelite_importer.hpp
@@ -151,7 +151,7 @@ struct treelite_importer {
     } else {
       result = 1.0;
     }
-    return static_cast<float>(result);
+    return result;
   }
 
   auto get_bias(treelite::Model const& tl_model) { return tl_model.base_scores.AsVector(); }

From d5cbde8f8e915f37e044ae899d600392a191a34a Mon Sep 17 00:00:00 2001
From: Simon Adorf <sadorf@nvidia.com>
Date: Wed, 29 Apr 2026 15:18:58 +0000
Subject: [PATCH 14/17] Enhance error handling in treelite_importer and
 simplify decision_forest_builder methods

- Added try-catch block in treelite_importer to provide detailed error messages for model import errors.
- Removed unnecessary tree_id parameter from add_categorical_node and add_leaf_vector_node methods in decision_forest_builder.
- Streamlined node addition logic in decision_forest_builder to improve clarity and maintainability.
---
 .../fil/detail/decision_forest_builder.hpp    | 34 ++-------
 cpp/include/cuml/fil/treelite_importer.hpp    | 70 ++++++++++---------
 2 files changed, 42 insertions(+), 62 deletions(-)

diff --git a/cpp/include/cuml/fil/detail/decision_forest_builder.hpp b/cpp/include/cuml/fil/detail/decision_forest_builder.hpp
index 6d3329030e..59829dc417 100644
--- a/cpp/include/cuml/fil/detail/decision_forest_builder.hpp
+++ b/cpp/include/cuml/fil/detail/decision_forest_builder.hpp
@@ -83,7 +83,6 @@ struct decision_forest_builder {
   void add_categorical_node(
     iter_t vec_begin,
     iter_t vec_end,
-    std::size_t tree_id,
     std::optional<int> tl_node_id                     = std::nullopt,
     std::size_t depth                                 = std::size_t{1},
     bool default_to_distant_child                     = false,
@@ -106,10 +105,7 @@ struct decision_forest_builder {
     auto max_cat = (vec_begin != vec_end) ? *std::max_element(vec_begin, vec_end) : cat_t{0};
     if constexpr (std::is_same_v<cat_t, index_t>) {
       if (max_cat == std::numeric_limits<index_t>::max()) {
-        auto node_id_repr =
-          tl_node_id.has_value() ? std::to_string(tl_node_id.value()) : std::string{"n/a"};
-        throw model_import_error{std::string{"Tree "} + std::to_string(tree_id) + ", Node " +
-                                 node_id_repr + ": Category index must be at most " +
+        throw model_import_error{std::string{"Category index must be at most "} +
                                  std::to_string(std::numeric_limits<index_t>::max() - 1)};
       }
     }
@@ -125,23 +121,14 @@ struct decision_forest_builder {
     auto set = bitset{set_storage, max_cat_plus_one};
     std::for_each(vec_begin, vec_end, [&set](auto&& cat_index) { set.set(cat_index); });
 
-    add_node(node_value,
-             tree_id,
-             tl_node_id,
-             depth,
-             false,
-             default_to_distant_child,
-             true,
-             feature,
-             offset,
-             false);
+    add_node(
+      node_value, tl_node_id, depth, false, default_to_distant_child, true, feature, offset, false);
   }
 
   /* Add a leaf node with vector output */
   template <typename iter_t>
   void add_leaf_vector_node(iter_t vec_begin,
                             iter_t vec_end,
-                            std::size_t tree_id,
                             std::optional<int> tl_node_id = std::nullopt,
                             std::size_t depth             = std::size_t{1})
   {
@@ -149,7 +136,6 @@ struct decision_forest_builder {
     std::copy(vec_begin, vec_end, std::back_inserter(vector_output_));
 
     add_node(leaf_index,
-             tree_id,
              tl_node_id,
              depth,
              true,
@@ -164,7 +150,6 @@ struct decision_forest_builder {
   template <typename value_t>
   void add_node(
     value_t val,
-    std::size_t tree_id,
     std::optional<int> tl_node_id                     = std::nullopt,
     std::size_t depth                                 = std::size_t{1},
     bool is_leaf_node                                 = true,
@@ -179,7 +164,7 @@ struct decision_forest_builder {
         if (cur_node_index_ % alignment_ != index_type{}) {
           auto padding = (alignment_ - cur_node_index_ % alignment_);
           for (auto i = index_type{}; i < padding; ++i) {
-            add_node(typename node_type::threshold_type{}, tree_id, std::nullopt);
+            add_node(typename node_type::threshold_type{}, std::nullopt);
           }
         }
       }
@@ -187,15 +172,8 @@ struct decision_forest_builder {
     }
 
     if (is_inclusive) { val = std::nextafter(val, std::numeric_limits<value_t>::infinity()); }
-    try {
-      nodes_.emplace_back(
-        val, is_leaf_node, default_to_distant_child, is_categorical_node, feature, offset);
-    } catch (const model_import_error& e) {
-      auto node_id_repr =
-        tl_node_id.has_value() ? std::to_string(tl_node_id.value()) : std::string{"n/a"};
-      throw model_import_error{std::string{"Tree "} + std::to_string(tree_id) + ", Node " +
-                               node_id_repr + ": " + e.what()};
-    }
+    nodes_.emplace_back(
+      val, is_leaf_node, default_to_distant_child, is_categorical_node, feature, offset);
     // 0 indicates the lack of ID mapping for a particular node
     node_id_mapping_.push_back(static_cast<index_type>(tl_node_id.value_or(0)));
     ++cur_node_index_;
diff --git a/cpp/include/cuml/fil/treelite_importer.hpp b/cpp/include/cuml/fil/treelite_importer.hpp
index 6fb917136f..853fc5d39b 100644
--- a/cpp/include/cuml/fil/treelite_importer.hpp
+++ b/cpp/include/cuml/fil/treelite_importer.hpp
@@ -258,42 +258,44 @@ struct treelite_importer {
           tl_model,
           [&builder, &offsets, &node_index](
             auto&& tree_id, auto&& node, auto&& depth, auto&& parent_index) {
-            if (node.is_leaf()) {
-              auto output = node.get_output();
-              builder.set_output_size(output.size());
-              if (output.size() > index_type{1}) {
-                builder.add_leaf_vector_node(
-                  std::begin(output), std::end(output), tree_id, node.get_treelite_id(), depth);
+            try {
+              if (node.is_leaf()) {
+                auto output = node.get_output();
+                builder.set_output_size(output.size());
+                if (output.size() > index_type{1}) {
+                  builder.add_leaf_vector_node(
+                    std::begin(output), std::end(output), node.get_treelite_id(), depth);
+                } else {
+                  builder.add_node(typename forest_model_t::io_type(output[0]),
+                                   node.get_treelite_id(),
+                                   depth,
+                                   true);
+                }
               } else {
-                builder.add_node(typename forest_model_t::io_type(output[0]),
-                                 tree_id,
-                                 node.get_treelite_id(),
-                                 depth,
-                                 true);
-              }
-            } else {
-              if (node.is_categorical()) {
-                auto categories = node.get_categories();
-                builder.add_categorical_node(std::begin(categories),
-                                             std::end(categories),
-                                             tree_id,
-                                             node.get_treelite_id(),
-                                             depth,
-                                             node.default_distant(),
-                                             node.get_feature(),
-                                             offsets[node_index]);
-              } else {
-                builder.add_node(typename forest_model_t::threshold_type(node.threshold()),
-                                 tree_id,
-                                 node.get_treelite_id(),
-                                 depth,
-                                 false,
-                                 node.default_distant(),
-                                 false,
-                                 node.get_feature(),
-                                 offsets[node_index],
-                                 node.is_inclusive());
+                if (node.is_categorical()) {
+                  auto categories = node.get_categories();
+                  builder.add_categorical_node(std::begin(categories),
+                                               std::end(categories),
+                                               node.get_treelite_id(),
+                                               depth,
+                                               node.default_distant(),
+                                               node.get_feature(),
+                                               offsets[node_index]);
+                } else {
+                  builder.add_node(typename forest_model_t::threshold_type(node.threshold()),
+                                   node.get_treelite_id(),
+                                   depth,
+                                   false,
+                                   node.default_distant(),
+                                   false,
+                                   node.get_feature(),
+                                   offsets[node_index],
+                                   node.is_inclusive());
+                }
               }
+            } catch (const model_import_error& e) {
+              throw model_import_error{std::string{"Tree "} + std::to_string(tree_id) + ", Node " +
+                                       std::to_string(node.get_treelite_id()) + ": " + e.what()};
             }
             ++node_index;
           });

From afce4e18a1a75f6bf9c91dead866bc707612d9ea Mon Sep 17 00:00:00 2001
From: Simon Adorf <sadorf@nvidia.com>
Date: Wed, 29 Apr 2026 15:54:10 +0000
Subject: [PATCH 15/17] Revert "Add bound check for tree_index"

This reverts commit 99d13f69c8b15b382fee24cfbf96b14b26376181.
---
 .../cuml/fil/detail/infer_kernel/gpu.cuh      | 95 +++++++++----------
 1 file changed, 46 insertions(+), 49 deletions(-)

diff --git a/cpp/include/cuml/fil/detail/infer_kernel/gpu.cuh b/cpp/include/cuml/fil/detail/infer_kernel/gpu.cuh
index e09dd2ad7c..95801160da 100644
--- a/cpp/include/cuml/fil/detail/infer_kernel/gpu.cuh
+++ b/cpp/include/cuml/fil/detail/infer_kernel/gpu.cuh
@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 #pragma once
@@ -125,59 +125,56 @@ CUML_KERNEL void __launch_bounds__(MAX_THREADS_PER_BLOCK, MIN_BLOCKS_PER_SM) inf
     // work within the loop if the task_index is below the actual task_count.
     auto const task_count_rounded_up = blockDim.x * raft_proto::ceildiv(task_count, blockDim.x);
 
-    // Ensure that tree_index doesn't overflow in the loop below
-    auto max_tree_index = (task_count - 1) / chunk_size;
-    if (max_tree_index < forest.tree_count()) {
-      // Infer on each tree and row
-      for (auto task_index = threadIdx.x; task_index < task_count_rounded_up;
-           task_index += blockDim.x) {
-        auto row_index = task_index % chunk_size;
-        auto real_task = task_index < task_count && row_index < rows_in_this_iteration;
-        row_index *= real_task;
-        auto tree_index  = task_index * real_task / chunk_size;
-        auto grove_index = (threadIdx.x / chunk_size) * (infer_type == infer_kind::default_kind);
-
-        auto tree_output  = std::conditional_t<has_vector_leaves,
-                                               typename node_t::index_type,
-                                               typename node_t::threshold_type>{};
-        auto leaf_node_id = index_type{};
-        if (infer_type == infer_kind::leaf_id) {
-          leaf_node_id =
-            evaluate_tree<has_vector_leaves, has_categorical_nodes, has_nonlocal_categories, true>(
-              forest, tree_index, input_data + row_index * col_count, categorical_data);
-        } else {
-          tree_output =
-            evaluate_tree<has_vector_leaves, has_categorical_nodes, has_nonlocal_categories, false>(
-              forest, tree_index, input_data + row_index * col_count, categorical_data);
-        }
+    // Infer on each tree and row
+    for (auto task_index = threadIdx.x; task_index < task_count_rounded_up;
+         task_index += blockDim.x) {
+      auto row_index = task_index % chunk_size;
+      auto real_task = task_index < task_count && row_index < rows_in_this_iteration;
+      row_index *= real_task;
+      auto tree_index  = task_index * real_task / chunk_size;
+      auto grove_index = (threadIdx.x / chunk_size) * (infer_type == infer_kind::default_kind);
+
+      auto tree_output  = std::conditional_t<has_vector_leaves,
+                                             typename node_t::index_type,
+                                             typename node_t::threshold_type>{};
+      auto leaf_node_id = index_type{};
+      if (infer_type == infer_kind::leaf_id) {
+        leaf_node_id =
+          evaluate_tree<has_vector_leaves, has_categorical_nodes, has_nonlocal_categories, true>(
+            forest, tree_index, input_data + row_index * col_count, categorical_data);
+      } else {
+        tree_output =
+          evaluate_tree<has_vector_leaves, has_categorical_nodes, has_nonlocal_categories, false>(
+            forest, tree_index, input_data + row_index * col_count, categorical_data);
+      }
 
-        if (infer_type == infer_kind::leaf_id) {
-          output_workspace[row_index * num_outputs * num_grove + tree_index * num_grove +
-                           grove_index] = static_cast<typename forest_t::io_type>(leaf_node_id);
-        } else {
-          if constexpr (has_vector_leaves) {
-            auto output_offset =
-              (row_index * num_outputs * num_grove +
-               tree_index * default_num_outputs * num_grove * (infer_type == infer_kind::per_tree) +
-               grove_index);
-            for (auto output_index = index_type{}; output_index < default_num_outputs;
-                 ++output_index) {
-              if (real_task) {
-                output_workspace[output_offset + output_index * num_grove] +=
-                  vector_output_p[tree_output * default_num_outputs + output_index];
-              }
+      if (infer_type == infer_kind::leaf_id) {
+        output_workspace[row_index * num_outputs * num_grove + tree_index * num_grove +
+                         grove_index] = static_cast<typename forest_t::io_type>(leaf_node_id);
+      } else {
+        if constexpr (has_vector_leaves) {
+          auto output_offset =
+            (row_index * num_outputs * num_grove +
+             tree_index * default_num_outputs * num_grove * (infer_type == infer_kind::per_tree) +
+             grove_index);
+          for (auto output_index = index_type{}; output_index < default_num_outputs;
+               ++output_index) {
+            if (real_task) {
+              output_workspace[output_offset + output_index * num_grove] +=
+                vector_output_p[tree_output * default_num_outputs + output_index];
             }
-          } else {
-            auto output_offset =
-              (row_index * num_outputs * num_grove +
-               (tree_index % default_num_outputs) * num_grove *
-                 (infer_type == infer_kind::default_kind) +
-               tree_index * num_grove * (infer_type == infer_kind::per_tree) + grove_index);
-            if (real_task) { output_workspace[output_offset] += tree_output; }
           }
+        } else {
+          auto output_offset =
+            (row_index * num_outputs * num_grove +
+             (tree_index % default_num_outputs) * num_grove *
+               (infer_type == infer_kind::default_kind) +
+             tree_index * num_grove * (infer_type == infer_kind::per_tree) + grove_index);
+          if (real_task) { output_workspace[output_offset] += tree_output; }
         }
-        __syncthreads();
       }
+
+      __syncthreads();
     }
 
     auto padded_num_groves = raft_proto::padded_size(num_grove, WARP_SIZE);

From ba786344c5e5d8df43fbcda042d1788c201610d8 Mon Sep 17 00:00:00 2001
From: Simon Adorf <sadorf@nvidia.com>
Date: Wed, 29 Apr 2026 17:30:21 +0000
Subject: [PATCH 16/17] Validate treelite tree before building decision forest.

---
 .../fil/detail/decision_forest_builder.hpp    | 53 +++++++++++++++++++
 1 file changed, 53 insertions(+)

diff --git a/cpp/include/cuml/fil/detail/decision_forest_builder.hpp b/cpp/include/cuml/fil/detail/decision_forest_builder.hpp
index 59829dc417..488e1be893 100644
--- a/cpp/include/cuml/fil/detail/decision_forest_builder.hpp
+++ b/cpp/include/cuml/fil/detail/decision_forest_builder.hpp
@@ -230,6 +230,59 @@ struct decision_forest_builder {
                            int device                       = 0,
                            raft_proto::cuda_stream stream   = raft_proto::cuda_stream{})
   {
+    // Validate forest invariants the inference kernel relies on. After this
+    // function returns, the forest is treated as trusted by the kernel.
+
+    // tree_index arithmetic in the kernel uses index_type, so the tree count
+    // must fit without narrowing.
+    if (root_node_indexes_.size() > std::numeric_limits<index_type>::max()) {
+      throw model_import_error{std::string{"Forest has "} +
+                               std::to_string(root_node_indexes_.size()) +
+                               " trees, which exceeds the maximum representable in index_type (" +
+                               std::to_string(std::numeric_limits<index_type>::max()) + ")"};
+    }
+
+    // forest::get_tree_root(tree_index) dereferences nodes_ + root_index.
+    // Ensure each root index points into the nodes buffer.
+    for (auto i = std::size_t{0}; i < root_node_indexes_.size(); ++i) {
+      if (root_node_indexes_[i] >= nodes_.size()) {
+        throw model_import_error{
+          std::string{"Tree "} + std::to_string(i) + ": root node index out of bounds (" +
+          std::to_string(root_node_indexes_[i]) + " >= " + std::to_string(nodes_.size()) + ")"};
+      }
+    }
+
+    auto constexpr const cat_bin_width =
+      typename node_type::index_type{sizeof(typename node_type::index_type) * 8};
+    if (max_num_categories_ > cat_bin_width) {
+      auto const storage_size = categorical_storage_.size();
+      for (auto i = std::size_t{0}; i < nodes_.size(); ++i) {
+        auto const& n = nodes_[i];
+        if (n.is_leaf() || !n.is_categorical()) { continue; }
+        auto const offset = n.index();
+
+        // evaluate_tree_impl() reads categorical_storage[offset] as the number
+        // of categories for this node; offset must be in-range.
+        if (offset >= storage_size) {
+          throw model_import_error{std::string{"Categorical node "} + std::to_string(i) +
+                                   ": storage offset out of bounds (" + std::to_string(offset) +
+                                   " >= " + std::to_string(storage_size) + ")"};
+        }
+        auto const stored_num_cats = categorical_storage_[offset];
+        auto const bins_required   = raft_proto::ceildiv(stored_num_cats, cat_bin_width);
+
+        // evaluate_tree_impl() reconstructs a bitset from
+        // [offset + 1, offset + 1 + bins_required). Compute this range using
+        // size_t to keep the arithmetic explicit and overflow-safe.
+        auto const bits_begin = static_cast<std::size_t>(offset) + std::size_t{1};
+        auto const bits_end   = bits_begin + static_cast<std::size_t>(bins_required);
+        if (bits_end > storage_size) {
+          throw model_import_error{std::string{"Categorical node "} + std::to_string(i) +
+                                   ": bitset extends past categorical_storage end"};
+        }
+      }
+    }
+
     // Safely cast average_factor_ and postproc_constant_ to node_type::threshold_type
     auto average_factor_casted    = typename node_type::threshold_type{};
     auto postproc_constant_casted = typename node_type::threshold_type{};

From 02e8e23b69ff2ef17ec54ecd63d8fb9cfc148bfa Mon Sep 17 00:00:00 2001
From: Simon Adorf <sadorf@nvidia.com>
Date: Wed, 29 Apr 2026 18:20:16 +0000
Subject: [PATCH 17/17] Add tests for decision forest builder to validate error
 handling

- Introduced a new test file for decision forest builder to check for out-of-bounds errors in categorical storage offsets and bitset extents.
- Updated CMakeLists.txt to include the new test file in the build configuration.
---
 cpp/tests/CMakeLists.txt                      |  6 +-
 ...decision_forest_builder_invalid_inputs.cpp | 82 +++++++++++++++++++
 2 files changed, 86 insertions(+), 2 deletions(-)
 create mode 100644 cpp/tests/sg/fil/decision_forest_builder_invalid_inputs.cpp

diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt
index 1f8844d008..89061b14b9 100644
--- a/cpp/tests/CMakeLists.txt
+++ b/cpp/tests/CMakeLists.txt
@@ -110,8 +110,10 @@ if(all_algo OR fil_algo)
   ConfigureTest(PREFIX SG NAME FOREST_TRAVERSAL_TEST sg/forest/traversal_forest.cpp ML_INCLUDE)
   ConfigureTest(PREFIX SG NAME TREELITE_TRAVERSAL_TEST sg/forest/treelite_traversal.cpp ML_INCLUDE)
   ConfigureTest(
-    PREFIX SG NAME TREELITE_IMPORTER_TEST sg/fil/treelite_importer.cpp
-                                          sg/fil/treelite_importer_invalid_inputs.cpp ML_INCLUDE
+    PREFIX SG
+    NAME TREELITE_IMPORTER_TEST
+         sg/fil/treelite_importer.cpp sg/fil/treelite_importer_invalid_inputs.cpp
+         sg/fil/decision_forest_builder_invalid_inputs.cpp ML_INCLUDE
   )
 endif()
 
diff --git a/cpp/tests/sg/fil/decision_forest_builder_invalid_inputs.cpp b/cpp/tests/sg/fil/decision_forest_builder_invalid_inputs.cpp
new file mode 100644
index 0000000000..5e508573a9
--- /dev/null
+++ b/cpp/tests/sg/fil/decision_forest_builder_invalid_inputs.cpp
@@ -0,0 +1,82 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#include <cuml/fil/decision_forest.hpp>
+#include <cuml/fil/detail/decision_forest_builder.hpp>
+#include <cuml/fil/exceptions.hpp>
+#include <cuml/fil/tree_layout.hpp>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+
+#include <array>
+#include <cstdint>
+
+namespace ML {
+namespace fil {
+namespace detail {
+
+using test_forest_t =
+  decision_forest<tree_layout::breadth_first, float, std::uint32_t, std::uint16_t, std::uint16_t>;
+
+TEST(DecisionForestBuilder, CategoricalStorageOffsetOutOfBounds)
+{
+  auto builder = decision_forest_builder<test_forest_t>(
+    /*max_num_categories=*/std::uint32_t{33}, /*align_bytes=*/std::uint32_t{0});
+
+  // Construct a malformed categorical node that references non-local category
+  // storage at an out-of-range offset. This should be rejected by the invariant
+  // checks in get_decision_forest().
+  builder.add_node(std::uint32_t{1234},
+                   /*tl_node_id=*/0,
+                   /*depth=*/0,
+                   /*is_leaf_node=*/false,
+                   /*default_to_distant_child=*/false,
+                   /*is_categorical_node=*/true,
+                   /*feature=*/0,
+                   /*offset=*/1);
+
+  ASSERT_THAT(
+    [&] { builder.get_decision_forest(/*num_feature=*/1, /*num_class=*/1); },
+    testing::ThrowsMessage<model_import_error>(testing::HasSubstr("storage offset out of bounds")));
+}
+
+TEST(DecisionForestBuilder, CategoricalBitsetExtentOutOfBounds)
+{
+  auto builder = decision_forest_builder<test_forest_t>(
+    /*max_num_categories=*/std::uint32_t{33}, /*align_bytes=*/std::uint32_t{0});
+
+  // Create a valid categorical node first, which allocates non-local storage:
+  // categorical_storage_ = [num_categories, packed_bin_data]
+  std::array<std::uint32_t, 1> categories{0};
+  builder.add_categorical_node(categories.begin(),
+                               categories.end(),
+                               /*tl_node_id=*/0,
+                               /*depth=*/0,
+                               /*default_to_distant_child=*/false,
+                               /*feature=*/0,
+                               /*offset=*/1);
+
+  // Construct another categorical node that points at offset=1, i.e. the first
+  // packed bin entry rather than the metadata entry. The value at offset=1 is
+  // interpreted as stored_num_cats, and the resulting bins_required exceeds the
+  // available headroom.
+  builder.add_node(std::uint32_t{1},
+                   /*tl_node_id=*/1,
+                   /*depth=*/1,
+                   /*is_leaf_node=*/false,
+                   /*default_to_distant_child=*/false,
+                   /*is_categorical_node=*/true,
+                   /*feature=*/0,
+                   /*offset=*/1);
+
+  ASSERT_THAT([&] { builder.get_decision_forest(/*num_feature=*/1, /*num_class=*/1); },
+              testing::ThrowsMessage<model_import_error>(
+                testing::HasSubstr("bitset extends past categorical_storage end")));
+}
+
+}  // namespace detail
+}  // namespace fil
+}  // namespace ML