Skip to content

[FEA] Multi-node Out of Core Streaming KMeans API#2066

Open
tarang-jain wants to merge 196 commits into
NVIDIA:mainfrom
tarang-jain:mnmg-streaming
Open

[FEA] Multi-node Out of Core Streaming KMeans API#2066
tarang-jain wants to merge 196 commits into
NVIDIA:mainfrom
tarang-jain:mnmg-streaming

Conversation

@tarang-jain

@tarang-jain tarang-jain commented May 7, 2026

Copy link
Copy Markdown
Contributor

Merge after #2015 and #2017

Allows a stream of input matrices per worker, that are further batched using the streaming_batch_size parameter. Reasoning: We should be able to supply dask partitions (on host) directly without having to concatenate them into one consolidated matrix.

As a part of this PR, we also unify the multi-GPU implementations into one (earlier the out of core implementation was separate).
Tests: We get rid of the separate out of core test file. The single MG testing unit is taking care of both out of core and on device matrices.

@tarang-jain

Copy link
Copy Markdown
Contributor Author

/ok to test d0fa3a9

@viclafargue viclafargue left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Great work! Pre-approving, just have two points that may need some attention.

Comment thread cpp/src/cluster/detail/kmeans_mg.cuh
Comment thread cpp/src/cluster/detail/kmeans_mg_batched_init.cuh
Comment thread cpp/src/cluster/detail/kmeans_mg_distributed_init.cuh Outdated
@tarang-jain

Copy link
Copy Markdown
Contributor Author

/ok to test 3321926

Comment thread cpp/src/cluster/detail/kmeans_mg_batched_init.cuh Outdated
Comment thread cpp/src/cluster/detail/kmeans_mg_batched_init.cuh
Comment thread cpp/src/cluster/detail/kmeans_mg_distributed_init.cuh Outdated
Comment thread cpp/src/cluster/detail/kmeans_mg_distributed_init.cuh Outdated
Comment thread cpp/src/cluster/detail/kmeans_mg_distributed_init.cuh Outdated
Comment thread cpp/CMakeLists.txt
@tarang-jain

Copy link
Copy Markdown
Contributor Author

/ok to test bd3e1de

@tarang-jain tarang-jain removed request for a team and msarahan June 23, 2026 13:55

@lowener lowener left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

@tarang-jain

Copy link
Copy Markdown
Contributor Author

/merge

@tarang-jain

Copy link
Copy Markdown
Contributor Author

/ok to test 40c77ca

@tarang-jain

Copy link
Copy Markdown
Contributor Author

/ok to test 079c216

@tarang-jain

Copy link
Copy Markdown
Contributor Author

/ok to test 7a38f62

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

breaking Introduces a breaking change feature request New feature or request

Projects

None yet

Development

Successfully merging this pull request may close these issues.

6 participants