[ogrisel]

Main results

The non-cuML cells are the estimator benchmark in journal 015. The cuML cells for those five tasks are journal 021: the same fits, with the training data and the shared random-row centers already on the device before the timer starts. The k256-d128 column, including its cuML cells, is journal 020: 150,000 × 128 samples and 256 clusters, chosen so both k and d are substantial while stock scikit-learn still finishes (6.60 s on the calibration fit). Every row in the L4 tables ran on one host: 4 CPU threads, one NVIDIA L4, float32, TF32 off, three seeds, max_iter=12, tol=1e-8, and a 60-second worker limit. The other tasks are k3-d1000 (950,000 × 1000), k17-d180 (2,100,000 × 180), k100-d32 (1,500,000 × 32), k550-d6 (525,000 × 6), and k3000-d2 (135,000 × 2).

Two array-API settings share the same estimator and the same arrays:

• Native uses a backend-specific centroid update: PyTorch index_add_ + bincount, or CuPy add.at + bincount. On CUDA the estimator keeps the portable update when the input is wide and k is small (16k < d for PyTorch, 64k < d for CuPy). In this mix that is only k3-d1000, so the two CUDA rows there execute the same update. k256-d128 uses the native update on both CUDA backends.
• Generic forces use_native_backends(False): the portable one-hot update on every namespace, including PyTorch and CuPy.

Array-API rows call KMeans.fit on data that is already on the device. Stock scikit-learn is 1.7.2. The cakedev0 build is scikit-learn 1.10.dev0 with PR 23 and PR 26. cuML is 26.10. Every cuML cell below uses the same placement as the array-API GPU rows: samples and shared centers are CuPy arrays before timing, and output_type="cupy" keeps the fitted centers on the device until after the timed region. Journal 015's cuML numbers included the host-to-device copy; journal 021 replaces those five tasks, and journal 020 already used the device-resident protocol for k256-d128. Values are seed means. In both timing tables the CPU rows come first (Torch CPU, stock scikit-learn, cakedev0), then the GPU rows (Torch CUDA, CuPy, cuML). Bold marks the fastest time on that task inside its hardware group. Tied minima are both bold.

Full fit: greedy k-means++ plus 12 Lloyd iterations

implementation	k3-d1000	k17-d180	k100-d32	k256-d128	k550-d6	k3000-d2
Torch CPU, native	20.24 s	11.08 s	16.86 s	5.81 s	17.94 s	40.03 s
Torch CPU, generic	18.23 s	11.74 s	20.06 s	6.61 s	34.15 s	timeout
scikit-learn 1.7.2	12.29 s	10.52 s	13.22 s	6.61 s	16.91 s	29.24 s
cakedev0 PRs 23+26	7.40 s	5.02 s	2.94 s	1.66 s	3.25 s	4.52 s
Torch CUDA, native	1.07 s	0.88 s	0.79 s	0.49 s	1.14 s	2.00 s
Torch CUDA, generic	1.04 s	0.86 s	0.86 s	0.51 s	1.23 s	2.14 s
CuPy, native	2.53 s	0.92 s	1.09 s	0.48 s	1.76 s	3.43 s
CuPy, generic	1.12 s	1.11 s	1.27 s	0.50 s	1.93 s	3.64 s
cuML 26.10, greedy	3.85 s	0.94 s	0.68 s	0.51 s	0.88 s	1.06 s
cuML 26.10, scalable	3.24 s	0.83 s	0.44 s	0.24 s	0.51 s	1.06 s

Three published means in the original non-cuML columns include one slow seed. Reruns of those cells: Torch CPU k3-d1000 is 16.5 s, and the native-over-generic ratio there is 1.10×; CuPy k3-d1000 is 1.11 s once the kernel cache is warm (the 2.53 s mean is a one-time compile, and that shape uses the portable update anyway). The k256-d128 CuPy and cuML seed-0 cells are likewise warm reruns; the first process paid 10.2 s and 3.6 s of one-time setup.

cuML scalable on k3-d1000 stops after 2 iterations on seed 1 (1.31 s). A repeat of that seed stops after 2 iterations again (1.30 s). Seeds 0 and 2 run all 12 iterations at 4.20 s and 4.21 s. The 3.24 s mean includes the early stop, as the journal 015 mean of 6.38 ± 1.67 s did.

