perf(index): parallelise the k-means centroid update over data and admit k=16 into the AMX GEMM

[xtangxtang]

Related: #6369. Touches the same function as the open #8560 -- the two cover disjoint shapes, see the section on that below.

The problem

train_ivf is hierarchical k-means followed by one closing assignment of every training row against the full centroid set (compute_loss). The closing assignment has had the AMX-FP16 GEMM under it since #8540. The hierarchical part has not, and its centroid update ran on one core.

The update in KMeansAlgoFloat::to_kmeans parallelised over the centroids: each thread owned a range of centroid rows and scanned the whole dataset for members in its range. That layout has a hard ceiling -- you cannot have more threads than centroids -- and the code made it explicit:

if k < num_cpus || k < 16 { num_cpus = 1; }

Hierarchical k-means trains every node of its tree at k <= hierarchical_k, which defaults to 16. So on any build host with more than 16 cores, every centroid update in the training loop ran single-threaded -- and for an f16 column, as scalar f16 adds, each a round trip through f32 -- while the distance computation next to it used every core. Sibling subtrees run in parallel, but at the top of the tree there are fewer nodes than cores, and that one core is the critical path.

A second, smaller thing sits next to it. The GEMM declined any centroid set below its 32-wide block. Everything in the training loop is k <= 16 -- the per-node training and the assignment of a node's rows to its children -- so the GEMM never ran there, only for the closing assignment and during shuffle.

The fix

Two commits, independent, in order of impact.

1. Partition the data, not the centroids. Each block sums its own rows into a private k × dimension f32 accumulator; the partials are folded in block order at the end. Every row is visited once -- the old layout walked the whole membership array once per block, p × n visits, to touch each vector from exactly one of them -- and the block count is bounded by the row count and thread count, not by k. The old strategy stays as a fallback for the shapes where accumulating in place is cheaper: f64; k × dimension too large to hold under a 512 MB budget; a non-f16 column with only one block's worth of rows; and, where the in-place update is itself parallel (k at least the thread count, and at least 16), any shape whose accumulators would cost more to fold than the work they save. That last rule is blocks × k × dimension <= 4 × n for f32, plus 6 × n × dimension / threads for f16 (whose in-place path also pays a scalar convert-add-convert per element); the constants come from the crossover measurements below and sit on their winning side. The test the in-place path makes to choose between one owner and many is shared as centroid_owner_blocks, so the routing cannot drift from the update it routes to.

