015 — native PyTorch and CuPy updates in the KMeans estimator

This entry adds backend-specific centroid updates to the portable KMeans estimator from 014 and reruns the full five-task benchmark through the public estimator. PyTorch uses index_add_ plus bincount; CuPy uses cupy.add.at plus bincount. Every other array API namespace keeps the portable one-hot update, and use_native_backends(False) forces the portable path everywhere so that each backend can be measured both ways on the same arrays.

The measured implementation is commit 04848dd:

• native_accumulate picks the backend callable, or None for the portable fallback;
• _torch_accumulate and _cupy_accumulate return per-cluster sums and counts;
• update_centroids_tiled calls the callable when one is given, and shares the empty-cluster relocation with the portable path through _relocate_empty.

def _torch_accumulate(X, labels, n_clusters):
    sums = torch.zeros((n_clusters, X.shape[1]), dtype=X.dtype, device=X.device)
    sums.index_add_(0, labels, X)
    counts = torch.bincount(labels, minlength=n_clusters).to(X.dtype)
    return sums, counts

Results: fit-results.json (300 full-fit rows), init-results.json (135 initializer rows), fit-first-run-array-api.json (the 180 array-API rows of the first run, superseded below), and the update crossover probes crossover-torch-cuda.json and crossover-cupy.json. The driver is bench_task_mix.py, the figures come from plot_estimator_benchmark.py, and the crossover from probe_update_crossover.py.

Setup

• Same five tasks, seeds, and limits as 011: float32, TF32 disabled, max_iter=12, tol=1e-8, seeds 0 to 2, 4 CPU threads, a fresh process per run, and a 60-second limit per worker. One NVIDIA L4.
• Array-API rows call KMeans(n_clusters=k, init=..., n_init=1).fit(X) on a tensor or array already on the device. Each run therefore includes input validation, the Lloyd iterations, a final assignment, and an inertia pass. cuML and scikit-learn receive NumPy input as before, so cuML timings include the host-to-device copy.
• "generic update" rows run the same estimator with use_native_backends(False).
• Two sections: greedy k-means++ followed by 12 Lloyd iterations, and 12 Lloyd iterations from the same random data rows for every implementation. The shared-start section includes estimator overhead that 011's Lloyd-only rows did not have.
• Memory meters: PyTorch and CuPy pool peaks on CUDA, VmHWM on CPU, and nvidia-smi for cuML.
• Initializer rows are unchanged code and come from the first run.

Two fixes found by the first run

The first full run, on commit f13c074, exposed two problems. Both are fixed in 04848dd and every array-API row was rerun.

Finiteness check memory. The estimator validated non-NumPy input with xp.all(xp.isfinite(X)). CuPy fuses that into one kernel, but eager PyTorch materializes a boolean array plus float temporaries as large as X. The check is now _all_finite: it tests column sums first and falls back to a tiled exact check only when a sum is not finite, which also covers finite data whose sums overflow.

Atomic contention at few clusters. On CUDA the scatter-add was slower than the portable update on the 3-cluster, 1000-feature task: Torch CUDA k-means++ fits took 1.56 s against 1.08 s generic. With 950000 samples and 3 clusters, each output address receives about 300000 atomic additions, while the one-hot product is a cheap bandwidth-bound GEMM. The probe below times one update with uniformly random labels on each task's data, sweeping k. Values are generic time divided by native time; italic cells are shapes where the final dispatch keeps the one-hot update.

The crossover depends on both k and the number of features d, and CuPy's add.at tolerates contention better than PyTorch's index_add_. The final dispatch keeps the one-hot update on CUDA when 16 * k < d for PyTorch and 64 * k < d for CuPy. Its worst misses in this sweep are at 2 clusters: PyTorch keeps the scatter-add at d=32, where it runs at 0.59x, and CuPy keeps the one-hot update at d=180, where the scatter-add would be 1.25x faster. On CPU the scatter-add won on every task, so Torch CPU always uses it. In the benchmark, only k3-d1000 falls on the one-hot side, so its CUDA "native" and "generic" rows run the same code.