Lloyd only, from the same random rows

This section removes the initializer. Stock scikit-learn and cakedev0 take the same time: cakedev0’s full-fit lead is k-means++, from PR 23, rather than the chunked Lloyd kernel from PR 26. On k256-d128 the two CPU builds tie at 0.50 s.

implementation	k3-d1000	k17-d180	k100-d32	k256-d128	k550-d6	k3000-d2
Torch CPU, native	14.81 s	7.42 s	9.24 s	2.03 s	4.41 s	20.86 s
Torch CPU, generic	16.30 s	8.09 s	8.56 s	2.77 s	20.34 s	44.45 s
scikit-learn 1.7.2	6.79 s	3.17 s	1.42 s	0.50 s	1.32 s	1.54 s
cakedev0 PRs 23+26	7.02 s	3.18 s	1.44 s	0.50 s	1.31 s	1.55 s
Torch CUDA, native	0.89 s	0.64 s	0.39 s	0.22 s	0.49 s	0.62 s
Torch CUDA, generic	0.89 s	0.63 s	0.46 s	0.23 s	0.59 s	0.75 s
CuPy, native	1.06 s	0.67 s	0.39 s	0.19 s	0.42 s	0.58 s
CuPy, generic	1.00 s	0.87 s	0.58 s	0.21 s	0.60 s	0.76 s
cuML 26.10	3.69 s	0.53 s	0.20 s	0.12 s	0.20 s	0.16 s

The Torch CPU k100-d32 native mean is one noisy seed (18 s against 4.8 s). A rerun of that seed is 4.8 s, and the native-over-generic ratio is then 1.78×.

CPU

The native update is a large win where many clusters make the one-hot product expensive, and a small one where assignment dominates.

• Shared-start Lloyd, native over generic, on the original five tasks: 1.10×, 1.09×, 1.78× (after the rerun above), 4.61×, and 2.13×. On k256-d128 it is 1.37× (2.03 s against 2.77 s), and the full fit is 1.14× (5.81 s against 6.61 s). The 3000-cluster k-means++ fit finishes in 40 s with the native update and exceeds 60 s with the portable update.
• Against stock scikit-learn 1.7.2, native Torch CPU full fits stay slower on the original five tasks, about 1.1–1.6× (using the 16.5 s rerun at k3-d1000). On k256-d128 the native fit is faster: 5.81 s against 6.61 s. Shared-start Lloyd stays slower on every task, including 2.03 s against 0.50 s at k256-d128; the original five are about 2.1×, 2.3×, 3.3×, 3.4×, and 13.5× slower than both scikit-learn builds.
• Against cakedev0, native Torch CPU full fits are 2.2–8.9× slower on the original five tasks and 3.5× slower on k256-d128 (5.81 s against 1.66 s). Greedy k-means++ does not use the centroid update, and that is where cakedev0 pulls away on the original tasks: Torch CPU init takes 2.20 / 3.62 / 12.73 / 13.79 / 21.54 s, stock scikit-learn 5.14 / 7.65 / 11.45 / 15.34 / 28.06 s, and cakedev0 1.30 / 2.04 / 1.53 / 1.96 / 3.00 s. The portable initializer already beats stock scikit-learn on four of those five tasks and is 1.7–8× slower than cakedev0. On k256-d128 the Lloyd kernels of the two scikit-learn builds tie, so the full-fit gap is again initialization.

GPU

On CUDA the native update adds a modest factor on top of a GEMM assignment, from 17 clusters up on CuPy and from 100 clusters up on PyTorch: shared-start speedups are 1.17–1.22× for Torch CUDA and 1.29–1.47× for CuPy on the original tasks, and 1.07× and 1.12× on k256-d128. Below that crossover the portable update is as fast, which is why k3-d1000 stays on it.

