dmlc / dmlc/xgboost

[epic] Use custom CUDA stream for the entire codebase.

Open
#12,122 5 comments 2 reactions 0 assignees View on GitHub
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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.