Full fit after k-means++

Three cells have one slow seed, and the means above include it:

• CuPy k3-d1000: seed 0 took 5.38 s against 1.11 s for seeds 1 and 2. It was the first CuPy process after the code change and paid CuPy's one-time kernel compilation for the new validation reductions. A rerun with an empty CUPY_CACHE_DIR took 8.36 s, and 1.11 s once warm.
• Torch CPU k3-d1000 (seed 2: 27.5 s against 16.6 s) and the shared-start Torch CPU k100-d32 (seed 2: 18.1 s against 4.8 s): reruns of those exact cells took 16.5 s and 4.8 s, so these were transient host noise.

Quality is mean inertia divided by the best mean in the same section. The native and generic paths agree to the shown precision except where seeded PyTorch CUDA k-means++ draws already differ across processes, as noted in 011; CuPy rows match exactly.

The scatter-add removes the one-hot tile, which saves 6 to 10 MiB at 3000 clusters; elsewhere the assignment tiles dominate and memory is unchanged.

Lloyd from shared random rows

Every implementation reaches the same objective within 0.07% on all tasks, except cuML at 3000 clusters, which keeps empty centers and ends 6.99% above the best, as in 011.

Native speedup over the generic update

Generic time divided by native time, from the seed means above.

The two Torch CPU cells below 1x come from the noisy seed-2 runs above: without seed 2 they are 1.10x and 1.78x. The CuPy k3-d1000 k-means++ cell reflects the one-time kernel compilation.

Comparison with 011

011 called the Lloyd loop directly, with the earlier tile schedule and the one-hot update; 015 goes through the estimator with the tile schedule chosen in 014 and the native update. Seed means, k-means++ full fit:

The shared-start rows, 011 Lloyd-only against 015 estimator fit:

From 17 clusters up, most of the GPU gain over 011 comes from the tile schedule: the shared-start generic rows are already 1.6x to 1.8x faster than 011. The native update adds another 1.2x to 1.5x from 100 clusters up. At 3 clusters, 015 is about 0.1 s slower because the estimator adds validation, a final assignment, and an inertia pass over the 3.8 GB input.

Initialization only

Greedy k-means++ does not use the centroid update, so these rows only check that the initializer still matches 011.

Every row is within 0.2 s or 10% of 011. Quality and memory are in init-quality.svg and init-memory.svg; they repeat 011's conclusions.

Conclusions

• The native PyTorch and CuPy updates are now the default for those namespaces, with the portable update kept for every other namespace and for wide CUDA inputs with few clusters.
• On CUDA, the native update speeds up k-means++ fits by 1.06x to 1.21x and shared-start fits by 1.17x to 1.47x, from 17 clusters up on CuPy and from 100 clusters up on PyTorch. Assignment now dominates the remaining time.
• Torch CPU gains the most: 1.9x on the 550-cluster k-means++ fit, 4.6x on its shared-start fit, and the 3000-cluster k-means++ fit now finishes in 40 s instead of timing out.
• PyTorch CUDA and CuPy now beat cuML's shared-start Lloyd at 100 clusters (0.39 s against 0.54 s) and stay faster at 3 and 17 clusters. cuML remains faster at 550 and 3000 clusters, by 1.2x to 1.9x. For k-means++ fits, cuML greedy is faster at 550 and 3000 clusters and cuML scalable from 100 clusters up; greedy initialization alone costs 0.5 to 2.9 s on the array-API backends there.
• Two estimator-level fixes came out of the first run: a finiteness check without input-sized temporaries, which recovers 5.3 GiB of PyTorch memory on the widest task, and a backend-specific k-versus-d crossover for the scatter-add.

Previous: portable array-API KMeans. Next: GEMM versus broadcast distances. Index: journal.