• Against stock scikit-learn, native Torch CUDA full fits are 11–17× faster on the original tasks (1.07 s against 12.29 s at 3 clusters, 2.00 s against 29.24 s at 3000) and 13× faster on k256-d128 (0.49 s against 6.61 s). Warm CuPy is 8–11× faster than stock scikit-learn on the original fits and 14× faster on k256-d128.
• Against cakedev0 on the host CPU, native Torch CUDA full fits are 6.9×, 5.7×, 3.7×, 2.9×, and 2.3× faster on the original tasks, and 3.4× faster on k256-d128 (0.49 s against 1.66 s). Warm CuPy is 6.7×, 5.5×, 2.7×, 1.8×, and 1.3× faster on the original tasks, and 3.5× faster on k256-d128.
• Against cuML greedy, with both sides device-resident: native Torch CUDA full fits are 3.6× and 1.07× faster at 3 and 17 clusters (1.07 s against 3.85 s, and 0.88 s against 0.94 s). cuML greedy is faster from 100 clusters up on the original tasks (0.68 s against 0.79 s, 0.88 s against 1.14 s, and 1.06 s against 2.00 s). On shared-start Lloyd, Torch CUDA still wins at 3 clusters (0.89 s against 3.69 s). cuML is faster from 17 clusters up: 0.53 s against 0.63 s, 0.20 s against 0.39 s, 0.20 s against 0.42 s, and 0.16 s against 0.58 s. Moving the copy out of the timer removes about 3.2 s on the 1000-feature task and about 1.6 s on the 180-feature task, which reverses the Lloyd ranking at 17 and 100 clusters. The remaining gap at 3 clusters is the Lloyd kernel.
• cuML scalable is the fastest GPU full fit from 17 clusters through 550 (0.83 s, 0.44 s, 0.24 s, and 0.51 s) and ties cuML greedy at 3000 clusters (both 1.06 s). At 3 clusters it is 3.24 s, still behind Torch CUDA generic at 1.04 s.
• On k256-d128 the copy was already excluded. Native CuPy is the fastest greedy full fit, at 0.48 s against 0.49 s for native Torch CUDA and 0.51 s for cuML greedy. Shared-start Lloyd goes the other way: cuML takes 0.12 s, native CuPy 0.19 s, and native Torch CUDA 0.22 s. cuML scalable is the fastest full fit at 0.24 s and ends 0.66% above the best greedy inertia.

Greedy k-means++ on Torch CUDA is within about 1.7× of cuML’s greedy initializer on every original task (faster at 17 clusters: 0.36 s against 0.40 s). Those initializer timings already started from device-resident data in journal 015, so they are unchanged. CuPy matches that at 3 and 17 clusters and falls behind as k grows (2.93 s against 0.92 s at 3000 clusters). At 100 clusters the full-fit gap is both pieces: CuPy initialization is 0.74 s against 0.48 s, and CuPy Lloyd is 0.39 s against 0.20 s, so the full fit is 1.09 s against 0.68 s.

Quality and memory

Greedy full fits on the original tasks land within about 0.4% of the best mean inertia. Shared-start Lloyd agrees within 0.07%, except cuML at 3000 clusters, which keeps empty centers and finishes about 7% higher than the other implementations. Native and generic inertias match at the reported precision. On k256-d128, greedy full fits stay between 118.919 and 118.930 inertia per point, and shared-start Lloyd agrees within 0.0005 around 118.960. cuML scalable finishes at 119.708, 0.66% higher. On the original tasks, device-resident cuML scalable takes 0.83 / 0.44 / 0.51 / 1.06 s from 17 clusters up and ends 5.4% above its own greedy inertia at 3000 clusters.

On k3-d1000 the array-API full fit peaks about 0.94 GiB above the input. Stock scikit-learn and cakedev0 peak about 7.1 GiB. Journal 015 recorded about 3.8 GiB of nvidia-smi growth for cuML there, which included allocating the device copy of the 3.54 GiB input. With that copy done before timing, the same greedy fit grows nvidia-smi by 54 MiB. The other device-resident greedy increments are 136 MiB, 98 MiB, 56 MiB, and 14 MiB. The native update removes the one-hot tile and saves 6–10 MiB at 3000 clusters; elsewhere the assignment tiles dominate. The k256-d128 input is about 73 MiB. Array-API full fits peak about 147 MiB above their allocator baseline, and both scikit-learn builds peak about 146 MiB above VmHWM.

Arc B390

