006 — blob structure, centroid count, and five implementations

Same budget for every full fit: at most 12 Lloyd iterations, or stop when the center shift falls under 1e-8 times the mean feature variance. Float32. The host is the L4 box (4 CPUs, 15 GiB RAM). The array-API code is eager PyTorch, CPU or CUDA, with no torch.compile. The request listed PyTorch CPU under the GPU heading; the GPU array-API run is PyTorch CUDA eager.

Scikit-learn 1.7.2 is the stock build. The optimized build is scikit-learn 1.10.dev0 from cakedev0/scikit-learn#23 at 7e26a5c, which also contains the chunked Lloyd kernel from cakedev0/scikit-learn#26. cuML is 26.10.00. cuML has no public k-means++ function, so that row is KMeans(max_iter=0, init="k-means++"), which leaves the initial centers in place and then assigns labels. That assignment is included in the cuML init time.

Seeds are not a shared stream. The array-API draws use NumPy Generator(0). Both scikit-learn builds use random_state=0. cuML uses random_state=0 as well. On the smaller matrices the CPU and CUDA array-API draws stay identical. On the two largest they diverge, because a float rounding difference changes a weighted sample and every later center follows.

Peak memory is an increment. PyTorch reports allocator bytes above the input. Scikit-learn reports the rise in VmHWM on the first call. cuML reports nvidia-smi, whose step is 1 MiB; a 2.0 MiB increment is that step, not a measured workspace.

Cases

Blob data places k random centers, scales them by separation, and adds Gaussian noise of scale std. The ratio separation / std is the blob-iness. noise is standard normal. Lloyd timing always starts from the first k rows, so that duration is not mixed with the initializer.

case samples features clusters structure
tight 20000 32 8 separation 8, std 0.4
mild 20000 32 8 separation 5, std 1
loose 20000 32 8 separation 2.5, std 1.3
overlap 20000 32 8 separation 1, std 1.6
noise 20000 32 8 standard normal
tight-d128 20000 128 8 separation 8, std 0.4
noise-d128 20000 128 8 standard normal
k4, k16, k64, k128 20000 32 4 to 128 separation 5, std 1
tall 100000 16 8 separation 5, std 1
wide 40000 256 16 separation 5, std 1
mixed 80000 64 32 separation 5, std 1
speed-m 65536 512 64 separation 5, std 1
speed-l 131072 1024 128 separation 5, std 1
speed-xl 131072 2048 256 separation 5, std 1

Quality

Mean squared distance to the assigned center. For well separated blobs this lands on features * std² once every blob owns a center (32 for the mild family, 512 for speed-m).

Inertia against blob structure

Iterations used on the structure sweep

Init, mean squared distance:

case Array API - Torch CPU scikit-learn 1.7.2 scikit-learn PRs 23+26 Array API - Torch CUDA cuML
tight 9.86 9.85 9.85 9.86 9.90
mild 59.63 64.18 64.18 59.63 55.41
loose 129.60 98.05 98.05 129.60 126.58
overlap 147.78 136.13 136.13 147.78 145.79
noise 41.50 41.88 41.88 41.50 41.80
k16 62.68 61.88 61.88 62.68 96.51
k64 70.94 71.75 71.75 70.94 62.94
k128 75.60 70.06 70.06 75.60 62.65
mixed 126.76 189.29 189.29 126.76 128.37
speed-m 1335.12 1331.29 1336.85 1334.99 1016.84
speed-l 2381.82 2704.90 2699.24 2381.82 2375.48
speed-xl 4410.30 4722.77 5042.97 4736.51 4085.00

Fit, mean squared distance and iterations actually run:

