From c674194dbd5e6cfcc08feb309b0cb1dfae7da95e Mon Sep 17 00:00:00 2001 From: Max Buckley Date: Thu, 30 Jul 2026 17:36:09 +0200 Subject: [PATCH] Enable fused 1-NN for small dim on consumer Blackwell (sm_120) use_fused() selected the unfused split GEMM/reduction path for all prop.major >= 10. That predicate covers both datacenter Blackwell (sm_100/103, major 10) and consumer Blackwell (sm_120/121, major 12), which differ in shared memory per SM and tensor core generation. On sm_120 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. Measured on RTX 5090 against the batched unfused path, fused is 1.13-1.34x faster for k <= 16 at n >= 1024 and 1.09-1.32x for k <= 8 at n = 256, and 0.77-0.97x slower from k = 20 upward. Datacenter Blackwell behaviour is unchanged; that branch was not measured. Co-Authored-By: Claude Opus 5 (1M context) --- cpp/src/cluster/detail/kmeans_common.cuh | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/cpp/src/cluster/detail/kmeans_common.cuh b/cpp/src/cluster/detail/kmeans_common.cuh index ab3ef0a05a..ccca38731e 100644 --- a/cpp/src/cluster/detail/kmeans_common.cuh +++ b/cpp/src/cluster/detail/kmeans_common.cuh @@ -62,7 +62,8 @@ namespace cuvs::cluster::kmeans::detail { * * On Ampere (SM <= 8.x) always use fused. * On Hopper (SM 9.x) use fused when m or n >= 4096. - * On Blackwell (SM >= 10.x) use unfused. + * On datacenter Blackwell (SM 10.x) use unfused. + * On consumer Blackwell (SM 12.x) use fused only for small k. */ template bool use_fused(const raft::resources& handle, IdxT m, IdxT n, IdxT k) @@ -75,8 +76,16 @@ bool use_fused(const raft::resources& handle, IdxT m, IdxT n, IdxT k) } else if (prop.major == 9 && (m >= 4096 || n >= 4096)) { // On Hopper if m, n are bigger than 4096, use fused return true; + } else if (prop.major == 12) { + // Consumer Blackwell (sm_120/sm_121) is a different machine from sm_100: + // 100 KB of shared memory per SM and no tcgen05. Here the fused kernel still + // wins for small k, where the unfused path's m x n distance matrix costs more + // traffic than the GEMM costs arithmetic. Measured crossover is k ~= 16-20, + // except at small n where the unfused reduction is already cheap and fused + // only pays off for the smallest k. + return k <= 16 && (n >= 512 || k <= 8); } else if (prop.major >= 10) { - // On Blackwell onwards, use unfused + // On datacenter Blackwell, use unfused return false; } return false;