From 6958e39c45c6a1857f4349c2f220cbc82d83281b Mon Sep 17 00:00:00 2001 From: Ron Wettenstein Date: Mon, 5 Oct 2026 15:15:15 +0300 Subject: [PATCH 1/2] Keep the current path's features in shared memory for GPU SHAP interactions The interaction kernel needs the distinct features on the current tree path at every return edge. Keep each depth's split feature, and whether a deeper split on the same feature shadows it, in per-warp shared memory, updated as the walk enters and leaves nodes. This replaces the O(D^2) rescan of the path in global memory with an O(D) pass over shared memory. Co-Authored-By: Claude Opus 5.5 --- src/predictor/interpretability/shap.cu | 97 +++++++++++++++----------- 1 file changed, 58 insertions(+), 39 deletions(-) diff --git a/src/predictor/interpretability/shap.cu b/src/predictor/interpretability/shap.cu index 7e9b56496a76..fd98fd725fdf 100644 --- a/src/predictor/interpretability/shap.cu +++ b/src/predictor/interpretability/shap.cu @@ -497,6 +497,26 @@ struct QuadratureSharedState { } }; +// Interactions also enumerate the features on the active path. All rows of a warp share the path, +// so keep each depth's split feature and whether a deeper split on the same feature shadows it. +template +struct QuadratureInteractionSharedState : public QuadratureSharedState { + bst_feature_t splits[kGpuQuadratureWarpsPerBlock][DepthCap]; + std::uint8_t shadowed[kGpuQuadratureWarpsPerBlock][DepthCap]; + + [[nodiscard]] XGBOOST_DEV_INLINE bst_feature_t& Split(int warp, int depth) { + return splits[warp][depth]; + } + + [[nodiscard]] XGBOOST_DEV_INLINE bool Shadowed(int warp, int depth) const { + return static_cast(shadowed[warp][depth]); + } + + XGBOOST_DEV_INLINE void SetShadowed(int warp, int depth, bool value) { + shadowed[warp][depth] = static_cast(value); + } +}; + template struct QuadratureShapTaskRunner { Loader loader; @@ -772,31 +792,30 @@ struct QuadratureShapInteractionTaskRunner { atomicAdd(phis + out_idx, contrib); } + XGBOOST_DEV_INLINE void EnterPathNode(CompressedNode const& node, int depth) { + shared.Split(warp, depth) = node.split_global; + shared.SetShadowed(warp, depth, false); + if (node.prev_same_offset_plus1 != 0) { + shared.SetShadowed(warp, depth - static_cast(node.prev_same_offset_plus1) + 1, true); + } + } + + XGBOOST_DEV_INLINE void ExitPathNode(CompressedNode const& node, int depth) { + if (node.prev_same_offset_plus1 == 0) { + return; + } + subgroup.Sync(); + if (subgroup.is_warp_leader) { + shared.SetShadowed(warp, depth - static_cast(node.prev_same_offset_plus1) + 1, false); + } + subgroup.Sync(); + } + template - XGBOOST_DEV_INLINE void ForEachUniquePartner(CompressedNode const* nodes_for_tree, - int current_depth, bst_feature_t current_split, - Fn&& fn) const { - bool skipped_current = false; - for (int depth = current_depth; depth >= 0; --depth) { - auto const& candidate = nodes_for_tree[shared.Node(warp, depth)]; - if (candidate.is_leaf) { - continue; - } - auto split = candidate.split_global; - if (!skipped_current && split == current_split) { - skipped_current = true; - continue; - } - bool shadowed = false; - for (int newer = current_depth; newer > depth; --newer) { - auto const& newer_node = nodes_for_tree[shared.Node(warp, newer)]; - if (!newer_node.is_leaf && newer_node.split_global == split) { - shadowed = true; - break; - } - } - if (!shadowed) { - fn(depth, split); + XGBOOST_DEV_INLINE void ForEachUniquePartner(int current_depth, Fn&& fn) const { + for (int depth = current_depth - 1; depth >= 0; --depth) { + if (!shared.Shadowed(warp, depth)) { + fn(depth, shared.Split(warp, depth)); } } } @@ -836,20 +855,18 @@ struct QuadratureShapInteractionTaskRunner { auto diag_contrib = subgroup.Sum(edge_delta_local); this->AddDiagonalContribution(row_idx, tree_group, node.split_global, diag_contrib); - this->ForEachUniquePartner( - nodes_for_tree, parent_depth, node.split_global, - [&](int partner_depth, bst_feature_t partner_split) { - float q_partner = 1.0f; - if (subgroup.is_leader && subgroup.RowActive()) { - q_partner = shared.PathProbability(warp, subgroup.row_slot, partner_depth); - } - q_partner = subgroup.Broadcast(q_partner); - auto pair_delta_local = - ExtractQuadratureInteractionDeltaLocal(quad_node, edge_delta_local, q_partner); - auto pair_contrib = subgroup.Sum(pair_delta_local); - this->AddPairContribution(row_idx, tree_group, node.split_global, partner_split, - pair_contrib); - }); + this->ForEachUniquePartner(parent_depth, [&](int partner_depth, bst_feature_t partner_split) { + float q_partner = 1.0f; + if (subgroup.is_leader && subgroup.RowActive()) { + q_partner = shared.PathProbability(warp, subgroup.row_slot, partner_depth); + } + q_partner = subgroup.Broadcast(q_partner); + auto pair_delta_local = + ExtractQuadratureInteractionDeltaLocal(quad_node, edge_delta_local, q_partner); + auto pair_contrib = subgroup.Sum(pair_delta_local); + this->AddPairContribution(row_idx, tree_group, node.split_global, partner_split, + pair_contrib); + }); if (child_idx == 0) { auto child_weight = node.right_weight; @@ -882,6 +899,7 @@ struct QuadratureShapInteractionTaskRunner { (*stack_size)++; *have_return = false; } else { + this->ExitPathNode(node, parent_depth); *ret_val += shared.Basis(warp, subgroup.row_slot, parent_depth, subgroup.point); (*stack_size)--; *have_return = true; @@ -908,6 +926,7 @@ struct QuadratureShapInteractionTaskRunner { if (child == 0) { if (subgroup.is_warp_leader) { shared.Stage(warp, depth) = 1; + this->EnterPathNode(node, depth); } subgroup.Sync(); } @@ -1095,7 +1114,7 @@ __global__ void __launch_bounds__(kGpuQuadratureTreeBlockThreads, 9) static_assert(kGpuQuadratureSegmentWidth == kGpuQuadraturePoints, "Full-tile specialization assumes every warp lane participates."); } - using SharedT = QuadratureSharedState; + using SharedT = QuadratureInteractionSharedState; __shared__ SharedT shared; From 231155f81a8f425d2f29d98d839ac99d48ba937a Mon Sep 17 00:00:00 2001 From: Rory Mitchell Date: Tue, 6 Oct 2026 12:52:41 +0200 Subject: [PATCH 2/2] Fix warp synchronization in GPU SHAP traversal Finish shared path reads before reusing node and stage slots, and synchronize stage reads before the warp leader advances the traversal. Apply the ordering to SHAP values and interactions, and cover full and partial row tiles with a CPU/GPU regression test suitable for Racecheck. --- src/predictor/interpretability/shap.cu | 10 ++++++ tests/cpp/predictor/test_shap.cu | 43 ++++++++++++++++++++++++++ 2 files changed, 53 insertions(+) diff --git a/src/predictor/interpretability/shap.cu b/src/predictor/interpretability/shap.cu index fd98fd725fdf..431fca92dd0c 100644 --- a/src/predictor/interpretability/shap.cu +++ b/src/predictor/interpretability/shap.cu @@ -600,6 +600,9 @@ struct QuadratureShapTaskRunner { contrib = subgroup.Sum(contrib); this->AddContribution(row_idx, tree_group, node.split_global, contrib); + // Row subgroups can advance independently in a partial tile. Finish reading the + // current path before the warp leader reuses the child node and stage slots. + subgroup.Sync(); if (child_idx == 0) { auto child_weight = node.right_weight; auto child_node = node.right; @@ -660,6 +663,8 @@ struct QuadratureShapTaskRunner { // stage == 0 explores the left child first. After the return path updates the parent state, // the second visit uses the cached go-left decision to push the right child. int child = static_cast(shared.Stage(warp, depth) != 0); + // All lanes must read the old stage before the warp leader changes it. + subgroup.Sync(); if (child == 0) { if (subgroup.is_warp_leader) { shared.Stage(warp, depth) = 1; @@ -868,6 +873,9 @@ struct QuadratureShapInteractionTaskRunner { pair_contrib); }); + // Row subgroups can advance independently in a partial tile. Finish reading the + // current path before the warp leader reuses the child node and stage slots. + subgroup.Sync(); if (child_idx == 0) { auto child_weight = node.right_weight; auto child_node = node.right; @@ -923,6 +931,8 @@ struct QuadratureShapInteractionTaskRunner { } int child = static_cast(shared.Stage(warp, depth) != 0); + // All lanes must read the old stage before the warp leader changes it. + subgroup.Sync(); if (child == 0) { if (subgroup.is_warp_leader) { shared.Stage(warp, depth) = 1; diff --git a/tests/cpp/predictor/test_shap.cu b/tests/cpp/predictor/test_shap.cu index a9fa19823cda..f8140f922bc3 100644 --- a/tests/cpp/predictor/test_shap.cu +++ b/tests/cpp/predictor/test_shap.cu @@ -85,6 +85,49 @@ TEST(GPUPredictor, CompareCPUShap) { } } +// Exercise independently advancing row subgroups under compute-sanitizer racecheck. +TEST(GPUPredictor, ShapPartialRowTiles) { + auto ctx = MakeCUDACtx(0); + bst_feature_t constexpr kCols{3}; + bst_idx_t constexpr kTrainRows{128}; + auto train = RandomDataGenerator(kTrainRows, kCols, 0.0).GenerateDMatrix(); + train->Info().labels.Reshape(kTrainRows, 1); + auto& labels = train->Info().labels.Data()->HostVector(); + for (bst_idx_t i = 0; i < kTrainRows; ++i) { + labels[i] = static_cast(i % 7); + } + std::unique_ptr cpu{Learner::Create({train})}; + cpu->Configure({{"tree_method", "hist"}, {"max_depth", "5"}, {"min_child_weight", "0"}}); + cpu->Configure(); + for (int i = 0; i < 3; ++i) { + cpu->UpdateOneIter(i, train); + } + Json model{Object{}}; + cpu->SaveModel(&model); + std::unique_ptr gpu{Learner::Create({})}; + gpu->LoadModel(model); + gpu->Configure({{"device", ctx.DeviceName()}}); + gpu->Configure(); + + for (bst_idx_t rows : {1, 2, 3, 4, 5, 7, 8, 9}) { + SCOPED_TRACE(rows); + auto data = RandomDataGenerator(rows, kCols, 0.2).GenerateDMatrix(); + for (bool interactions : {false, true}) { + SCOPED_TRACE(interactions); + HostDeviceVector expected, actual; + cpu->Predict(data, false, &expected, 0, 0, false, false, !interactions, false, interactions); + gpu->Predict(data, false, &actual, 0, 0, false, false, !interactions, false, interactions); + ASSERT_TRUE(actual.DeviceCanRead()); + auto const& h_expected = expected.ConstHostVector(); + auto const& h_actual = actual.ConstHostVector(); + ASSERT_EQ(h_expected.size(), h_actual.size()); + for (std::size_t i = 0; i < h_expected.size(); ++i) { + ASSERT_NEAR(h_expected[i], h_actual[i], 1e-4); + } + } + } +} + TEST(GPUPredictor, ShapOutputCasesGPU) { auto ctx = MakeCUDACtx(0); auto cases = BuildShapTestCases(&ctx);