case Array API - Torch CPU scikit-learn 1.7.2 scikit-learn PRs 23+26 Array API - Torch CUDA cuML
tight 5.13, 2 it 5.13, 2 it 5.13, 2 it 5.13, 2 it 5.13, 2 it
mild 31.89, 2 it 31.89, 2 it 31.89, 2 it 31.89, 2 it 31.89, 2 it
loose 73.65, 12 it 53.99, 2 it 53.99, 2 it 73.65, 12 it 72.43, 12 it
overlap 81.77, 12 it 81.77, 9 it 81.77, 9 it 81.77, 12 it 84.63, 12 it
noise 29.54, 12 it 29.55, 12 it 29.55, 12 it 29.54, 12 it 29.56, 12 it
tight-d128 20.46, 2 it 20.46, 2 it 20.46, 2 it 20.46, 2 it 20.46, 2 it
noise-d128 125.25, 12 it 125.24, 12 it 125.24, 12 it 125.25, 12 it 125.24, 12 it
k4 32.02, 2 it 32.02, 2 it 32.02, 2 it 32.02, 2 it 32.02, 2 it
k16 31.94, 2 it 31.94, 2 it 31.94, 2 it 31.94, 2 it 52.76, 12 it
k64 38.45, 12 it 36.90, 12 it 36.90, 12 it 38.45, 12 it 31.96, 2 it
k128 38.63, 12 it 36.52, 12 it 36.03, 9 it 38.63, 12 it 31.73, 2 it
tall 15.97, 2 it 15.97, 2 it 15.97, 2 it 15.97, 2 it 15.97, 2 it
wide 255.97, 2 it 255.97, 2 it 255.97, 2 it 255.97, 2 it 255.97, 2 it
mixed 63.99, 2 it 94.62, 12 it 94.62, 12 it 63.99, 2 it 63.99, 2 it
speed-m 675.90, 12 it 665.32, 11 it 677.33, 7 it 675.90, 12 it 511.36, 2 it
speed-l 1199.17, 10 it 1368.80, 10 it 1359.15, 11 it 1199.17, 12 it 1188.53, 6 it
speed-xl 2214.23, 3 it 2385.00, 4 it 2380.32, 3 it 2398.48, 4 it 2043.96, 2 it

Where the table shows one inertia repeated across a row, the relative gap is under 0.1%. That happens for tight, mild, noise, tight-d128, noise-d128, k4, tall, and wide. On those, every implementation uses the same small iteration count when the blobs are separated, and the full budget of 12 when the data is noise.

Disclaimer: several gaps are the initializer, and they are not all in one direction. Extending the array-API run on loose from 12 to 40 iterations only moved the inertia from 73.65 to 73.64, so the extra iterations are not going to catch scikit-learn's 53.99. That fit started from a worse k-means++ draw (129.60 against 98.05). The same pattern, with the opposite winner:

Empty clusters are not the cause of the loose gap: the array-API labeling at iteration 12 used every center. The implementation still does not relocate an empty center, which scikit-learn does. That difference did not appear on the rows where the inertias match.

Duration and memory

Calls below a few tens of milliseconds are real second-call times, but cuML's one-iteration measurement sits on a 40 ms floor on every small shape, so it is not resolved under that floor. The speed-m fit is the first row where several implementations take longer than a second.

Duration on the larger shapes

Fit memory and the centroid sweep

k-means++, seconds and peak increment:

case Array API - Torch CPU scikit-learn 1.7.2 scikit-learn PRs 23+26 Array API - Torch CUDA cuML
k64 0.078 s, 14.1 MiB 0.109 s, 4.5 MiB 0.015 s, 2.3 MiB 0.028 s, 2.8 MiB 0.041 s, 2.0 MiB
k128 0.158 s, 14.1 MiB 0.221 s, 4.5 MiB 0.028 s, 2.5 MiB 0.055 s, 2.8 MiB 0.039 s, 6.0 MiB
wide 0.079 s, 43.7 MiB 0.257 s, 10.4 MiB 0.041 s, 4.1 MiB 0.008 s, 40.2 MiB 0.041 s, 44.0 MiB
mixed 0.149 s, 24.3 MiB 0.284 s, 8.3 MiB 0.038 s, 5.9 MiB 0.015 s, 20.3 MiB 0.040 s, 24.0 MiB
speed-m 1.256 s, 69.3 MiB 2.696 s, 21.2 MiB 0.620 s, 6.2 MiB 0.066 s, 64.5 MiB 0.163 s, 134.0 MiB
speed-l 8.863 s, 135.3 MiB 24.113 s, 58.4 MiB 5.609 s, 9.8 MiB 0.370 s, 129.5 MiB 1.067 s, 524.0 MiB
speed-xl 41.786 s, 135.3 MiB 118.235 s, 135.0 MiB 24.074 s, 12.7 MiB 1.345 s, 129.5 MiB 3.944 s, 1042.0 MiB

