[epic] Use custom CUDA stream for the entire codebase.
- Dominant language
- C++
- Stars
- 28.8k
- Forks
- 8.9k
- Avg merge
- 1d 12h
- Merged PRs (30d)
- 54
Description
This will be a long refactoring task. The objective is to enable the use of a custom CUDA stream to improve control over asynchronous memory allocation and to enable stream-specific device.
We have support for device ordinal `cuda:1`. This has been a pain point for XGBoost, yet it's a widely used feature. In CUDA 13, streams are implicitly attached to the device during creation. As a result, if we can use a custom stream, we can remove the C API guard and avoid initializing the CUDA context.
Plan:
- [ ] Provide optional context parameter in all storage class, including:
+ [ ] DeviceUVector
+ [ ] HostDeviceVector
+ [ ] Tensor
+ [ ] TemporaryArray
- [ ] Use stream-oriented memory allocation in the booster class.
- [ ] Use stream-oriented memory allocation in the DMatrix classes.
- [ ] Provide synchronization between the DMatrix and the booster.
- [ ] Remove the set device when a device ordinal is not provided.
- [ ] Remove C API guard. Verify that XGBoost doesn't initialize the CUDA context when CUDA is not used.
PRs:
- https://github.com/dmlc/xgboost/pull/12121
Related:
- https://github.com/dmlc/xgboost/issues/12116
- https://github.com/dmlc/xgboost/issues/11884
Contributor guide
No contributing guide indexed for this repository
Research direction
Start by reviewing PR #12121 and the checklist in this issue, beginning with the listed storage classes: DeviceUVector, HostDeviceVector, Tensor, and TemporaryArray. Trace the planned stream-oriented allocation and synchronization work across the booster and DMatrix classes. Done means the checklist is completed, including removing the device setting and C API guard and verifying CUDA is not initialized when unused.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp
- Domain
- machine-learning, performance
- Issue type
- Refactor
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Stale
- Clarity
- Needs clarification
- Newbie friendliness
- 20/100