Lloyd k-means in the Python array API

What this PR explores

This experiment studies nearest-centroid assignment, Lloyd updates, empty clusters, and centroid initialization through the Python array API. The eager implementation runs with PyTorch on CPU, CUDA, and XPU, with CuPy on CUDA, and with dpnp on a Level Zero GPU. Comparisons cover stock scikit-learn, scikit-learn with k-means++ performance patches, Cython, and cuML.

The journal progresses from kernel-fusion probes to bounded-memory tiled assignment, full fits, repeated-seed and shared-start comparisons, empty-center relocation, and greedy versus scalable initialization. The current six-task mix spans 3 to 3000 clusters, 2 to 1000 features, and 135000 to 2100000 samples, including a balanced 256-cluster, 128-feature case. Runs use float32, eager execution, TF32 disabled, three seeds, and a 60-second worker limit.

Latest results

023 adds the portable array-API path on dpnp GPU arrays to the 022 Arc B390 measurements. 022 holds PyTorch XPU (native and one-hot), PyTorch CPU, stock scikit-learn 1.9.1, and cakedev0. CuPy and cuML are absent on this host.

Main conclusions:

• dpnp, which always uses the one-hot update, is faster than native Torch CPU on every full fit. It is also faster than the XPU one-hot update at 17 clusters (2.57 s against 2.83 s) and at 256 (1.11 s against 1.87 s).
• The XPU one-hot update remains the fastest full fit at 3 clusters (3.22 s; dpnp is 4.83 s). From 17 clusters up, cakedev0 is still the fastest full fit. dpnp is second at 17 clusters (2.57 s against 2.35 s) and at 256 (1.11 s against 1.06 s).
• At 256 clusters dpnp Lloyd takes 0.47 s, level with stock scikit-learn and cakedev0. At 550 and 3000 clusters both XPU updates and Cython Lloyd are faster than dpnp.
• On Arc B390 the one-hot PyTorch update is faster than index_add_ through 256 clusters (022). The CUDA few-cluster cutoff would keep index_add_ on k17-d180, where the one-hot update is the faster XPU path.

Journal

• 001 — fusion probe: eager and compiled distance-plus-argmin kernels.
• 002 — large-shape OOM: which schedules avoid the full distance matrix.
• 003 — three shapes: sample-, centroid-, and feature-heavy fusion behavior.
• 004 — tile trade-offs: bounded-memory tiles, Cython Lloyd, and k-means++.
• 005 — full k-means: eager array API, Cython, patched scikit-learn, and cuML.
• 006 — benchmark breadth: structure and cluster-count sensitivity across five implementations.
• 007 — repeated seeds and shared centers: uncertainty bars and initializer-independent Lloyd comparisons.
• 008 — empty-center task mix: stochastic relocation across five bounded tasks.
• 009 — k-means++ only: timing, memory, and quality before Lloyd.
• 010 — scalable initialization: farthest relocation, cuML initializer options, and array-API k-means||.
• 011 — consolidated benchmark: initialization, Lloyd-only, and full-fit results on one page, plus the scalable-init ablation.
• 012 — source link audit: source links pinned to the commits that each report measured.
• 013 — large-k Lloyd optimization: scalable-init sensitivity, cuML kernel analysis, optimization probes, and next steps.
• 014 — portable array-API KMeans: sklearn-compatible estimator extraction and CuPy-native centroid-update probes.
• 015 — native backend updates in the estimator: PyTorch and CuPy scatter-add updates, the CUDA contention crossover, and the full benchmark rerun.
• 016 — GEMM versus broadcast distances: time, memory, and quality across initialization, Lloyd, and full fit on CPU and CUDA.
• 017 — GEMM cancellation fallback: translated-data failures, threshold sensitivity, and selective direct-distance repair.
• 018 — backend-neutral CPU optimizations: labels-only assignment, tiling, interleaving, augmented GEMM, compilation, and sampling experiments.
• 019 — CPU and GPU optimization rerun: the 018 optimizations on three backends and five tasks, beside cakedev0 on the same host.
• 020 — balanced clusters and features: a k256-d128 task, stock-scikit-learn calibration, and a device-resident cuML comparison.
• 021 — device-resident cuML on the original five tasks: the same placement protocol as 020, applied to the cuML cells that 015 timed from host memory.
• 022 — PyTorch XPU on Arc B390: the six-task mix on an Intel Arc B390, with PyTorch CPU, stock scikit-learn 1.9.1, and cakedev0 on the same laptop.
• 023 — dpnp array API on Arc B390: the portable one-hot update on dpnp GPU arrays, beside the 022 measurements.