Journal 022 repeats this mix on a laptop with an Intel Arc B390, and journal 023 adds dpnp on that GPU. These rows are a different host from the L4 tables above: Intel Core Ultra X7 358H, 30 GiB RAM, integrated Arc B390. Float32, PyTorch matmul precision highest, three seeds, max_iter=12, tol=1e-8. CPU rows use 4 threads pinned to CPUs 0–3. XPU and dpnp workers are unpinned. Stock scikit-learn here is 1.9.1. cakedev0 is the same commit as above. On XPU, native is index_add_ + bincount at every shape, and generic forces the portable one-hot update. dpnp has no keyed reduction, so native and generic execute that same portable update and the table has one dpnp row. Bold is the fastest time in the CPU group and in the GPU group.

Full fit: greedy k-means++ plus 12 Lloyd iterations

implementation	k3-d1000	k17-d180	k100-d32	k256-d128	k550-d6	k3000-d2
Torch CPU, native	4.99 s	3.82 s	5.12 s	1.85 s	5.98 s	17.06 s
Torch CPU, generic	5.01 s	4.16 s	5.79 s	2.32 s	7.61 s	20.49 s
scikit-learn 1.9.1	33.12 s	7.31 s	8.95 s	5.72 s	10.26 s	13.67 s
cakedev0 PRs 23+26	9.01 s	2.35 s	1.92 s	1.06 s	2.18 s	3.01 s
Torch XPU, native	4.45 s	3.36 s	2.98 s	2.05 s	4.02 s	4.88 s
Torch XPU, generic	3.22 s	2.83 s	2.78 s	1.87 s	4.07 s	4.97 s
dpnp, portable	4.83 s	2.57 s	2.71 s	1.11 s	5.74 s	11.28 s

• Torch XPU, native full fits take 4.45 / 3.36 / 2.98 / 2.05 / 4.02 / 4.88 s, and shared-start Lloyd takes 4.18 / 2.95 / 2.24 / 1.66 / 2.68 / 3.08 s.
• Torch XPU, generic is the faster XPU full fit through 256 clusters (3.22 s against 4.45 s at 3 clusters) and the fastest full fit on this host at 3 clusters; at 550 and 3000 clusters native is ahead by less than 0.1 s.
• dpnp full fits take 4.83 / 2.57 / 2.71 / 1.11 / 5.74 / 11.28 s: faster than native Torch CPU on every task, faster than Torch XPU generic at 17 and 256 clusters, and second to cakedev0 there (2.57 s against 2.35 s, and 1.11 s against 1.06 s).

Lloyd only, from the same random rows

implementation	k3-d1000	k17-d180	k100-d32	k256-d128	k550-d6	k3000-d2
Torch CPU, native	4.55 s	2.66 s	1.69 s	1.13 s	1.92 s	10.62 s
Torch CPU, generic	4.33 s	2.83 s	2.34 s	1.57 s	3.37 s	13.23 s
scikit-learn 1.9.1	8.53 s	1.80 s	0.88 s	0.47 s	1.13 s	1.59 s
cakedev0 PRs 23+26	8.64 s	1.80 s	0.95 s	0.46 s	1.20 s	1.63 s
Torch XPU, native	4.18 s	2.95 s	2.24 s	1.66 s	2.68 s	3.08 s
Torch XPU, generic	3.15 s	2.23 s	2.05 s	1.48 s	2.73 s	3.29 s
dpnp, portable	4.20 s	2.11 s	1.95 s	0.47 s	3.56 s	4.83 s

• Torch XPU, generic Lloyd is fastest at 3 clusters (3.15 s); native Lloyd is faster at 550 and 3000 clusters (2.68 s and 3.08 s).
• dpnp Lloyd matches both scikit-learn builds at 256 clusters (0.47 s) and is ahead of Torch XPU generic there (1.48 s); from 17 clusters up, Cython Lloyd is faster on the other tasks.

Later, still experimental, backend-neutral changes

Journal 019 reran the 018 prototypes on Torch CPU, Torch CUDA, and CuPy. Those changes are an augmented [X, 1] GEMM and a labels-only assignment. They sit on top of the native update, and that run does not include stock scikit-learn, cuML, or k256-d128.

