004 — tile sizes, a Cython Lloyd step, and k-means++

The compiled running tile from 003 removes the distance matrix and is slower than the Gram GEMM. This iteration asks three things:

• How sample-tile and centroid-tile sizes move the speed and memory of a Gram assignment.
• Whether that tiled Gram step, after profiling, gets closer to a chunked GEMM Lloyd iteration.
• What k-means++ initialization costs in the same array-API code.

The core functions stay in Python array-API syntax and run without torch.compile. The compiler, when used, is wrapped around those functions by the benchmark driver.

Schedules added

||x||² is omitted inside a tile because it does not change the argmin. That is the same reduction the Cython kernel uses. Centroid norms are computed once per centroid tile, then reused for every sample tile.

Cython reference

The comparison target is lloyd_iter_chunked_dense from cakedev0/scikit-learn#26. Each chunk has 256 samples. The chunk computes ||c||² - 2 X Cᵀ with a GEMM, takes argmin, and scatters the samples into per-thread center buffers. Those buffers are reduced after the parallel region.

The branch is built on the L4 host as scikit-learn 1.10.dev0 (g++ 13.3). The host has 4 CPUs, so the iteration uses 4 threads.

Gram tiles on feature-heavy

65536 samples, 4096 centroids, 16384 features, float32, TF32 off. Median of two calls after one warmup. Compile time is not included in the steady median. The first compiled tile took 9.0 s to compile; the others took 0.4–1.4 s.

For every tile whose distance matrix is smaller than the centroid tile, the measured peak is tile_k * 16384 * 4 bytes plus 528 KiB. That is one float32 copy of the centroid tile, the buffer a transposed GEMM needs, plus a fixed half megabyte. It does not grow with the sample tile. The 256-sample chunk used by the Cython kernel is the same 256.5 MiB as a 4096-sample tile when every centroid is in the product, and it is slightly slower (1.08 s against 0.92 s).

The distance matrix exceeds that copy only for the full 65536 by 4096 product, which is 1.00 GiB. Eager execution peaks at 2.00 GiB there. torch.compile leaves the 1.00 GiB product and drops the extra gigabyte. It does not change the steady time of the tiled products. The fastest compiled tile in this grid is 4096 by 1024: 0.928 s and 64.5 MiB, against 1.042 s and 4.00 GiB for the eager full Gram.

Compiled running tiles

Same feature-heavy problem. One torch.compile block covers tile_k centroids, then the driver walks every centroid. Steady time excludes compile. tiled merges the running minimum outside the compiled block. stateful merges inside it. tile_k 256 was not run: at 128 the compile was already 28 s and the distance tile was back.

At 16 and 32 the stateful peak is 1.50 MiB, two copies of the n * 12 byte label-and-distance state. At 64 and 128 the peak jumps to the (n, tile_k) distance tile (16 MiB and 32 MiB) plus about 2 MiB, and the assignment gets slower. The fastest running schedule is the outside merge at 32 centroids: 5.98 s and 2.56 MiB. The fastest Gram tile, 0.928 s and 64.5 MiB, is still about six times faster.

Putting the merge inside the compiled block saves the extra copy at 16 and 32 and does not make the assignment faster.

Kernels

One compiled Gram tile of 4096 samples by 1024 centroids is ampere_sgemm_128x128_tn for 10.3 ms, then triton_per_fused_amin_argmin_mul_sub_unsqueeze for 0.01 ms. The argmin reads the GEMM output in a second kernel.

One compiled running tile of 32 centroids, merge included, is a single kernel, triton_red_fused_add_full_lift_fresh_lt_mul_select_sub_sum_where_zeros, for 49.5 ms. Distance, comparison, and the running-minimum update are in that kernel. Covering all 4096 centroids takes 128 of those launches, which is the 6.6 s stateful row above.

Lloyd step

Feature-heavy, eager, float32. The assignment and the update are timed separately. Each phase resets the device peak counter.

The tiled mean update matches the one-hot update in time and lowers the update peak from 1.75 GiB to 432 MiB. The 4096 x 1024 assignment is the schedule to keep: the whole iteration is 1.70 s, and neither phase allocates the 4 GiB Gram temporary.

Same arrays as the Cython iteration

65536 samples, 1024 centroids, 2048 features, float32, on this host. The GPU call is an eager lloyd_step with a 256 by 1024 Gram tile and an update tile of 256. The CPU call is lloyd_iter_chunked_dense with 4 threads.

The labels differ in 1 of 65536 rows. The GPU step is short because this shape was sized so the 4-thread CPU iteration would last more than a second. On the feature-heavy problem the same array-API step is the 1.70 s row above.

k-means++

65536 samples, 1024 centers, 8192 features, float32. The trial count is 2 + log(n_clusters).

The first version recomputed X * X for every new center and pulled candidate indices back to the host. That took 37.5 s and peaked at 4.00 GiB above the inputs.

The version in the tree computes the sample norm once and gathers candidate rows with take. It takes 10.3 s and peaks at 2.00 GiB, which is one X * X the size of the sample matrix.

Progress

• Tile grid, running-tile sweep, kernel names, Lloyd step, Cython comparison, and both k-means++ timings are recorded. results.json and cython.json in this folder are the sources.
• The Gram tile of 4096 by 1024 plus an update tile of 256 is the Lloyd step that keeps the feature-heavy iteration at 1.70 s without the 4 GiB assignment temporary.

Cython timing

Same host as the GPU runs. One dense iteration on 65536 samples, 1024 centroids, 2048 features, float32. Chunk length 256. Median of 3 calls after one warmup. The assignment call sets update_centers=False. The Lloyd call also updates centers and relocates empty clusters.

The thread-local distance tiles and center buffers are 36.0 MiB together (4 threads times a 256 by 1024 distance tile, plus the center sums). Process ru_maxrss after the run was 1.63 GiB, which includes the 512 MiB sample matrix. The matching PyTorch lloyd_step on this shape is not timed yet.

Progress

• Gram tiles and compiled running tiles are both in the tables above. results.json in this folder is the source.
• Cython Lloyd on 65536×1024×2048 is in. The same shape has not been timed with lloyd_step yet.
• The kernel profile, the tiled Lloyd update, and k-means++ are running on the L4.