Add weighted k-means and k-means++ init (stacked on #22)

[zeyuyang8]

Summary

Stacked on #22. Per @andy-yang-1's suggestion, the flattened-launch-grid fix is split into its own PR (#22). This PR contains only the opt-in feature work:

1. Weighted k-means. New Triton kernel triton_centroid_update_sorted_euclid_weighted + batch_kmeans_Euclid_weighted, exposed via FlashKMeans.fit(data, weights=...). Torch-native fallback supports weights too. The kernel streams D in BLOCK_D tiles with masking (matching _centroid_update_chunk_kernel), so non-power-of-two D works and large D + wide dtype stays within a bounded footprint.

2. k-means++ initialization. init={random, scalable-kmeans++, standard-kmeans++} (K-Means‖ matches cuML; standard matches scikit-learn), n_init best-of-N restarts by per-batch inertia.

Weighted centroid-update kernel — correctness + throughput (H100 80GB)

benchmarks/weighted/bench_weighted.py + tests/test_weighted_centroid.py (15 cases).

Correctness vs a pure-torch weighted-mean reference:

D	result
64 / 128 / 256 (pow2)	✅ max abs diff ≤ 6.1e-5 (fp16)
80 / 96 / 192 (non-pow2)	✅ max abs diff ≤ 1.2e-4 (fp16)
1024 (fp32)	✅ max abs diff 1.3e-7
all-ones weights == unweighted kernel	✅ exact
empty cluster → old_centroids	✅

Throughput (weighted vs unweighted centroid update), overhead ≤ 6%:

shape (B,N,K,D)	unweighted	weighted	overhead
(1, 1M, 256, 128)	3105 M/s	2976 M/s	+4.3%
(8, 131K, 256, 128)	2766 M/s	2632 M/s	+5.1%
(4, 262K, 512, 256)	2144 M/s	2122 M/s	+1.1%
(4, 262K, 512, 96)	2832 M/s	2675 M/s	+5.9%
(2, 262K, 512, 192)	1719 M/s	1711 M/s	+0.4%

scalable-kmeans++ quality fix

The scalable (K-Means‖) reduction used a plain non-greedy kmeans++, while the baseline is greedy — so on ~25% of instances scalable cost exceeded 2x standard (rare 9–10x blow-ups). Switching the reduction to the greedy weighted variant (matching cuML) removes the tail: across 20 seeds × {N=3200, 32000}, max ratio 10x → 1.2x, 0/20 exceed 2x. Fixes test_scalable_vs_sequential_quality.

Tests

pytest tests/ → 24 passed (k-means++ init + weighted-kernel correctness incl. non-pow2 D).

Note on the diff

Until #22 merges, this PR's diff against main also shows the grid-fix files — same commit as #22. After #22 lands I'll rebase and this shows feature-only.

[andy-yang-1]

@zeyuyang8 Thanks for your contribution!

A couple of suggestions:

Benchmarks. There aren't any perf numbers yet It'd be great to see (a) before/after throughput for the flattened-grid change on the existing euclid/cosine paths to confirm no regression, and (b) throughput plus a numerical correctness check for the new weighted centroid-update kernel.

Possibly splitting the PR. The flattened-grid fix is independent of the new features and has a wider blast radius — it modifies the shared _centroid_update_chunk_kernel and the non-split assign kernels that all current users hit, whereas weighted k-means and k-means++ are opt-in. It might be cleaner to land the grid fix as its own PR (easy to benchmark and diff against main in isolation) and keep the feature work separate. Not a hard blocker if you'd rather keep them together, but it would de-risk the shared-path change.

Glad to help benchmark on an H200 if useful.

[zeyuyang8]

Thanks for the review — both suggestions done, and digging in surfaced two real bugs I've also fixed.

Split. The grid fix is now its own PR, #22 (grid-only diff vs main, benchmarks included). This PR is stacked on it and holds just the opt-in weighted-kmeans / k-means++ work.

Benchmarks (H100 80GB, fp16). Scripts are in both PRs so they're reproducible.

Grid fix — no regression on the non-split euclid/cosine assign + centroid-update paths (AFTER/BEFORE latency ratio, worst case 1.02x):

shape (B,N,K,D)	euclid_assign	cosine_assign	euclid_update	cosine_update
(1, 1M, 256, 128)	0.99x	0.99x	1.00x	1.00x
(8, 131K, 256, 128)	1.00x	0.97x	1.02x	1.02x
(64, 16K, 256, 128)	1.01x	0.99x	1.01x	1.01x
(1, 4M, 1024, 64)	1.02x	1.01x	1.00x	1.00x
(256, 8K, 128, 128)	0.99x	0.98x	1.01x	1.01x
(4, 262K, 512, 256)	1.00x	1.01x	1.01x	1.01x

And it's a real bug fix: at B=70000 (> 65535) main fails all three kernels with Triton Error [CUDA]: invalid argument; the flattened grid launches fine.

Weighted kernel. While benchmarking I found the first version used a single untiled full-D load — it crashed on non-power-of-two D (96/192/…) and was ~48% slower at D=256. Fixed by streaming D in BLOCK_D tiles (matching _centroid_update_chunk_kernel). Now: correct vs a torch weighted-mean reference across pow2 + non-pow2 + fp32 (max abs diff ≤ 1.2e-4 fp16, 1.3e-7 fp32), all-ones weights == the unweighted kernel exactly, and throughput overhead is ≤ 6% (D=256: 48% → 1.1%).

scalable-kmeans++. Its reduction used plain (non-greedy) kmeans++ while the baseline is greedy, so ~25% of random instances blew up (rare 9–10x cost vs standard, independent of dataset size). Switched to the greedy weighted reduction (matching cuML) — across 20 seeds × {N=3200, 32000}, the max ratio drops from ~10x to 1.2x and scalable is on par with / slightly better than standard on average. pytest tests/ is now green (24 passed).

All numbers are on an H100 — I don't have an H200 handy, but benchmarks/grid_fix/ and benchmarks/weighted/ produce the same tables on one. If your H200 offer still stands, a cross-arch confirmation would be great. 🙏