Augmented GEMM speeds Lloyd when k >= d: 1.25–3.15× on CPU, 1.37–1.79× on Torch CUDA, and 1.18–1.38× on CuPy, and is neutral below that. It stores another input-sized array, raising the k3-d1000 full-fit peak from about 0.95 GiB to 4.5 GiB. After these changes the best eager CPU method is still 2.1–7.0× slower than cakedev0 on init, 2.0–4.3× on Lloyd, and 2.1–6.1× on full fit (the pre-018 Lloyd gap at 3000 clusters was 13.5×). The best Torch CUDA full fit in that run is 2.9–10× faster than cakedev0, and the best CuPy fit is 1.4–8× faster. Augmented GEMM is not the estimator default: it needs a memory budget and the cancellation repair measured in journal 017.

Journal 016 found that the GEMM distance wins every Lloyd and full-fit cell against a bounded broadcast, which raised CUDA peak memory by up to 3.9×. Journal 017 found that a translation of 1e3 makes plain float32 GEMM mislabel about 97% of rows; cakedev0’s relative fallback restores the labels, at a 1.11–1.91× cost on ordinary data.

What the branch adds

• A scikit-learn-compatible dense KMeans estimator in kmeans_lloyd/src/kmeans_lloyd/kmeans.py. fit is the entry point. It keeps NumPy, PyTorch, CuPy, and strict array-API namespaces and devices, with random initialization and Lloyd iterations.
• Greedy k-means++ is kmeans_plusplus in kmeans_lloyd/src/kmeans_lloyd/iteration.py.
• Native centroid updates for PyTorch and CuPy, the portable one-hot update for every other namespace, and the CUDA k-versus-d cutoff above.
• A finiteness check that tests column sums first. The previous xp.isfinite(X) allocated 5.3 GiB of PyTorch temporaries on the 3.8 GB input.
• The benchmark mix now includes k256-d128. _worker_cuml starts the cuML timer after the training data and shared centers are resident CuPy arrays. Journal 021 applies that contract to the five tasks whose cuML cells in journal 015 included the copy.
• Estimator checks, strict array-API checks, optional CUDA tests, and the journals linked below.

Journal

Rendered index: https://ogrisel.github.io/cursor-sandbox/pr-17/

• 023 — dpnp array API on Arc B390
• 022 — PyTorch XPU on Arc B390
• 021 — device-resident cuML on the original five tasks
• 020 — balanced k256-d128, with device-resident cuML input
• 015 — native versus portable updates, with scikit-learn, cakedev0, and cuML
• 011 — consolidated benchmark before the native update
• 019 — CPU and GPU rerun of the backend-neutral prototypes beside cakedev0
• 016 — GEMM versus broadcast
• 017 — catastrophic cancellation
• 018 — backend-neutral CPU prototypes

  

[cakedev0]

Also, the benchmark cases don't include large k & large d cases. Which is a case I had in my previous position, and I struggled with the speed of k-means at the time. I was clustering embeddings of collaborative filtering models (dim 64 usually) and of texts (dim ~100-1000 depending on the model and if PCA was applied beforehand).

[ogrisel]

Good catch. I will try to spawn a VM to update my bench repo with consolidated timings.

[ogrisel]

Also, the benchmark cases don't include large k & large d cases.

I agree. I will also add that to the mix.

[cakedev0]

FYI, I opened two RAFT PRs that speed up cuML k-means with n_clusters <= 4: NVIDIA/raft#3194 (small change) and NVIDIA/raft#3195 (new kernel, faster in more cases), alternatives to each other. On k3-d1000, 84% of cuML's GPU time goes to RAFT's centroid update kernel for nkeys <= 4 (reduce_rows_by_key), which reads row-major data with uncoalesced loads: ~250 ms per Lloyd iteration on the L4, against ~36 ms for the distance + argmin pass.

