Enable fused 1-NN for small dim on consumer Blackwell (sm_120)

[maxwbuckley]

Enable the fused 1-NN kernel for small dim on consumer Blackwell (sm_120)

Problem

use_fused() in cpp/src/cluster/detail/kmeans_common.cuh selects the unfused
split GEMM/reduction path for prop.major >= 10:

} else if (prop.major >= 10) {
  // On Blackwell onwards, use unfused
  return false;
}

That predicate covers two different machines:

	prop.major	shared mem / SM	tensor cores
datacenter Blackwell (B200)	10	228 KB	tcgen05
consumer/workstation Blackwell (GB202)	12	100 KB	no tcgen05

The rule was introduced in #1768 along with the unfused path, benchmarked on
IVF-PQ. On sm_120 it is too broad: the fused kernel still wins for small dim,
where the unfused path's m x n distance matrix costs more memory traffic than
the GEMM costs arithmetic.

This matters for a real workload. IVF-PQ trains its PQ codebooks by running
kmeans_balanced on pq_len-dimensional subspaces with pq_book_size clusters
(ivf_pq_build.cuh:407), and pq_len = dim / pq_dim is typically 4-8 — squarely
in the regime where the gate currently picks the slower path.

Change

Split sm_120/121 out of the major >= 10 branch and enable fused for small
dim there. sm_100/103 behaviour is unchanged — see Caveats.

Measurements

RTX 5090 (sm_120), CUDA 12.9. Fused vs. the current batched unfused path (post
#2346), interleaved A/B, 40 reps, medians, dedicated GPU. Speedup > 1 means fused
is faster. All entries significant at |Welch t| > 3 unless noted.

dim	n=256	n=1024	n=4096
4	1.09-1.32x	1.13-1.17x	1.27-1.28x
8	1.13-1.21x	1.14-1.18x	1.27-1.32x
12	1.01-1.05x	1.14-1.16x	1.28-1.34x
16	0.94-1.29x	1.14-1.18x	1.28-1.32x
20	0.82-0.97x	0.82-0.83x	0.94-0.95x
24	0.77-0.94x	0.82-0.90x	0.94-0.96x

m ∈ {200K, 1M, 4M}; ranges span m. The crossover is sharp between dim 16 and 20.
At n=256 the gain is gone by dim 16, hence the n >= 512 || k <= 8 condition.

Above dim 20 the existing rule is correct and is left alone — fused degrades to
0.22-0.30x by dim 512, which is what #1768 measured.

End-to-end effect is small

Being upfront: the primitive-level win does not translate into a large
IVF-PQ build speedup, because PQ codebook training is only a modest share of
build time. ivf_pq::build, N=1M, dim=128, n_lists=1024, interleaved A/B,
30 reps:

pq_dim	pq_len	baseline	this PR	Welch t
32	4	625 ms	0.997x	-1.37 (not significant)
64	2	1142 ms	1.053x	4.01

The effect appears only once pq_dim is large enough that codebook training is a
meaningful fraction of the build. So this is a small, no-regression improvement
rather than a headline one — it is worth taking because the dispatch is free, not
because it transforms build times.

Correctness

Labels differ from the unfused path at very small dim (up to 25% of rows at
dim=2, n=4096). These are ties, not errors: the achieved objective (total
min-distance) agrees to 5.3e-4 relative with sign-varying difference, so neither
path finds systematically worse minima. The per-row gaps are the cancellation
error of expanded-L2 under TF32 when ||x||^2 is large relative to the distance,
which affects both paths equally.

Caveats

• Only sm_120 was measured. I do not have sm_100 hardware. Extending this
  conclusion to datacenter Blackwell without data would repeat the overreach this
  PR is fixing, so that branch is untouched. If someone can run the same sweep on
  B200, the same treatment may well apply there.
• Unrelated but noticed while investigating: minClusterDistanceCompute (the
  distance-only variant used by k-means++ init) calls fusedDistanceNNMinReduce
  with no use_fused() gate at all. So on Blackwell, k-means++ init already runs
  the fused kernel while the Lloyd loop refuses to. That asymmetry looks
  unintentional — happy to file it separately.

Reproducing

Benchmark harness used for the table is bench_fused_vs_unfused.cu (not included
in this PR); it calls fusedDistanceNNMinReduce and unfusedDistanceNNMinReduce
directly at the shapes minClusterAndDistanceCompute uses.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.