One Lloyd iteration from the first k rows:

case Array API - Torch CPU scikit-learn 1.7.2 scikit-learn PRs 23+26 Array API - Torch CUDA cuML
k128 0.017 s, 27.3 MiB 0.010 s, 4.8 MiB 0.010 s, 4.8 MiB 0.002 s, 19.9 MiB 0.040 s, 2.0 MiB
wide 0.016 s, 11.3 MiB 0.038 s, 77.8 MiB 0.037 s, 78.0 MiB 0.002 s, 5.6 MiB 0.040 s, 2.0 MiB
speed-m 0.090 s, 38.7 MiB 0.141 s, 255.9 MiB 0.142 s, 255.9 MiB 0.003 s, 33.3 MiB 0.040 s, 130.0 MiB
speed-l 0.592 s, 136.2 MiB 0.701 s, 1023.9 MiB 0.700 s, 1024.0 MiB 0.008 s, 130.5 MiB 0.162 s, 518.0 MiB
speed-xl 1.855 s, 164.2 MiB 2.029 s, 2048.0 MiB 2.035 s, 2048.0 MiB 0.026 s, 130.6 MiB 0.285 s, 1036.1 MiB

Full fit:

case Array API - Torch CPU scikit-learn 1.7.2 scikit-learn PRs 23+26 Array API - Torch CUDA cuML
k128 0.386 s, 41.3 MiB 0.246 s, 7.9 MiB 0.048 s, 5.3 MiB 0.069 s, 20.6 MiB 0.040 s, 6.0 MiB
wide 0.135 s, 49.9 MiB 0.335 s, 78.0 MiB 0.077 s, 78.0 MiB 0.013 s, 59.5 MiB 0.040 s, 44.0 MiB
mixed 0.188 s, 52.5 MiB 0.342 s, 38.9 MiB 0.096 s, 39.0 MiB 0.023 s, 40.0 MiB 0.040 s, 24.0 MiB
speed-m 2.414 s, 74.9 MiB 3.174 s, 255.7 MiB 0.882 s, 255.9 MiB 0.099 s, 76.1 MiB 0.163 s, 134.0 MiB
speed-l 15.240 s, 160.7 MiB 27.135 s, 1023.9 MiB 7.912 s, 1023.9 MiB 0.509 s, 132.0 MiB 1.141 s, 524.0 MiB
speed-xl 49.614 s, 175.0 MiB 122.993 s, 2048.0 MiB 30.310 s, 2048.1 MiB 1.517 s, 133.6 MiB 4.000 s, 1042.0 MiB

On speed-xl (131072 samples, 2048 features, 256 clusters, a 1.00 GiB matrix) one Lloyd step is 0.026 s on CUDA, 0.285 s in cuML, 1.855 s for the array API on CPU, and 2.03 s for both scikit-learn builds. The scikit-learn step peaks 2.00 GiB above the input, two extra copies of the matrix. The array-API CPU step peaks at 164 MiB and the CUDA step at 131 MiB. cuML peaks about 1.02 GiB above the memory already in use.

The full fit on that shape is 1.52 s for array-API CUDA, 4.00 s for cuML, 30.3 s for the patched scikit-learn, 49.6 s for array-API CPU, and 123 s for stock scikit-learn. cuML's fit is the one that sits on the blob floor. The patched initializer is the fastest of the CPU inits (24.1 s against 41.8 s for the array API and 118 s for stock scikit-learn) and its draw is the worst of the five.

A PyTorch CPU matmul of this width uses the MKL build with 4 threads. The per-center product is a short N dimension, and the CPU initializer remains slower than the patched scikit-learn kernel. No further rewrite beat that gap inside the array API.