[FEA] Truly Decoupled Double Buffering for SG KMeans

[tarang-jain]

Profile: H100 + Intel(R) Xeon(R) Platinum 8480CL

Dataset mem type = pinned memory. This speeds up training by a lot. nsys gave throughput for H2D transfers around 50 GB/s. This is a 114 GB dataset. The screenshot above shows the overlap between the two streams with prefetch.

Metric	Baseline (2M batch)	Prefetch (2M batch)	Baseline (4M batch)
profile_mode	baseline	prefetch	baseline
dataset	base_falcon_1024_30M.fbin	base_falcon_1024_30M.fbin	base_falcon_1024_30M.fbin
rows	30,000,000	30,000,000	30,000,000
dimensions	1,024	1,024	1,024
clusters	2,048	2,048	2,048
batch_rows	2,000,000	2,000,000	4,000,000
requested_iterations	20	20	20
pinned_load_seconds	66.3182	59.2441	60.1955
initialization_seconds	0.0169335	0.0170245	0.0172878
fit_seconds	101.064	59.8278	101.737
fit speedup vs. 4M	1.007x	1.700x	1.000x
e2e_seconds	168.422	120.27	162.921
e2e speedup vs. 4M	0.967x	1.355x	1.000x
iterations	20	20	20
inertia	2.14326e+07	2.14325e+07	2.14325e+07

Also the use of pinned memory shows > 50 GB/s of throughput for H2D transfers. This gets much lower with regular host mem.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[viclafargue]

Great work, the reported speedup looks promising! To make the changes easier to review and validate, could we keep this PR focused on double-buffered prefetching? The workspace-resource selection, final-inertia optimization, and cluster-cost refactoring could be handled in separate follow-up PRs, allowing each change to be tested and benchmarked independently.

[tarang-jain]

The use of the detail namespace API is part of this optimization. The public API currently synchronizes the stream, so it is not allowing concurrent prefetch during inertia computation.