The accumulators are f32 whatever the element type. They are fresh memory either way, so widening them costs one extra k × dimension per block and nothing in the pass over the data; what it buys is that an f16 column no longer depends on the size of its clusters -- an f16 running total past 2048 cannot represent an increment of 1. train_kmeans caps its input at k × 512 rows precisely so the in-place f16 path never hits that (#4977); the new path keeps the property without leaning on the cap.

2. Lower the GEMM's k floor from 32 to 16. PackedCentroidsF16 already pads to a multiple of 32 with zero rows, so a 16-centroid set half-fills every pass. It is still the better trade: the alternative re-reads the whole centroid table per row and is memory-bound. The zero columns score a distance of exactly 1.0 and only the first k columns are consulted, which the padding test now pins at k = 16. prefers_flat_amx_assignment in utils.rs moves with it, as its doc comment requires; at k = 16 that predicate is not reached in practice (may_train_index returns first below 1_000_000 centroid elements), so the change there keeps the invariant true rather than changing behaviour.

Relationship to #6369 and #8560

#6369 is the p × n membership rescan in to_kmeans. #8560 (open) addresses it with a stable membership-to-row index for the case k >= available_parallelism, and keeps a single-threaded accumulation for everything below that:

if available_parallelism == 1 || k < available_parallelism || k < 16 { /* one thread */ }

That excluded case is every node of hierarchical k-means on any host with more than 16 cores, and it is the case this PR measures. The two mostly do not compete on shapes: #8560's indexed path applies where k is at least the thread count; commit 1 here applies unconditionally where k is below it -- where an owner-partitioned layout has no parallelism to offer at all -- and above it only where the accumulators are small next to the membership (a PQ sub-vector, or f16 at high dimension), handing everything else back to the in-place update unchanged. They do edit the same function, so whichever lands second rebases; if #8560 goes first I will fold this onto its recompute_float_centroids as the branch its serial arm takes. I have not marked this as closing #6369 while #8560 is open against it.

Commit 2 is independent of both.

Performance

LAION 100M, 768-dim fp16, dot distance, IVF_PQ, 10000 partitions, sample_rate = 256 (2.56M training rows). 32 physical cores / 64 threads of a Xeon 6972P (Granite Rapids), pinned with numactl + taskset, LANCE_CPU_THREADS=64. Three binaries from the same toolchain (clang-16, -mprefer-vector-width=512, fp16kernels): main at 8c392b2, main + commit 1, main + both commits. Two runs per cell, in alternating order.

The figure is the Trained IVF model in ... log line: hierarchical k-means plus the closing assignment, without the sampling I/O in front of them (that stage is 23--43 s here and drifts with the page cache, which is why I do not quote whole-stage took= numbers).

Trained IVF model in, seconds	main	+ commit 1	+ commit 2 (this PR)	main → PR
AMX enabled	22.0 / 22.9	10.0 / 10.2	8.4 / 8.6	2.6×
LANCE_DISABLE_AMX=1	80.7 / 76.2	67.0 / 68.2	64.8 / 65.6	1.2×

Reading across the columns:

• Commit 1 is worth ~12 s on either row (22.4 → 10.1 with AMX, 78.5 → 67.6 without). Same saving whether or not the GEMM is running, which is what a CPU-independent change should show. It is the hierarchical stage's single-core update, and it is most of what main spends in that stage.
• Commit 2 is worth ~1.6 s with AMX (10.1 → 8.5, both runs of each arm within 0.3 s), and by construction nothing without it -- the gate is never consulted when AMX is off, and the bottom row's 67.6 vs 65.2 is the run-to-run scatter of the AMX-off arms (2--4 s). It is the smaller commit and I would rather say so than let the ratio suggest otherwise.
• The remaining gap between the rows (~57 s) is the closing assignment, GEMM vs per-vector. That was already in main and is not credited to this PR.

So on an AMX host main spends ~19 s of its 22 s in the hierarchical stage and this PR brings that stage to ~5 s; on a host without AMX the closing assignment dominates and the PR is a 1.2× on the total.

Three notes on how solid these numbers are.

• Every cell is two runs. The AMX-on arms repeat to within 0.3 s; the AMX-off arms scatter by 2--4 s. Ratios above are on the means.
• The AMX-off arms are the same binaries with LANCE_DISABLE_AMX=1, confirmed by the startup banner (amx_usable=false) and by the timing itself. The per-vector dot kernel in those arms was confirmed by disassembly to be 512-bit AVX-512 (zmm registers present), not a scalar fallback, so the baseline is what it claims to be.
• These numbers are on current main. The branch was first developed against 706b941, before the hierarchical k-means rewrite; there the single-core update was ~266 s of a ~320 s train_ivf and commit 1 saved ~245 s. The rewrite (training each node on at most k × 512 rows, subtrees in parallel) already took most of that away, which is why the same fix is now worth ~12 s. The structure of the result did not change -- commit 1 large, commit 2 small -- only the base it sits on.

Correctness

Both commits keep the per-vector and in-place paths as fallbacks and pin the new paths against them:

• the two centroid-summing strategies agree to accumulation-order tolerance on random f32 and f16 data across block counts;
• an f16 cluster too large to total in f16 comes out at the right mean on one block and on many -- the single-block case matters because it is what a CI container gets;
• routing declines exactly the f64, over-budget and single-block shapes and nothing else, on every element type;
• against a parallel in-place update, routing keeps a small accumulator and hands back a large one at both element types' allowances, and ignores the allowance wherever that update is serial (k below the thread count, or below 16);
• an empty input totals to zeros rather than reaching par_chunks with a zero chunk size;
• an out-of-range membership id is ignored rather than written past the accumulator;
• the AMX gate still declines k = 15 and dim = 31, on any host;
• the existing GEMM-vs-per-vector agreement and padding-never-wins tests now include k = 16.

Within a host the new update is reproducible bit for bit (block-order fold). Across hosts with different core counts, the last bits of a centroid differ; train_kmeans already seeds from SmallRng::from_os_rng, so cross-host reproducibility was not on offer before either.

Checks

cargo fmt --all -- --check
cargo clippy --all --tests --benches -- -D warnings
cargo test -p lance-index --release --lib kmeans

The AMX-dependent tests ran the kernel on the development host (Xeon 6972P, Granite Rapids, clang-16; CPUID reports AMX-FP16 and the XTILEDATA grant succeeds) and skip on hosts without it.

Where the in-place update is already parallel

The open question when this went up as a draft was whether the data path regresses the shapes the in-place update already served well, k at or above the thread count. It did, on some of them, and the routing rule above is the result. Method: a temporary #[ignore] test (not in the PR) calling sum_clusters_over_data and sum_clusters_over_centroids directly on random data, in a 64-thread rayon pool pinned with taskset to 32 cores / 64 threads of the Xeon 6972P, median of 3--200 repetitions per cell after a warm-up. "data" is the data path forced on regardless of routing; "→" is what the routing now picks.

The five shapes of benches/kmeans_recompute.rs (from #8560), microseconds per update:

shape	f32 in-place	f32 data	→	f16 in-place	f16 data	→
N=512, dim=1024, k=256	283	(one block)	in-place	178	91 (2.0×)	data, single block
N=16384, dim=1024, k=64	463	1159 (0.40×)	in-place	3124	956 (3.3×)	data
N=65536, dim=1024, k=4096	8471	64474 (0.13×)	in-place	12355	53878 (0.23×)	in-place
N=65536, dim=64, k=256	695	902 (0.77×)	in-place	1519	889 (1.7×)	in-place¹
N=131072, dim=128, k=256	1569	1798 (0.87×)	in-place	4179	1806 (2.3×)	data

¹ A conservative miss, not a regression: at 16 folded floats per row and dimension 64 the f16 allowance (4 + 6 × 64 / 64 = 10) declines a shape the data path would have won by 1.2--1.7× across runs. Run-to-run scatter on the in-place f16 cells is 1.5--2×, hence the ranges.

The shapes this PR is for, and a PQ sub-vector, same units:

shape	f32 in-place	f32 data	→	f16 in-place	f16 data	→
hierarchical node N=8192, dim=768, k=16	1072	290 (3.7×)	data	32167	230 (140×)	data
hierarchical node N=131072, dim=768, k=16	19264	1869 (10×)	data	515110	1227 (420×)	data
PQ sub-vector N=65536, dim=8, k=256	564	205 (2.8×)	data	659	169 (3.9×)	data
flat N=2.56M, dim=768, k=10000	62861	140121 (0.45×)	in-place	202893	118079 (1.7×)	data

Where the constants come from -- N=32768, 32 blocks, k=256 (or 128 where marked), data path relative to in-place:

folded floats per row	f32	f16	f16 dimension
4	1.63×	2.0×	16
8	1.08×	1.68×	32
16	0.67×	1.24×	64
32	0.31×	0.86×	128
64 (k=128)	0.26×	1.01×	512
128 (k=128)	0.34×	1.01×	1024
256 (k=128)	0.17×	1.01×	2048
256	0.31×	0.58×	1024

f32 crosses over between 8 and 16 folded floats per row, independent of dimension; the allowance is 4. f16 crosses over at roughly 32, 64, 128 and 256 per row for dimension 128, 512, 1024 and 2048 -- proportional to dimension / threads, because what it additionally saves is one scalar half-precision add per element per owner; the allowance is 4 + 6 × dimension / threads, against a measured slope of about 8. With the rule in place, none of the nineteen shapes above is slower than main; the train_ivf numbers are unaffected, since every shape in that build is k <= 16.

[lance-gatekeeper]

✅ Gate recommendation: approve.

The effective PR patch is unchanged after the latest merge from main. The data-partitioned f32 accumulators remove the hierarchical k-means bottleneck while bounding memory and retaining the prior centroid-owner path for unsuitable shapes; the k=16 AMX threshold remains aligned with exact-assignment routing and excludes padded centroids from reduction.