Make GPU sketch prune emit a correct summary#12090
Conversation
There was a problem hiding this comment.
Pull request overview
This PR updates the GPU quantile sketch pruning path to avoid producing duplicated cut entries (especially under highly repeated input values) and extends the GPU quantile tests to validate the stronger post-prune invariants.
Changes:
- Refactors pruning in
SketchContainer::Pruneto select/gather pruned entries via index selection and a segmented-unique pass, instead of pruning then callingUnique(). - Renames the device-side binary search helper to return an index (
BinarySearchQueryIndex) and adapts pruning to use it. - Expands GPU quantile tests to assert per-column no-duplicate cut values after pruning, including a new repeated-values test case.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
src/common/quantile.cu |
Refactors GPU sketch pruning to de-duplicate by selected source indices and shrinks output accordingly. |
tests/cpp/common/test_quantile.cu |
Strengthens prune validation (per-column duplicate checks) and adds a test targeting repeated-value inputs. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| template <typename EntryFromIndex> | ||
| void GatherPruneEntries(Span<size_t const> selected_idx, Span<SketchEntry> out_cuts, | ||
| EntryFromIndex entry_from_index) { | ||
| dh::LaunchN(selected_idx.size(), |
There was a problem hiding this comment.
Please pass an explicit stream to the LaunchN. The context returns the default per-thread stream for now, but I plan to move torward self-managed CUDA stream in the future to aid device ordinal selection, among other benefits.
| return in[abs_idx]; | ||
| }; // NOLINT | ||
| SelectPruneIndices(d_columns_ptr_out, d_columns_ptr_in, ft, d_selected_idx, entry_from_index); | ||
| auto n_selected = dh::SegmentedUnique( |
There was a problem hiding this comment.
Not sure what the difference is after this PR? It's still running an unique op, and the existing one is still here.
There was a problem hiding this comment.
Turns out I can't remove unique yet I need a couple more steps. What does PR really does is separate out a prune implementation that kind of does 2 things, processing raw inputs and processing sketches, into two different implementations.
This is a focused cleanup in the GPU sketch
Prune(...)path.The main change is to separate two roles that were previously mixed together:
For the summary-side
Prune(...)path, this PR now:SketchEntryvalues into the output summaryThis is closer to the CPU
SetPrunestructure, where prune selection is based on selected records instead of immediately materializing output and then relying on that materialization for all cleanup.This PR does not remove the later value-based
Unique()step. It turns out that is still needed for now, especially in distributed / merged-summary cases. The value of this PR is separating out the summary prune implementation so later cleanup steps can be changed more safely and incrementally.This is also one part of a larger refactoring of the GPU sketching path. The broader direction is to separate:
Testing
make -C build-make -j35 testxgboost./build-make/testxgboost --gtest_filter='GPUQuantile.Prune:GPUQuantile.PruneDuplicated:GPUQuantile.MergeBasic:GPUQuantile.MultiMerge:Quantile.DistributedBasic:Quantile.Distributed'All 6 tests passed.