diff --git a/include/xgboost/predictor.h b/include/xgboost/predictor.h
index 78d029d199df..4320141c62a9 100644
--- a/include/xgboost/predictor.h
+++ b/include/xgboost/predictor.h
@@ -12,6 +12,8 @@
 #include <xgboost/context.h>
 #include <xgboost/data.h>
 #include <xgboost/host_device_vector.h>
+#include <xgboost/linalg.h>
+#include <xgboost/span.h>
 
 #include <functional>  // for function
 #include <memory>      // for shared_ptr
@@ -24,6 +26,7 @@ struct GBTreeModel;
 }  // namespace xgboost::gbm
 
 namespace xgboost {
+class RegTree;
 /**
  * \brief Contains pointer to input matrix and associated cached predictions.
  */
@@ -141,6 +144,18 @@ class Predictor {
   virtual void PredictLeaf(DMatrix* dmat, HostDeviceVector<float>* out_preds,
                            gbm::GBTreeModel const& model, bst_tree_t tree_end = 0) const = 0;
 
+  /**
+   * \brief Add prediction contributions from known leaf ids. The leaf ids are encoded with
+   *        tree::SamplePosition, where invalid rows are still decoded for prediction.
+   *
+   * \param leaf_ids Leaf ids for each tree, one vector per tree.
+   * \param trees Trees corresponding to the leaf id vectors.
+   * \param out_preds Prediction output to be incremented.
+   */
+  virtual void PredictFromLeafIds(common::Span<HostDeviceVector<bst_node_t> const> leaf_ids,
+                                  common::Span<RegTree const*> trees,
+                                  linalg::MatrixView<float> out_preds) const = 0;
+
   /**
    * \brief feature contributions to individual predictions; the output will be
    * a vector of length (nfeats + 1) * num_output_group * nsample, arranged in
diff --git a/include/xgboost/tree_updater.h b/include/xgboost/tree_updater.h
index 083b52b5fbdb..f90e9c796e5c 100644
--- a/include/xgboost/tree_updater.h
+++ b/include/xgboost/tree_updater.h
@@ -72,24 +72,6 @@ class TreeUpdater : public Configurable {
                       common::Span<HostDeviceVector<bst_node_t>> out_position,
                       std::vector<RegTree*> const& out_trees) = 0;
 
-  /**
-   * @brief Determines whether updater has enough knowledge about a given dataset to
-   *        quickly update prediction cache for the training data and performs the update
-   *        if possible.
-   *
-   * @param p_fmat data matrix
-   * @param out_preds prediction cache to be updated
-   *
-   * @return boolean indicating whether updater has capability to update the prediction
-   *         cache. If true, the prediction cache will have been updated by the time this
-   *         function returns.
-   */
-  virtual bool UpdatePredictionCache(DMatrix const* /*data*/,
-                                     common::Span<HostDeviceVector<bst_node_t>> /*out_position*/,
-                                     linalg::MatrixView<float> /*out_preds*/) {
-    return false;
-  }
-
   [[nodiscard]] virtual char const* Name() const = 0;
 
   /**
diff --git a/plugin/sycl/predictor/predictor.cc b/plugin/sycl/predictor/predictor.cc
index 3e81b4dc118a..87514a605116 100644
--- a/plugin/sycl/predictor/predictor.cc
+++ b/plugin/sycl/predictor/predictor.cc
@@ -10,8 +10,8 @@
 
 #include <cstddef>
 #include <limits>
-#include <mutex>
 #include <sycl/sycl.hpp>
+#include <vector>
 
 #include "../../../src/common/timer.h"
 #include "../data.h"
@@ -25,6 +25,7 @@
 #pragma GCC diagnostic pop
 #include "../../src/common/math.h"
 #include "../../src/gbm/gbtree_model.h"
+#include "../../src/tree/sample_position.h"
 #include "../device_manager.h"
 #include "../device_properties.h"
 #include "node.h"
@@ -217,6 +218,55 @@ class Predictor : public xgboost::Predictor {
     cpu_predictor->PredictLeaf(p_fmat, out_preds, model, ntree_limit);
   }
 
+  void PredictFromLeafIds(common::Span<HostDeviceVector<bst_node_t> const> leaf_ids,
+                          common::Span<RegTree const*> trees,
+                          linalg::MatrixView<float> out_preds) const override {
+    CHECK_EQ(leaf_ids.size(), trees.size());
+    if (out_preds.Device().IsCPU()) {
+      cpu_predictor->PredictFromLeafIds(leaf_ids, trees, out_preds);
+      return;
+    }
+    CHECK(out_preds.Device().IsSycl());
+    qu_ = device_manager.GetQueue(out_preds.Device());
+
+    for (std::size_t tree_idx = 0; tree_idx < trees.size(); ++tree_idx) {
+      auto const* p_tree = trees[tree_idx];
+      CHECK(p_tree);
+      CHECK(!p_tree->IsMultiTarget())
+          << "Multi-target leaf-id prediction is not implemented for SYCL.";
+      CHECK_EQ(out_preds.Shape(1), 1);
+
+      auto d_leaf_ids = leaf_ids[tree_idx].ConstDeviceSpan();
+      CHECK_EQ(d_leaf_ids.size(), out_preds.Shape(0));
+
+      auto const h_nodes = p_tree->GetNodes(DeviceOrd::CPU());
+      std::vector<float> h_leaf_values(h_nodes.size());
+      for (std::size_t nidx = 0; nidx < h_nodes.size(); ++nidx) {
+        if (h_nodes[nidx].IsLeaf()) {
+          h_leaf_values[nidx] = h_nodes[nidx].LeafValue();
+        }
+      }
+      USMVector<float, MemoryType::on_device> leaf_values;
+      leaf_values.Resize(qu_, h_leaf_values.size());
+      qu_->memcpy(leaf_values.Data(), h_leaf_values.data(),
+                  h_leaf_values.size() * sizeof(h_leaf_values.front()))
+          .wait_and_throw();
+
+      auto const* d_leaf_values = leaf_values.DataConst();
+      auto* out = out_preds.Values().data();
+      auto out_stride = out_preds.Stride(0);
+      qu_->submit([&](::sycl::handler& cgh) {
+           cgh.parallel_for<>(::sycl::range<1>(d_leaf_ids.size()), [=](::sycl::id<1> pid) {
+             auto row_idx = pid[0];
+             auto position = d_leaf_ids[row_idx];
+             auto nidx = position >= 0 ? position : ~position;
+             out[row_idx * out_stride] += d_leaf_values[nidx];
+           });
+         })
+          .wait_and_throw();
+    }
+  }
+
   void PredictContribution(DMatrix* p_fmat, HostDeviceVector<float>* out_contribs,
                            const gbm::GBTreeModel& model, bst_tree_t ntree_limit, bool approximate,
                            int condition, unsigned condition_feature) const override {
diff --git a/plugin/sycl/tree/hist_updater.cc b/plugin/sycl/tree/hist_updater.cc
index 3c92ef80b30b..7dee32eb0776 100644
--- a/plugin/sycl/tree/hist_updater.cc
+++ b/plugin/sycl/tree/hist_updater.cc
@@ -11,6 +11,7 @@
 #include "../../src/collective/allreduce.h"
 #include "../../src/tree/common_row_partitioner.h"
 #include "../../src/tree/driver.h"
+#include "../../src/tree/sample_position.h"
 #include "../common/hist_util.h"
 #include "xgboost/linalg.h"
 
@@ -309,10 +310,11 @@ void HistUpdater<GradientSumT>::ExpandWithLossGuide(const common::GHistIndexMatr
 }
 
 template <typename GradientSumT>
-void HistUpdater<GradientSumT>::Update(
-    xgboost::tree::TrainParam const* param, const common::GHistIndexMatrix& gmat,
-    const HostDeviceVector<GradientPair>& gpair, DMatrix* p_fmat,
-    xgboost::common::Span<HostDeviceVector<bst_node_t>> out_position, RegTree* p_tree) {
+void HistUpdater<GradientSumT>::Update(xgboost::tree::TrainParam const* param,
+                                       const common::GHistIndexMatrix& gmat,
+                                       const HostDeviceVector<GradientPair>& gpair, DMatrix* p_fmat,
+                                       HostDeviceVector<bst_node_t>* p_out_position,
+                                       RegTree* p_tree) {
   builder_monitor_.Start("Update");
 
   tree_evaluator_.Reset(qu_, param_, p_fmat->Info().num_col_);
@@ -330,54 +332,42 @@ void HistUpdater<GradientSumT>::Update(
     p_tree->Stat(nid).base_weight = snode_host_[nid].weight;
     p_tree->Stat(nid).sum_hess = static_cast<float>(snode_host_[nid].stats.GetHess());
   }
+  this->FinalizePosition(p_fmat->Info().num_row_, *p_tree, p_out_position);
 
   builder_monitor_.Stop("Update");
 }
 
 template <typename GradientSumT>
-bool HistUpdater<GradientSumT>::UpdatePredictionCache(
-    const DMatrix* data, ::xgboost::linalg::MatrixView<float> out_preds) {
-  CHECK(out_preds.Device().IsSycl());
-  // p_last_fmat_ is a valid pointer as long as UpdatePredictionCache() is called in
-  // conjunction with Update().
-  if (!p_last_fmat_ || !p_last_tree_ || data != p_last_fmat_) {
-    return false;
+void HistUpdater<GradientSumT>::FinalizePosition(std::size_t n_samples, RegTree const& tree,
+                                                 HostDeviceVector<bst_node_t>* p_out_position) {
+  CHECK(p_out_position);
+  if (param_.subsample < 1.0f || row_set_collection_.Data().Size() != n_samples) {
+    p_out_position->Resize(0);
+    return;
   }
-  builder_monitor_.Start("UpdatePredictionCache");
-  CHECK_GT(out_preds.Size(), 0U);
-
-  size_t n_nodes = row_set_collection_.Size();
-  std::vector<::sycl::event> events(n_nodes);
-  auto tree = p_last_tree_->HostScView();
-  for (size_t node = 0; node < n_nodes; node++) {
-    const common::RowSetCollection::Elem& rowset = row_set_collection_[node];
-    if (rowset.begin != nullptr && rowset.end != nullptr && rowset.Size() != 0) {
-      int nid = rowset.node_id;
-      // if a node is marked as deleted by the pruner, traverse upward to locate
-      // a non-deleted leaf.
-      if (tree.IsDeleted(nid)) {
-        while (tree.IsDeleted(nid)) {
-          nid = tree.Parent(nid);
-        }
-        CHECK(tree.IsLeaf(nid));
-      }
-      bst_float leaf_value = tree.LeafValue(nid);
-      const size_t* rid = rowset.begin;
-      const size_t num_rows = rowset.Size();
-
-      events[node] = qu_->submit([&](::sycl::handler& cgh) {
-        cgh.parallel_for<>(::sycl::range<1>(num_rows), [=](::sycl::item<1> pid) {
-          size_t row_id = rid[pid.get_id(0)];
-          float& val = const_cast<float&>(out_preds(row_id));
-          val += leaf_value;
-        });
-      });
-    }
+  p_out_position->SetDevice(ctx_->Device());
+  p_out_position->Resize(n_samples);
+  if (n_samples == 0) {
+    return;
   }
-  qu_->wait();
+  auto d_position = p_out_position->DeviceSpan();
 
-  builder_monitor_.Stop("UpdatePredictionCache");
-  return true;
+  auto const t = tree.HostScView();
+  for (std::size_t nid = 0; nid < row_set_collection_.Size(); ++nid) {
+    auto const& row_set = row_set_collection_[nid];
+    if (row_set.node_id < 0 || !t.IsLeaf(row_set.node_id) || row_set.Size() == 0) {
+      continue;
+    }
+    auto const encoded = xgboost::tree::SamplePosition::Encode(row_set.node_id, true);
+    auto const* rows = row_set.begin;
+    auto* position = d_position.data();
+    auto n_rows = row_set.Size();
+    qu_->submit([&](::sycl::handler& cgh) {
+         cgh.parallel_for<>(::sycl::range<1>(n_rows),
+                            [=](::sycl::id<1> pid) { position[rows[pid[0]]] = encoded; });
+       })
+        .wait_and_throw();
+  }
 }
 
 template <typename GradientSumT>
diff --git a/plugin/sycl/tree/hist_updater.h b/plugin/sycl/tree/hist_updater.h
index 575dee8ccba9..ef49d76fefa9 100644
--- a/plugin/sycl/tree/hist_updater.h
+++ b/plugin/sycl/tree/hist_updater.h
@@ -77,9 +77,7 @@ class HistUpdater {
   // update one tree, growing
   void Update(xgboost::tree::TrainParam const* param, const common::GHistIndexMatrix& gmat,
               const HostDeviceVector<GradientPair>& gpair, DMatrix* p_fmat,
-              xgboost::common::Span<HostDeviceVector<bst_node_t>> out_position, RegTree* p_tree);
-
-  bool UpdatePredictionCache(const DMatrix* data, ::xgboost::linalg::MatrixView<float> p_out_preds);
+              HostDeviceVector<bst_node_t>* p_out_position, RegTree* p_tree);
 
   void SetHistSynchronizer(HistSynchronizer<GradientSumT>* sync);
   void SetHistRowsAdder(HistRowsAdder<GradientSumT>* adder);
@@ -164,6 +162,9 @@ class HistUpdater {
   void ExpandWithLossGuide(const common::GHistIndexMatrix& gmat, RegTree* p_tree,
                            const HostDeviceVector<GradientPair>& gpair);
 
+  void FinalizePosition(std::size_t n_samples, RegTree const& tree,
+                        HostDeviceVector<bst_node_t>* p_out_position);
+
   void ReduceHists(const std::vector<int>& sync_ids, size_t nbins);
 
   inline static bool LossGuide(ExpandEntry lhs, ExpandEntry rhs) {
diff --git a/plugin/sycl/tree/updater_quantile_hist.cc b/plugin/sycl/tree/updater_quantile_hist.cc
index 14fa701e62fa..5e8db2e93ea7 100644
--- a/plugin/sycl/tree/updater_quantile_hist.cc
+++ b/plugin/sycl/tree/updater_quantile_hist.cc
@@ -2,8 +2,8 @@
  * Copyright 2017-2024 by Contributors
  * \file updater_quantile_hist.cc
  */
-#include <vector>
 #include <memory>
+#include <vector>
 
 #pragma GCC diagnostic push
 #pragma GCC diagnostic ignored "-Wtautological-constant-compare"
@@ -25,7 +25,7 @@ DMLC_REGISTRY_FILE_TAG(updater_quantile_hist_sycl);
 
 DMLC_REGISTER_PARAMETER(HistMakerTrainParam);
 
-void QuantileHistMaker::Configure(const Args& args) {
+void QuantileHistMaker::Configure(const Args &args) {
   const DeviceOrd device_spec = ctx_->Device();
   qu_ = device_manager.GetQueue(device_spec);
 
@@ -43,14 +43,9 @@ void QuantileHistMaker::Configure(const Args& args) {
   }
 }
 
-template<typename GradientSumT>
-void QuantileHistMaker::SetPimpl(std::unique_ptr<HistUpdater<GradientSumT>>* pimpl,
-                                 DMatrix *dmat) {
-  pimpl->reset(new HistUpdater<GradientSumT>(
-                ctx_,
-                qu_,
-                param_,
-                int_constraint_, dmat));
+template <typename GradientSumT>
+void QuantileHistMaker::SetPimpl(std::unique_ptr<HistUpdater<GradientSumT>> *pimpl, DMatrix *dmat) {
+  pimpl->reset(new HistUpdater<GradientSumT>(ctx_, qu_, param_, int_constraint_, dmat));
   if (collective::IsDistributed()) {
     (*pimpl)->SetHistSynchronizer(new DistributedHistSynchronizer<GradientSumT>());
     (*pimpl)->SetHistRowsAdder(new DistributedHistRowsAdder<GradientSumT>());
@@ -66,8 +61,9 @@ void QuantileHistMaker::CallUpdate(const std::unique_ptr<HistUpdater<GradientSum
                                    ::xgboost::linalg::Matrix<GradientPair> *gpair, DMatrix *dmat,
                                    xgboost::common::Span<HostDeviceVector<bst_node_t>> out_position,
                                    const std::vector<RegTree *> &trees) {
-  for (auto tree : trees) {
-    pimpl->Update(param, gmat_, *(gpair->Data()), dmat, out_position, tree);
+  CHECK_EQ(out_position.size(), trees.size());
+  for (std::size_t tree_idx = 0; tree_idx < trees.size(); ++tree_idx) {
+    pimpl->Update(param, gmat_, *(gpair->Data()), dmat, &out_position[tree_idx], trees[tree_idx]);
   }
 }
 
@@ -105,30 +101,9 @@ void QuantileHistMaker::Update(xgboost::tree::TrainParam const *param, GradientC
   p_last_dmat_ = dmat;
 }
 
-bool QuantileHistMaker::UpdatePredictionCache(const DMatrix *data,
-                                              xgboost::common::Span<HostDeviceVector<bst_node_t>>,
-                                              ::xgboost::linalg::MatrixView<float> out_preds) {
-  if (param_.subsample < 1.0f) return false;
-
-  if (hist_precision_ == HistPrecision::fp32) {
-    if (pimpl_fp32) {
-      return pimpl_fp32->UpdatePredictionCache(data, out_preds);
-    } else {
-      return false;
-    }
-  } else {
-    if (pimpl_fp64) {
-      return pimpl_fp64->UpdatePredictionCache(data, out_preds);
-    } else {
-      return false;
-    }
-  }
-}
-
 XGBOOST_REGISTER_TREE_UPDATER(QuantileHistMaker, "grow_quantile_histmaker_sycl")
-.describe("Grow tree using quantized histogram with SYCL.")
-.set_body(
-    [](Context const* ctx, ObjInfo const * task) {
+    .describe("Grow tree using quantized histogram with SYCL.")
+    .set_body([](Context const *ctx, ObjInfo const *task) {
       return new QuantileHistMaker(ctx, task);
     });
 }  // namespace tree
diff --git a/plugin/sycl/tree/updater_quantile_hist.h b/plugin/sycl/tree/updater_quantile_hist.h
index 5419c6ae825f..26886b13d31d 100644
--- a/plugin/sycl/tree/updater_quantile_hist.h
+++ b/plugin/sycl/tree/updater_quantile_hist.h
@@ -29,21 +29,20 @@ namespace sycl {
 namespace tree {
 
 // training parameters specific to this algorithm
-struct HistMakerTrainParam
-    : public XGBoostParameter<HistMakerTrainParam> {
+struct HistMakerTrainParam : public XGBoostParameter<HistMakerTrainParam> {
   bool single_precision_histogram = false;
   // declare parameters
   DMLC_DECLARE_PARAMETER(HistMakerTrainParam) {
-    DMLC_DECLARE_FIELD(single_precision_histogram).set_default(false).describe(
-        "Use single precision to build histograms.");
+    DMLC_DECLARE_FIELD(single_precision_histogram)
+        .set_default(false)
+        .describe("Use single precision to build histograms.");
   }
 };
 
 /*! \brief construct a tree using quantized feature values with SYCL backend*/
-class QuantileHistMaker: public TreeUpdater {
+class QuantileHistMaker : public TreeUpdater {
  public:
-  QuantileHistMaker(Context const* ctx, ObjInfo const * task) :
-                             TreeUpdater(ctx), task_{task} {
+  QuantileHistMaker(Context const* ctx, ObjInfo const* task) : TreeUpdater(ctx), task_{task} {
     updater_monitor_.Init("SYCLQuantileHistMaker");
   }
   void Configure(const Args& args) override;
@@ -52,10 +51,6 @@ class QuantileHistMaker: public TreeUpdater {
               xgboost::common::Span<HostDeviceVector<bst_node_t>> out_position,
               const std::vector<RegTree*>& trees) override;
 
-  bool UpdatePredictionCache(const DMatrix* data,
-                             xgboost::common::Span<HostDeviceVector<bst_node_t>>,
-                             ::xgboost::linalg::MatrixView<float> out_preds) override;
-
   void LoadConfig(Json const& in) override {
     auto const& config = get<Object const>(in);
     FromJson(config.at("train_param"), &this->param_);
@@ -68,9 +63,7 @@ class QuantileHistMaker: public TreeUpdater {
     out["sycl_hist_train_param"] = ToJson(hist_maker_param_);
   }
 
-  char const* Name() const override {
-    return "grow_quantile_histmaker_sycl";
-  }
+  char const* Name() const override { return "grow_quantile_histmaker_sycl"; }
 
  protected:
   HistMakerTrainParam hist_maker_param_;
@@ -80,23 +73,22 @@ class QuantileHistMaker: public TreeUpdater {
   common::GHistIndexMatrix gmat_;
   // (optional) data matrix with feature grouping
   // column accessor
-  DMatrix const* p_last_dmat_ {nullptr};
-  bool is_gmat_initialized_ {false};
+  DMatrix const* p_last_dmat_{nullptr};
+  bool is_gmat_initialized_{false};
 
   xgboost::common::Monitor updater_monitor_;
 
-  template<typename GradientSumT>
-  void SetPimpl(std::unique_ptr<HistUpdater<GradientSumT>>*, DMatrix *dmat);
+  template <typename GradientSumT>
+  void SetPimpl(std::unique_ptr<HistUpdater<GradientSumT>>*, DMatrix* dmat);
 
-  template<typename GradientSumT>
+  template <typename GradientSumT>
   void CallUpdate(const std::unique_ptr<HistUpdater<GradientSumT>>& builder,
-                  xgboost::tree::TrainParam const *param,
-                  ::xgboost::linalg::Matrix<GradientPair> *gpair,
-                  DMatrix *dmat,
+                  xgboost::tree::TrainParam const* param,
+                  ::xgboost::linalg::Matrix<GradientPair>* gpair, DMatrix* dmat,
                   xgboost::common::Span<HostDeviceVector<bst_node_t>> out_position,
-                  const std::vector<RegTree *> &trees);
+                  const std::vector<RegTree*>& trees);
 
-  enum class HistPrecision {fp32, fp64};
+  enum class HistPrecision { fp32, fp64 };
   HistPrecision hist_precision_;
 
   std::unique_ptr<HistUpdater<float>> pimpl_fp32;
@@ -106,10 +98,9 @@ class QuantileHistMaker: public TreeUpdater {
 
   ::sycl::queue* qu_;
   DeviceManager device_manager;
-  ObjInfo const *task_{nullptr};
+  ObjInfo const* task_{nullptr};
 };
 
-
 }  // namespace tree
 }  // namespace sycl
 }  // namespace xgboost
diff --git a/src/gbm/gbtree.cc b/src/gbm/gbtree.cc
index 00908a9f3379..1529dd87ecb8 100644
--- a/src/gbm/gbtree.cc
+++ b/src/gbm/gbtree.cc
@@ -231,6 +231,11 @@ void GBTree::DoBoost(DMatrix* p_fmat, GradientContainer* in_gpair, PredictionCac
   }
 
   predt->predictions.SetDevice(ctx_->Device());
+  auto const& predictor = this->GetPredictor(false, &predt->predictions, p_fmat);
+  if (predt->predictions.Size() == 0 && p_fmat->Info().num_row_ != 0) {
+    CHECK_EQ(predt->version, 0);
+    predictor->InitOutPredictions(p_fmat->Info(), &predt->predictions, model_);
+  }
   auto out = linalg::MakeTensorView(ctx_, &predt->predictions, p_fmat->Info().num_row_,
                                     model_.learner_model_param->OutputLength());
   CHECK_NE(n_groups, 0);
@@ -238,27 +243,41 @@ void GBTree::DoBoost(DMatrix* p_fmat, GradientContainer* in_gpair, PredictionCac
   // The node position for each row, 1 HDV for each tree in the forest.  Note that the
   // position is negated if the row is sampled out.
   std::vector<HostDeviceVector<bst_node_t>> node_position;
+  auto update_prediction_cache = [&](std::vector<HostDeviceVector<bst_node_t>>& positions,
+                                     TreesOneGroup const& trees,
+                                     linalg::MatrixView<float> out_preds) {
+    CHECK_EQ(positions.size(), trees.size());
+    for (auto& position : positions) {
+      if (out_preds.Shape(0) != 0 && position.Size() != out_preds.Shape(0)) {
+        return false;
+      }
+      position.SetDevice(predt->predictions.Device());
+    }
+    std::vector<RegTree const*> tree_ptrs;
+    tree_ptrs.reserve(trees.size());
+    for (auto const& tree : trees) {
+      tree_ptrs.push_back(tree.get());
+    }
+    predictor->PredictFromLeafIds(common::Span{positions}, common::Span{tree_ptrs}, out_preds);
+    return true;
+  };
 
   if (model_.learner_model_param->IsVectorLeaf()) {
     // Multi-target, vector leaf
     TreesOneGroup ret;
     BoostNewTrees(in_gpair, p_fmat, 0, &node_position, &ret);
-    std::size_t num_new_trees = ret.size();
-    new_trees.push_back(std::move(ret));
-    if (updaters_.size() > 0 && num_new_trees == 1 && predt->predictions.Size() > 0 &&
-        updaters_.back()->UpdatePredictionCache(p_fmat, common::Span{node_position}, out)) {
+    if (update_prediction_cache(node_position, ret, out)) {
       predt->Update(1);
     }
+    new_trees.push_back(std::move(ret));
   } else if (model_.learner_model_param->OutputLength() == 1u) {
     // Single target
     TreesOneGroup ret;
     BoostNewTrees(in_gpair, p_fmat, 0, &node_position, &ret);
-    const size_t num_new_trees = ret.size();
-    new_trees.push_back(std::move(ret));
-    if (updaters_.size() > 0 && num_new_trees == 1 && predt->predictions.Size() > 0 &&
-        updaters_.back()->UpdatePredictionCache(p_fmat, common::Span{node_position}, out)) {
+    if (update_prediction_cache(node_position, ret, out)) {
       predt->Update(1);
     }
+    new_trees.push_back(std::move(ret));
   } else {
     // Multi-target, scalar leaf
     CHECK_EQ(in_gpair->gpair.Size() % n_groups, 0U)
@@ -266,23 +285,17 @@ void GBTree::DoBoost(DMatrix* p_fmat, GradientContainer* in_gpair, PredictionCac
     GradientContainer tmp;
     tmp.gpair = linalg::Matrix<GradientPair>{
         {in_gpair->gpair.Shape(0), static_cast<std::size_t>(1ul)}, ctx_->Device()};
-    bool update_predict = true;
+    bool cache_updated{true};
     for (bst_target_t gid = 0; gid < n_groups; ++gid) {
       node_position.clear();
       CopyGradient(ctx_, &in_gpair->gpair, gid, &tmp.gpair);
       TreesOneGroup ret;
       BoostNewTrees(&tmp, p_fmat, gid, &node_position, &ret);
-      const size_t num_new_trees = ret.size();
-      new_trees.push_back(std::move(ret));
       auto v_predt = out.Slice(linalg::All(), linalg::Range(gid, gid + 1));
-      // random forest doesn't support the prediction cache yet.
-      if (!(updaters_.size() > 0 && predt->predictions.Size() > 0 && num_new_trees == 1 &&
-            updaters_.back()->UpdatePredictionCache(p_fmat, common::Span{node_position},
-                                                    v_predt))) {
-        update_predict = false;
-      }
+      cache_updated = update_prediction_cache(node_position, ret, v_predt) && cache_updated;
+      new_trees.push_back(std::move(ret));
     }
-    if (update_predict) {
+    if (cache_updated) {
       predt->Update(1);
     }
   }
diff --git a/src/predictor/cpu_predictor.cc b/src/predictor/cpu_predictor.cc
index 07480210b47e..d29c0c23c847 100644
--- a/src/predictor/cpu_predictor.cc
+++ b/src/predictor/cpu_predictor.cc
@@ -21,6 +21,7 @@
 #include "../data/gradient_index.h"          // for GHistIndexMatrix
 #include "../data/proxy_dmatrix.h"           // for DMatrixProxy
 #include "../gbm/gbtree_model.h"             // for GBTreeModel, GBTreeModelParam
+#include "../tree/sample_position.h"         // for SamplePosition
 #include "array_tree_layout.h"               // for ProcessArrayTree
 #include "data_accessor.h"                   // for GHistIndexMatrixView, SparsePageView
 #include "dmlc/registry.h"                   // for DMLC_REGISTRY_FILE_TAG
@@ -567,6 +568,40 @@ class CPUPredictor : public Predictor {
     });
   }
 
+  void PredictFromLeafIds(common::Span<HostDeviceVector<bst_node_t> const> leaf_ids,
+                          common::Span<RegTree const *> trees,
+                          linalg::MatrixView<float> out_preds) const override {
+    CHECK_EQ(leaf_ids.size(), trees.size());
+    CHECK(out_preds.Device().IsCPU());
+
+    for (std::size_t tree_idx = 0; tree_idx < trees.size(); ++tree_idx) {
+      auto const *p_tree = trees[tree_idx];
+      CHECK(p_tree);
+      auto const h_leaf_ids = leaf_ids[tree_idx].ConstHostSpan();
+      CHECK_EQ(h_leaf_ids.size(), out_preds.Shape(0));
+
+      if (!p_tree->IsMultiTarget()) {
+        CHECK_EQ(out_preds.Shape(1), 1);
+        auto const tree = p_tree->HostScView();
+        common::ParallelFor(out_preds.Shape(0), ctx_->Threads(), [&](std::size_t row_idx) {
+          auto nidx = tree::SamplePosition::Decode(h_leaf_ids[row_idx]);
+          out_preds(row_idx, 0) += tree.LeafValue(nidx);
+        });
+      } else {
+        auto const tree = p_tree->HostMtView();
+        auto n_targets = tree.NumTargets();
+        CHECK_EQ(out_preds.Shape(1), n_targets);
+        common::ParallelFor(out_preds.Shape(0), ctx_->Threads(), [&](std::size_t row_idx) {
+          auto nidx = tree::SamplePosition::Decode(h_leaf_ids[row_idx]);
+          auto weight = tree.LeafValue(nidx);
+          for (bst_target_t target_idx = 0; target_idx < n_targets; ++target_idx) {
+            out_preds(row_idx, target_idx) += weight(target_idx);
+          }
+        });
+      }
+    }
+  }
+
   void PredictContribution(DMatrix *p_fmat, HostDeviceVector<float> *out_contribs,
                            const gbm::GBTreeModel &model, bst_tree_t ntree_limit, bool approximate,
                            int condition, unsigned condition_feature) const override {
diff --git a/src/predictor/gpu_predictor.cu b/src/predictor/gpu_predictor.cu
index 0db5cc752c77..a8d88ca9adf2 100644
--- a/src/predictor/gpu_predictor.cu
+++ b/src/predictor/gpu_predictor.cu
@@ -21,6 +21,7 @@
 #include "../data/proxy_dmatrix.cuh"  // for DispatchAny
 #include "../data/proxy_dmatrix.h"
 #include "../gbm/gbtree_model.h"
+#include "../tree/sample_position.h"
 #include "../tree/tree_view.h"
 #include "gbtree_view.h"  // for GBTreeModelView
 #include "gpu_data_accessor.cuh"
@@ -601,6 +602,45 @@ class GPUPredictor : public xgboost::Predictor {
       });
     });
   }
+
+  void PredictFromLeafIds(common::Span<HostDeviceVector<bst_node_t> const> leaf_ids,
+                          common::Span<RegTree const*> trees,
+                          linalg::MatrixView<float> out_preds) const override {
+    CHECK_EQ(leaf_ids.size(), trees.size());
+    CHECK(out_preds.Device().IsCUDA());
+    CHECK_EQ(out_preds.Device().ordinal, ctx_->Ordinal());
+    auto stream = ctx_->CUDACtx()->Stream();
+
+    for (std::size_t tree_idx = 0; tree_idx < trees.size(); ++tree_idx) {
+      auto const* p_tree = trees[tree_idx];
+      CHECK(p_tree);
+      auto d_leaf_ids = leaf_ids[tree_idx].ConstDeviceSpan();
+      CHECK_EQ(d_leaf_ids.size(), out_preds.Shape(0));
+
+      if (!p_tree->IsMultiTarget()) {
+        CHECK_EQ(out_preds.Shape(1), 1);
+        dh::CachingDeviceUVector<RegTree::Node> nodes;
+        dh::CopyTo(p_tree->GetNodes(DeviceOrd::CPU()), &nodes, stream);
+        common::Span<RegTree::Node> d_nodes = dh::ToSpan(nodes);
+        dh::LaunchN(d_leaf_ids.size(), stream, [=] XGBOOST_DEVICE(std::size_t row_idx) mutable {
+          auto nidx = tree::SamplePosition::Decode(d_leaf_ids[row_idx]);
+          out_preds(row_idx, 0) += d_nodes[nidx].LeafValue();
+        });
+      } else {
+        auto mt_tree = tree::MultiTargetTreeView{ctx_->Device(), false, p_tree};
+        auto n_targets = mt_tree.NumTargets();
+        CHECK_EQ(out_preds.Shape(1), n_targets);
+        thrust::for_each_n(ctx_->CUDACtx()->CTP(), thrust::make_counting_iterator(0ul),
+                           out_preds.Size(), [=] XGBOOST_DEVICE(std::size_t i) mutable {
+                             auto [row_idx, target_idx] =
+                                 linalg::UnravelIndex(i, out_preds.Shape());
+                             auto nidx = tree::SamplePosition::Decode(d_leaf_ids[row_idx]);
+                             auto weight = mt_tree.LeafValue(nidx);
+                             out_preds(row_idx, target_idx) += weight(target_idx);
+                           });
+      }
+    }
+  }
 };
 
 XGBOOST_REGISTER_PREDICTOR(GPUPredictor, "gpu_predictor")
diff --git a/src/tree/hist/evaluate_splits.h b/src/tree/hist/evaluate_splits.h
index 762ef9150211..f3c4d62faba6 100644
--- a/src/tree/hist/evaluate_splits.h
+++ b/src/tree/hist/evaluate_splits.h
@@ -18,7 +18,6 @@
 #include "../../common/random.h"       // for ColumnSampler
 #include "../constraints.h"            // for FeatureInteractionConstraintHost
 #include "../param.h"                  // for TrainParam
-#include "../sample_position.h"        // for SamplePosition
 #include "../split_evaluator.h"        // for TreeEvaluator
 #include "../tree_view.h"              // for MultiTargetTreeView
 #include "expand_entry.h"              // for MultiExpandEntry
@@ -825,49 +824,5 @@ class HistMultiEvaluator {
   [[nodiscard]] auto Evaluator() const { return tree_evaluator_.GetEvaluator(); }
 };
 
-/**
- * @brief CPU implementation of update prediction cache, which calculates the leaf value
- *        for the last tree and accumulates it to prediction vector.
- *
- * @param last_tree The last tree being updated by tree updater
- */
-inline void UpdatePredictionCacheImpl(Context const *ctx, ScalarTreeView const &last_tree,
-                                      common::Span<bst_node_t const> node_position,
-                                      linalg::VectorView<float> out_preds) {
-  CHECK(out_preds.Device().IsCPU());
-  common::ParallelFor(out_preds.Size(), ctx->Threads(), [&](std::size_t idx) {
-    bst_node_t nidx = node_position[idx];
-    nidx = SamplePosition::Decode(nidx);
-    auto weight = last_tree.LeafValue(nidx);
-    out_preds(idx) += weight;
-  });
-}
-
-inline void UpdatePredictionCacheImpl(Context const *ctx, RegTree const *p_last_tree,
-                                      common::Span<bst_node_t const> node_position,
-                                      linalg::MatrixView<float> out_preds) {
-  CHECK_GT(out_preds.Size(), 0U);
-  CHECK(p_last_tree);
-
-  auto const &tree = *p_last_tree;
-  if (!tree.IsMultiTarget()) {
-    return UpdatePredictionCacheImpl(ctx, p_last_tree->HostScView(), node_position,
-                                     out_preds.Slice(linalg::All(), 0));
-  }
-
-  auto const mt_tree = tree.HostMtView();
-  auto n_targets = mt_tree.NumTargets();
-  CHECK_EQ(out_preds.Shape(1), n_targets);
-  CHECK(out_preds.Device().IsCPU());
-
-  common::ParallelFor(out_preds.Shape(0), ctx->Threads(), [&](std::size_t sample_idx) {
-    bst_node_t nidx = node_position[sample_idx];
-    nidx = SamplePosition::Decode(nidx);
-    auto weight = mt_tree.LeafValue(nidx);
-    for (bst_target_t target_idx = 0; target_idx < n_targets; ++target_idx) {
-      out_preds(sample_idx, target_idx) += weight(target_idx);
-    }
-  });
-}
 }  // namespace xgboost::tree
 #endif  // XGBOOST_TREE_HIST_EVALUATE_SPLITS_H_
diff --git a/src/tree/updater_approx.cc b/src/tree/updater_approx.cc
index 910942a321f7..21a89eb48cf0 100644
--- a/src/tree/updater_approx.cc
+++ b/src/tree/updater_approx.cc
@@ -19,7 +19,7 @@
 #include "common_row_partitioner.h"          // for CommonRowPartitioner
 #include "dmlc/registry.h"                   // for DMLC_REGISTRY_FILE_TAG
 #include "driver.h"                          // for Driver
-#include "hist/evaluate_splits.h"            // for HistEvaluator, UpdatePredictionCacheImpl
+#include "hist/evaluate_splits.h"            // for HistEvaluator
 #include "hist/expand_entry.h"               // for CPUExpandEntry
 #include "hist/hist_param.h"                 // for HistMakerTrainParam
 #include "hist/histogram.h"                  // for MultiHistogramBuilder
@@ -128,17 +128,6 @@ class GlobalApproxBuilder {
     return nodes.front();
   }
 
-  void UpdatePredictionCache(DMatrix const *p_fmat, common::Span<bst_node_t const> node_position,
-                             linalg::MatrixView<float> out_preds) const {
-    monitor_->Start(__func__);
-    // Caching prediction seems redundant for approx tree method, as sketching takes up
-    // majority of training time.
-    CHECK_EQ(out_preds.Size(), p_fmat->Info().num_row_);
-    CHECK_EQ(node_position.size(), p_fmat->Info().num_row_);
-    UpdatePredictionCacheImpl(ctx_, p_last_tree_, node_position, out_preds);
-    monitor_->Stop(__func__);
-  }
-
   void BuildHistogram(DMatrix *p_fmat, RegTree *p_tree,
                       std::vector<CPUExpandEntry> const &valid_candidates,
                       std::vector<GradientPair> const &gpair, common::Span<float> hess) {
@@ -308,19 +297,6 @@ class GlobalApproxUpdater : public TreeUpdater {
       ++t_idx;
     }
   }
-
-  bool UpdatePredictionCache(DMatrix const *p_fmat,
-                             common::Span<HostDeviceVector<bst_node_t>> out_position,
-                             linalg::MatrixView<float> out_preds) override {
-    if (p_fmat != cached_ || !pimpl_) {
-      return false;
-    }
-    if (out_position.size() > 1) {
-      return false;
-    }
-    this->pimpl_->UpdatePredictionCache(p_fmat, out_position.front().ConstHostSpan(), out_preds);
-    return true;
-  }
 };
 
 DMLC_REGISTRY_FILE_TAG(grow_histmaker);
diff --git a/src/tree/updater_gpu_hist.cu b/src/tree/updater_gpu_hist.cu
index 4621ce3ab3ff..50c906ec56d4 100644
--- a/src/tree/updater_gpu_hist.cu
+++ b/src/tree/updater_gpu_hist.cu
@@ -504,34 +504,6 @@ struct GPUHistMakerDevice {
     }
   }
 
-  bool UpdatePredictionCache(linalg::MatrixView<float> out_preds_d,
-                             common::Span<HostDeviceVector<bst_node_t>> out_position,
-                             RegTree const* p_tree) {
-    CHECK(p_tree);
-    CHECK(out_preds_d.Device().IsCUDA());
-    CHECK_EQ(out_preds_d.Device().ordinal, ctx_->Ordinal());
-
-    CHECK_EQ(out_position.size(), 1);
-    auto d_position = out_position.front().ConstDeviceSpan();
-    CHECK_EQ(out_preds_d.Size(), d_position.size());
-
-    // Use the nodes from tree, the leaf value might be changed by the objective since the
-    // last update tree call.
-    dh::CachingDeviceUVector<RegTree::Node> nodes;
-    // We can remove the CPU copy once we refactor the GPU hist to use the device tree.
-    dh::CopyTo(p_tree->GetNodes(DeviceOrd::CPU()), &nodes, this->ctx_->CUDACtx()->Stream());
-    common::Span<RegTree::Node> d_nodes = dh::ToSpan(nodes);
-    CHECK_EQ(out_preds_d.Shape(1), 1);
-    dh::LaunchN(d_position.size(), ctx_->CUDACtx()->Stream(),
-                [=] XGBOOST_DEVICE(std::size_t idx) mutable {
-                  bst_node_t nidx = d_position[idx];
-                  nidx = SamplePosition::Decode(nidx);
-                  auto weight = d_nodes[nidx].LeafValue();
-                  out_preds_d(idx, 0) += weight;
-                });
-    return true;
-  }
-
   void ApplySplit(const GPUExpandEntry& candidate, RegTree* p_tree) {
     RegTree& tree = *p_tree;
 
@@ -762,26 +734,6 @@ class GPUHistMaker : public TreeUpdater {
     p_scimpl_->UpdateTree(gpair_hdv, p_fmat, p_tree, p_out_position);
   }
 
-  bool UpdatePredictionCache(DMatrix const* p_fmat,
-                             common::Span<HostDeviceVector<bst_node_t>> out_position,
-                             linalg::MatrixView<float> p_out_preds) override {
-    if (p_scimpl_ == nullptr || p_last_fmat_ == nullptr || p_last_fmat_ != p_fmat) {
-      return false;
-    }
-    if (out_position.size() > 1) {
-      return false;
-    }
-
-    xgboost_NVTX_FN_RANGE();
-
-    if (this->p_last_tree_->IsMultiTarget()) {
-      CHECK(p_mtimpl_);
-      return p_mtimpl_->UpdatePredictionCache(p_out_preds, out_position, p_last_tree_);
-    } else {
-      return p_scimpl_->UpdatePredictionCache(p_out_preds, out_position, p_last_tree_);
-    }
-  }
-
   [[nodiscard]] char const* Name() const override { return "grow_gpu_hist"; }
 
  private:
@@ -900,21 +852,6 @@ class GPUGlobalApproxMaker : public TreeUpdater {
     maker_->UpdateTree(gpair, p_fmat, p_tree, p_out_position);
   }
 
-  bool UpdatePredictionCache(DMatrix const* p_fmat,
-                             common::Span<HostDeviceVector<bst_node_t>> out_position,
-                             linalg::MatrixView<float> p_out_preds) override {
-    if (maker_ == nullptr || p_last_fmat_ == nullptr || p_last_fmat_ != p_fmat) {
-      return false;
-    }
-    if (out_position.size() > 1) {
-      return false;
-    }
-    monitor_.Start(__func__);
-    bool result = maker_->UpdatePredictionCache(p_out_preds, out_position, p_last_tree_);
-    monitor_.Stop(__func__);
-    return result;
-  }
-
   [[nodiscard]] char const* Name() const override { return "grow_gpu_approx"; }
 
  private:
diff --git a/src/tree/updater_gpu_hist.cuh b/src/tree/updater_gpu_hist.cuh
index 830c7813f825..a2f3d33df288 100644
--- a/src/tree/updater_gpu_hist.cuh
+++ b/src/tree/updater_gpu_hist.cuh
@@ -636,27 +636,6 @@ class MultiTargetHistMaker {
     }
   }
 
-  bool UpdatePredictionCache(linalg::MatrixView<float> out_preds_d,
-                             common::Span<HostDeviceVector<bst_node_t>> out_position,
-                             RegTree const* p_tree) {
-    xgboost_NVTX_FN_RANGE();
-
-    CHECK_EQ(out_position.size(), 1);
-    auto d_position = out_position.front().ConstDeviceSpan();
-    CHECK_EQ(out_preds_d.Shape(0), d_position.size());
-    auto mt_tree = MultiTargetTreeView{this->ctx_->Device(), false, p_tree};
-    thrust::for_each_n(this->ctx_->CUDACtx()->CTP(), thrust::make_counting_iterator(0ul),
-                       out_preds_d.Size(), [=] XGBOOST_DEVICE(std::size_t i) mutable {
-                         auto [sample_idx, target_idx] =
-                             linalg::UnravelIndex(i, out_preds_d.Shape());
-                         bst_node_t nidx = d_position[sample_idx];
-                         nidx = SamplePosition::Decode(nidx);
-                         auto weight = mt_tree.LeafValue(nidx);
-                         out_preds_d(sample_idx, target_idx) += weight(target_idx);
-                       });
-    return true;
-  }
-
   void UpdateTree(GradientContainer* gpair, DMatrix* p_fmat, ObjInfo const* task, RegTree* p_tree,
                   HostDeviceVector<bst_node_t>* p_out_position) {
     xgboost_NVTX_FN_RANGE();
diff --git a/src/tree/updater_quantile_hist.cc b/src/tree/updater_quantile_hist.cc
index e0a0ce937f9b..82f09ecc24ed 100644
--- a/src/tree/updater_quantile_hist.cc
+++ b/src/tree/updater_quantile_hist.cc
@@ -400,21 +400,6 @@ class MultiTargetHistBuilder {
         ctx_{ctx} {
     monitor_->Init(__func__);
   }
-
-  bool UpdatePredictionCache(DMatrix const *p_fmat, common::Span<bst_node_t const> node_position,
-                             linalg::MatrixView<float> out_preds) const {
-    // p_last_fmat_ is a valid pointer as long as UpdatePredictionCache() is called in
-    // conjunction with Update().
-    if (!p_last_fmat_ || !p_last_tree_ || p_fmat != p_last_fmat_) {
-      return false;
-    }
-    monitor_->Start(__func__);
-    CHECK_EQ(out_preds.Size(), p_fmat->Info().num_row_ * p_last_tree_->NumTargets());
-    CHECK_EQ(node_position.size(), p_fmat->Info().num_row_);
-    UpdatePredictionCacheImpl(ctx_, p_last_tree_, node_position, out_preds);
-    monitor_->Stop(__func__);
-    return true;
-  }
 };
 
 /**
@@ -451,21 +436,6 @@ class HistUpdater {
     monitor_->Init(__func__);
   }
 
-  bool UpdatePredictionCache(DMatrix const *data, common::Span<bst_node_t const> node_position,
-                             linalg::MatrixView<float> out_preds) const {
-    // p_last_fmat_ is a valid pointer as long as UpdatePredictionCache() is called in
-    // conjunction with Update().
-    if (!p_last_fmat_ || !p_last_tree_ || data != p_last_fmat_) {
-      return false;
-    }
-    monitor_->Start(__func__);
-    CHECK_EQ(out_preds.Size(), data->Info().num_row_);
-    CHECK_EQ(node_position.size(), data->Info().num_row_);
-    UpdatePredictionCacheImpl(ctx_, p_last_tree_, node_position, out_preds);
-    monitor_->Stop(__func__);
-    return true;
-  }
-
  public:
   // initialize temp data structure
   void InitData(DMatrix *fmat, RegTree const *p_tree, linalg::MatrixView<GradientPair const>) {
@@ -687,22 +657,6 @@ class QuantileHistMaker : public TreeUpdater {
       hist_param_.CheckTreesSynchronized(ctx_, *tree_it);
     }
   }
-
-  bool UpdatePredictionCache(DMatrix const *p_fmat,
-                             common::Span<HostDeviceVector<bst_node_t>> node_position,
-                             linalg::MatrixView<float> out_preds) override {
-    if (node_position.size() > 1) {
-      return false;
-    }
-    auto position = node_position.front().ConstHostSpan();
-    if (out_preds.Shape(1) > 1) {
-      CHECK(p_mtimpl_);
-      return p_mtimpl_->UpdatePredictionCache(p_fmat, position, out_preds);
-    } else {
-      CHECK(p_impl_);
-      return p_impl_->UpdatePredictionCache(p_fmat, position, out_preds);
-    }
-  }
 };
 
 XGBOOST_REGISTER_TREE_UPDATER(QuantileHistMaker, "grow_quantile_histmaker")
diff --git a/tests/cpp/predictor/test_cpu_predictor.cc b/tests/cpp/predictor/test_cpu_predictor.cc
index 21b1c2f1f086..8bfad3452e73 100644
--- a/tests/cpp/predictor/test_cpu_predictor.cc
+++ b/tests/cpp/predictor/test_cpu_predictor.cc
@@ -167,7 +167,7 @@ TEST(CpuPredictor, InplacePredict) {
 }
 
 namespace {
-void TestUpdatePredictionCache(bool use_subsampling) {
+void TestTrainingPredictionCache(bool use_subsampling) {
   std::size_t constexpr kRows = 64, kCols = 16, kClasses = 4;
   LearnerModelParam mparam{MakeMP(kCols, .0, kClasses)};
   Context ctx;
@@ -226,9 +226,9 @@ TEST(CPUPredictor, CategoricalPredictLeaf) {
   TestCategoricalPredictLeaf(&ctx);
 }
 
-TEST(CpuPredictor, UpdatePredictionCache) {
-  TestUpdatePredictionCache(false);
-  TestUpdatePredictionCache(true);
+TEST(CpuPredictor, TrainingPredictionCache) {
+  TestTrainingPredictionCache(false);
+  TestTrainingPredictionCache(true);
 }
 
 TEST(CpuPredictor, LesserFeatures) {
diff --git a/tests/cpp/predictor/test_predictor.cc b/tests/cpp/predictor/test_predictor.cc
index 2ffcd073237e..77380cf4c535 100644
--- a/tests/cpp/predictor/test_predictor.cc
+++ b/tests/cpp/predictor/test_predictor.cc
@@ -57,6 +57,26 @@ void TestBasic(DMatrix *dmat, Context const *ctx) {
   for (auto v : h_leaf_out_predictions) {
     ASSERT_EQ(v, 0);
   }
+
+  std::vector<HostDeviceVector<bst_node_t>> leaf_ids(1);
+  leaf_ids.front().SetDevice(ctx->Device());
+  leaf_ids.front().Resize(h_leaf_out_predictions.size());
+  auto &h_leaf_ids = leaf_ids.front().HostVector();
+  for (std::size_t i = 0; i < h_leaf_out_predictions.size(); ++i) {
+    h_leaf_ids[i] = static_cast<bst_node_t>(h_leaf_out_predictions[i]);
+  }
+  std::vector<RegTree const *> trees{model.trees.front().get()};
+  PredictionCacheEntry from_leaf_ids;
+  predictor->InitOutPredictions(dmat->Info(), &from_leaf_ids.predictions, model);
+  auto from_leaf_view =
+      linalg::MakeTensorView(ctx, &from_leaf_ids.predictions, dmat->Info().num_row_,
+                             model.learner_model_param->OutputLength());
+  predictor->PredictFromLeafIds(common::Span{leaf_ids}, common::Span{trees}, from_leaf_view);
+  auto const &h_from_leaf_ids = from_leaf_ids.predictions.ConstHostVector();
+  ASSERT_EQ(h_from_leaf_ids.size(), out_predictions_h.size());
+  for (std::size_t i = 0; i < h_from_leaf_ids.size(); ++i) {
+    ASSERT_EQ(h_from_leaf_ids[i], out_predictions_h[i]);
+  }
 }
 
 void TestBatchPredictionWithWeights(Context const *ctx) {
diff --git a/tests/cpp/tree/test_gpu_hist.cu b/tests/cpp/tree/test_gpu_hist.cu
index 46c4103e12a8..d40e0c4205e2 100644
--- a/tests/cpp/tree/test_gpu_hist.cu
+++ b/tests/cpp/tree/test_gpu_hist.cu
@@ -7,6 +7,7 @@
 #include <xgboost/gradient.h>            // for GradientContainer
 #include <xgboost/host_device_vector.h>  // for HostDeviceVector
 #include <xgboost/json.h>                // for Json
+#include <xgboost/predictor.h>           // for Predictor
 #include <xgboost/task.h>                // for ObjInfo
 #include <xgboost/tree_model.h>          // for RegTree
 #include <xgboost/tree_updater.h>        // for TreeUpdater
@@ -17,6 +18,7 @@
 
 #include "../../../src/tree/param.h"  // for TrainParam
 #include "../helpers.h"
+#include "../predictor/test_predictor.h"  // for CreatePredictorForTest
 
 namespace xgboost::tree {
 namespace {
@@ -43,7 +45,9 @@ void UpdateTree(Context const* ctx, GradientContainer* gpair, DMatrix* dmat, Reg
   hist_maker->Update(&param, gpair, dmat, common::Span<HostDeviceVector<bst_node_t>>{position},
                      {tree});
   auto cache = linalg::MakeTensorView(ctx, preds->DeviceSpan(), preds->Size(), 1);
-  ASSERT_TRUE(hist_maker->UpdatePredictionCache(dmat, common::Span{position}, cache));
+  std::unique_ptr<Predictor> predictor{CreatePredictorForTest(ctx)};
+  std::vector<RegTree const*> tree_ptrs{tree};
+  predictor->PredictFromLeafIds(common::Span{position}, common::Span{tree_ptrs}, cache);
 }
 }  // anonymous namespace
 
diff --git a/tests/cpp/tree/test_prediction_cache.h b/tests/cpp/tree/test_prediction_cache.h
index 39073360de32..2db0a250cd76 100644
--- a/tests/cpp/tree/test_prediction_cache.h
+++ b/tests/cpp/tree/test_prediction_cache.h
@@ -4,15 +4,16 @@
 #pragma once
 
 #include <gtest/gtest.h>
-
 #include <xgboost/host_device_vector.h>
+#include <xgboost/predictor.h>
 #include <xgboost/tree_updater.h>
 
 #include <memory>
 
 #include "../../../src/tree/param.h"  // for TrainParam
 #include "../helpers.h"
-#include "xgboost/task.h"             // for ObjInfo
+#include "../predictor/test_predictor.h"  // for CreatePredictorForTest
+#include "xgboost/task.h"                 // for ObjInfo
 
 namespace xgboost {
 class TestPredictionCache : public ::testing::Test {
@@ -83,7 +84,11 @@ class TestPredictionCache : public ::testing::Test {
       out_prediction_cached.Resize(n_samples_);
       auto cache =
           linalg::MakeTensorView(ctx, &out_prediction_cached, out_prediction_cached.Size(), 1);
-      ASSERT_TRUE(updater->UpdatePredictionCache(Xy_.get(), common::Span{position}, cache));
+      if (position.front().Size() == n_samples_) {
+        std::unique_ptr<Predictor> predictor{CreatePredictorForTest(ctx)};
+        std::vector<RegTree const*> tree_ptrs{&tree};
+        predictor->PredictFromLeafIds(common::Span{position}, common::Span{tree_ptrs}, cache);
+      }
     }
 
     for (auto policy : {"depthwise", "lossguide"}) {