Expected effect on the 015 tasks, estimated from kernel timings on the L4 (I didn't rebuild cuVS to measure it end to end):

• k3-d1000: the cuML fit should go from ~3.6 s to ~0.8 s with either PR. Avoid unintended FP64 promotion in float kernels NVIDIA/raft#3133, already merged in RAFT but not in the 26.10 wheels, should bring it to ~2 s on its own.
• The other tasks are unaffected by these two PRs.

NVIDIA/raft#3176 (open, by an NVIDIA engineer) optimizes the same update step for more than 4 keys. On the 015 tasks it accounts for 18% of cuML's GPU time at k17-d180, 9% at k100-d32 and ~1% at 550 and 3000 clusters, so that's the most it can save there. The rest is almost all the fused distance + argmin kernel.

[ogrisel]

@cakedev0 the results table in the descriptions to account for the problems you found in the prior benchmarking code.

[cakedev0]

It matches roughly what my agent measured 👍

[cakedev0]

I ran the GPU rows of the task mix on an Intel Arc B580 (xe driver, Ubuntu 26.04, Core Ultra 9 285) with torch XPU and dpnp. The estimator code is this PR as of 4bc8139, unchanged. All 108 fits (6 tasks, 2 inits, 3 seeds, 3 implementations) finished with 12 iterations, and the inertias match the L4 torch-cuda rows of journal 015 to 3 or 4 significant digits.

The driver imports bench_task_mix.py for the tasks, data, shared random starts and inertia, and keeps its protocol: fresh process per fit, 60 s limit, 4 threads, float32, max_iter=12, tol=1e-8, data on the device before timing. Two differences:

• measure_call uses CUDA events, so the timer is device synchronize + perf_counter.
• On XPU the first fit in a process includes ~1 s of one-time kernel compilation, so each worker also times a second identical fit ("warm").

dpnp has no native update in backends.py, so it always runs the one-hot update. Versions: torch 2.14.1+xpu, dpnp 0.20.0, dpctl 0.22.1, scikit-learn 1.9.1, array-api-compat 1.16.0, numpy 2.5.4. Values are seed means. The L4 row is copied from the PR description for reference: it is a first fit timed with CUDA events, so it's not directly comparable to the warm rows.

Full fit: greedy k-means++ plus 12 Lloyd iterations

implementation	k3-d1000	k17-d180	k100-d32	k256-d128	k550-d6	k3000-d2
Torch XPU, native, first fit	1.79 s	1.70 s	1.66 s	1.30 s	2.20 s	2.75 s
Torch XPU, generic, first fit	1.54 s	1.54 s	1.65 s	1.31 s	2.38 s	2.83 s
dpnp (generic), first fit	1.23 s	1.30 s	1.35 s	0.72 s	2.88 s	6.36 s
Torch XPU, native, warm	0.80 s	0.69 s	0.60 s	0.20 s	0.91 s	1.53 s
Torch XPU, generic, warm	0.53 s	0.51 s	0.60 s	0.21 s	1.03 s	1.68 s
dpnp (generic), warm	1.10 s	1.13 s	1.23 s	0.53 s	2.53 s	6.14 s
L4 Torch CUDA, native (PR description)	1.07 s	0.88 s	0.79 s	0.49 s	1.14 s	2.00 s

Lloyd only, from the same random rows

implementation	k3-d1000	k17-d180	k100-d32	k256-d128	k550-d6	k3000-d2
Torch XPU, native, first fit	1.45 s	1.33 s	1.14 s	0.89 s	1.35 s	1.60 s
Torch XPU, generic, first fit	1.19 s	1.16 s	1.14 s	0.89 s	1.55 s	1.67 s
dpnp (generic), first fit	1.11 s	1.07 s	0.96 s	0.32 s	1.67 s	2.30 s
Torch XPU, native, warm	0.73 s	0.59 s	0.38 s	0.06 s	0.44 s	0.57 s
Torch XPU, generic, warm	0.46 s	0.41 s	0.37 s	0.08 s	0.55 s	0.72 s
dpnp (generic), warm	0.99 s	0.96 s	0.93 s	0.17 s	1.47 s	1.98 s
L4 Torch CUDA, native (PR description)	0.89 s	0.64 s	0.39 s	0.22 s	0.49 s	0.62 s

What I take from this:

• The native vs one-hot crossover also exists on XPU, but native_accumulate only keeps the one-hot update when X.device.type == "cuda", so XPU always gets index_add_. Warm Lloyd, one-hot is 1.6x faster on k3-d1000 (0.46 s vs 0.73 s) and 1.4x on k17-d180 (0.41 s vs 0.59 s). k17-d180 would still take the native path under the CUDA rule (16 * 17 > 180), so the XPU threshold looks higher. index_add_ wins at large k (1.25x on k550-d6, 1.26x on k3000-d2).
• Torch XPU pays ~0.7 to 1.3 s of compilation on the first fit in a process. dpnp pays less (~0.1 to 0.35 s), which is why it wins most first-fit cells, but once warm it is 2x to 3.5x slower than the faster Torch XPU update (4x on the k3000-d2 full fit, mostly k-means++).
• Warm Torch XPU on the B580 is in the same range as the L4 numbers, but with a different timer and warm vs first fit I wouldn't read more into it than that.

Torch XPU peak extra device memory (from torch.xpu.max_memory_allocated) is the same for native and generic: 928 MB, 2052 MB, 368 MB, 148 MB, ~75 MB and ~70 MB in the column order above (Lloyd only).

Driver (bench_task_mix_xpu.py) and pixi manifest

[workspace]
name = "kmeans-xpu"
channels = ["conda-forge"]
platforms = ["linux-64"]

[dependencies]
python = "3.12.*"
numpy = ">=2.2,<3"
scikit-learn = ">=1.6"
array-api-compat = ">=1.11"
pytest = "*"
setuptools = "<82"

[pypi-dependencies]
dpctl = "*"
dpnp = "*"
triton-xpu = { version = "*", index = "https://download.pytorch.org/whl/xpu" }
torch = { version = "==2.14.*", index = "https://download.pytorch.org/whl/xpu" }

pixi run python bench_task_mix_xpu.py --sandbox ~/cursor-sandbox/kmeans_lloyd --out results/fit.jsonl

"""Intel XPU variant of cursor-sandbox PR 17's ``bench_task_mix.py``.

Same tasks, data, shared random-row starts, protocol (fresh process per run,
60 s limit, 4 threads, ``max_iter=12``, ``tol=1e-8``) and inertia as the CUDA
driver, which it imports. Only the GPU contenders change: torch on ``xpu``
(native ``index_add_`` update, and the portable one) and dpnp (portable update
only, ``native_accumulate`` returns None for it).

``measure_call`` uses CUDA events, so the timer here is device synchronize +
``perf_counter`` around the fit, with the data already on the device.

Usage: python bench_task_mix_xpu.py --sandbox ~/cursor-sandbox/kmeans_lloyd \
           --out results.jsonl [--tasks ...] [--impls ...] [--inits ...]
"""

from __future__ import annotations

import argparse
import json
import os
import subprocess
import sys
import time
from pathlib import Path

import numpy as np

IMPLS = ("torch-xpu", "torch-xpu-generic", "dpnp")


def _load_driver(sandbox):
    sys.path.insert(0, str(sandbox / "scripts"))
    sys.path.insert(0, str(sandbox / "src"))
    import bench_task_mix

    return bench_task_mix


def _to_numpy(array):
    module = type(array).__module__
    if module.startswith("torch"):
        return array.detach().cpu().numpy()
    if module.startswith("dpnp"):
        import dpnp

        return dpnp.asnumpy(array)
    return np.asarray(array)


def worker(args, driver):
    samples = np.load(args.data)
    centers0 = np.load(args.centers)
    native = not args.impl.endswith("-generic")
    if args.impl.startswith("torch"):
        import torch

        torch.set_num_threads(4)
        data = torch.as_tensor(samples, device="xpu")
        fixed = torch.as_tensor(centers0, device="xpu")
        torch.ones((2, 2), device="xpu") @ torch.ones((2, 2), device="xpu")
        sync = torch.xpu.synchronize
        sync()
        torch.xpu.reset_peak_memory_stats()
        before = torch.xpu.memory_allocated()
    else:
        import dpnp

        data = dpnp.asarray(samples)
        fixed = dpnp.asarray(centers0)
        dpnp.ones((2, 2), dtype=dpnp.float32) @ dpnp.ones((2, 2), dtype=dpnp.float32)
        sync = data.sycl_queue.wait
        sync()

    from kmeans_lloyd import KMeans
    from kmeans_lloyd.backends import use_native_backends

    use_native_backends(native)
    model = KMeans(
        n_clusters=args.k,
        init="k-means++" if args.init == "plusplus" else fixed,
        n_init=1,
        max_iter=driver.MAX_ITER,
        tol=driver.TOL,
        random_state=args.seed,
    )
    started = time.perf_counter()
    model.fit(data)
    sync()
    seconds = time.perf_counter() - started
    memory = None
    if args.impl.startswith("torch"):
        memory = int(torch.xpu.max_memory_allocated() - before)
    # ``seconds`` matches the CUDA driver (first fit in a fresh process, so it
    # includes kernel JIT); a second identical fit separates that cost out.
    started = time.perf_counter()
    model.fit(data)
    sync()
    warm_seconds = time.perf_counter() - started
    print(
        json.dumps(
            {
                "seconds": seconds,
                "warm_seconds": warm_seconds,
                "mem_increment_bytes": memory,
                "mem_meter": "torch-xpu-allocator" if memory is not None else None,
                "n_iter": int(model.n_iter_),
                "inertia_per_point": driver._mean_inertia(
                    samples, _to_numpy(model.cluster_centers_)
                ),
                "native_update": native and args.impl.startswith("torch"),
                "impl": args.impl,
                "init": args.init,
                "phase": "fit",
                "seed": args.seed,
                "k": args.k,
            }
        ),
        flush=True,
    )


def drive(args, driver):
    data_dir = Path.home() / "task-mix-data-v4"
    env = os.environ.copy()
    env.update(
        NVIDIA_TF32_OVERRIDE="0",
        OMP_NUM_THREADS="4",
        MKL_NUM_THREADS="4",
        OPENBLAS_NUM_THREADS="4",
    )
    tasks = {t for t in args.tasks.split(",") if t}
    impls = [i for i in (args.impls.split(",") if args.impls else IMPLS)]
    inits = [i for i in args.inits.split(",") if i]
    seeds = [int(s) for s in args.seeds.split(",") if s]
    out = Path(args.out)
    for task in driver.TASKS:
        if tasks and task["name"] not in tasks:
            continue
        data_path = data_dir / f"{task['name']}.npy"
        driver.generate(task, data_path)
        for seed in seeds:
            indices_path, centers_path = driver.save_random_start(data_path, task, seed)
            for impl in impls:
                for init in inits:
                    command = [
                        sys.executable,
                        str(Path(__file__).resolve()),
                        "--worker",
                        "--sandbox", str(args.sandbox),
                        "--impl", impl,
                        "--init", init,
                        "--data", str(data_path),
                        "--centers", str(centers_path),
                        "--k", str(task["k"]),
                        "--seed", str(seed),
                    ]
                    print(f"start {task['name']} seed={seed} {impl} {init}", flush=True)
                    row = {"task": task, "seed": seed, "impl": impl, "init": init, "phase": "fit"}
                    wall_started = time.perf_counter()
                    try:
                        completed = subprocess.run(
                            command, env=env, text=True, capture_output=True,
                            timeout=driver.TIMEOUT,
                        )
                    except subprocess.TimeoutExpired as error:
                        row.update(error="timeout", timeout_seconds=driver.TIMEOUT)
                        print("  timeout", flush=True)
                    else:
                        lines = completed.stdout.strip().splitlines()
                        if completed.returncode or not lines:
                            row.update(
                                error=f"exit {completed.returncode}",
                                stderr=completed.stderr[-2000:],
                            )
                            print(f"  exit {completed.returncode}", flush=True)
                        else:
                            row.update(json.loads(lines[-1]))
                            row["wall_seconds"] = time.perf_counter() - wall_started
                            print(
                                f"  {row['seconds']:.3f}s warm={row['warm_seconds']:.3f}s inertia={row['inertia_per_point']:.3f} "
                                f"iter={row['n_iter']}",
                                flush=True,
                            )
                    with out.open("a") as stream:
                        stream.write(json.dumps(row) + "\n")


def main():
    parser = argparse.ArgumentParser(allow_abbrev=False)
    parser.add_argument("--worker", action="store_true")
    parser.add_argument("--sandbox", type=Path, required=True)
    parser.add_argument("--impl")
    parser.add_argument("--init")
    parser.add_argument("--inits", default="plusplus,random")
    parser.add_argument("--data")
    parser.add_argument("--centers")
    parser.add_argument("--k", type=int)
    parser.add_argument("--seed", type=int)
    parser.add_argument("--tasks", default="")
    parser.add_argument("--impls", default="")
    parser.add_argument("--seeds", default="0,1,2")
    parser.add_argument("--out", default="results.jsonl")
    args = parser.parse_args()
    args.sandbox = args.sandbox.expanduser().resolve()
    driver = _load_driver(args.sandbox)
    if args.worker:
        worker(args, driver)
    else:
        drive(args, driver)


if __name__ == "__main__":
    main()

Run and written up with Claude Code.