Introduce fused GEMM + 1-NN primitive using cuTile by divyegala · Pull Request #2249 · NVIDIA/cuvs

divyegala · 2026-06-17T21:00:49Z

This PR adds infrastructure built on top of existing JIT LTO architecture to generate kernels using cutile-python at build time, and embed them in the C++ library to make them callable from C++.

Merge #2548 first.

copy-pr-bot · 2026-06-17T21:00:52Z

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

huuanhhuyn

Many thanks @divyegala for the comprehensive PR!

I double-verified the functional correctness of the kernel from this PR. The kernel also matches one-to-one the performance with the one from gitlab.

dantegd

Could we make the path selection and workspace allocation agree on whether the cuTile launch is guaranteed to succeed? When minClusterDistanceCompute.cu selects FusedCutile, it doesn’t allocate the CUTLASS KVP scratch or mutex workspace. But if try_fused_1nn_tile() returns false, cpp/src/distance/detail/fused_distance_nn.cuh falls through to CUTLASS and expects those buffers to exist. That leaves us without a valid fallback if launcher construction fails, even though the try_ API suggests falling back is supported. Would it make sense to build or probe the launcher before committing to FusedCutile, or keep the fallback workspace allocated for as long as this path can still return false?

dantegd · 2026-08-27T16:58:59Z

+                                                                           cutlass_kvp_scratch,
+                                                                           stream);
      break;
+    case cuvs::distance::DistanceType::InnerProduct: break;


Could we route InnerProduct to a real non-cuTile implementation when cuTile isn't built or available?

On CUDA 12, try_fused_1nn_tile is compiled as false, so this function initializes the scratch output to (0, max), hits this empty case, and then unpacks that untouched sentinel. The new input_fp32_fused test includes InnerProduct unconditionally, so the supported CUDA 12 configuration returns incorrect results. What do you think about selecting the unfused path upstream for InnerProduct when cuTile is unavailable, or implementing the fallback here?

The unfused path does not support InnerProduct for some reason. Let me see if we can support or just throw for now.

dantegd · 2026-08-27T18:51:44Z

+  static_assert(std::is_same_v<IdxT, int> || std::is_same_v<IdxT, int64_t>);
+
+  constexpr int strict_pitch_elements = 16 / sizeof(DataT);
+  const bool use_strict_abi           = k % strict_pitch_elements == 0;


Both the strict and relaxed signatures declare base_addr_divisible_by=16 for matrices, norms, indices, and distances, but this selection only checks row pitch. Valid RAFT subviews or DLPack byte offsets can therefore launch under a false compiler assumption; the int64 loop below also makes xn + max_i32 and nearest_dist + max_i32 non-16-byte-aligned on its second batch. cuTile documents violating an exported signature assumption as undefined behavior: “Launching the exported kernel with an array that doesn’t satisfy this assumption would then result in undefined behavior.”

Should we validate base alignment for every non-null array before launching, or export a base-alignment-1 fallback? What do you think about making alignment part of eligibility before choosing either ABI?

Oh I did not think at all about subviews or offsets, good catch. For now, let's just validate base pointer alignment. I can add to the follow-up issue to validate 1-byte alignment and what kind of cubin is produced by cutile?

…e-python-to-cpp

dantegd · 2026-09-02T14:44:53Z

    case cuvs::distance::DistanceType::InnerProduct: {
-      // TODO: pass buffer
-      rmm::device_uvector<MathT> distances(n_rows * n_clusters, stream, mr);
+      if (uses_fused_distance_nn(


Could we make this decision using the same pointer-aware eligibility check as min_cluster_and_distance()? On CUDA 13, use_fused() can enter this branch before checking the actual X, centroid, and labels pointers. If any of those rejects cuTile, min_cluster_and_distance() falls back through use_legacy_fused(); for InnerProduct that means Unfused, which eventually calls pairwise_distance_kmeans() and throws because it only supports L2 metrics. CUDA 12 instead reaches the GEMM+argmin path below. What do you think about resolving final cuTile eligibility here and taking this existing GEMM fallback whenever it isn’t eligible?

dantegd · 2026-09-02T14:45:33Z

+
+      switch (fused_path) {
+        case FusedDistancePath::FusedCutile:
+          // Conservatively budget the fallback in case the eventual pointer-aware probe fails.


Could we budget the actual architecture-specific fallback here? On Blackwell—and on Hopper when both the row and cluster counts are below the legacy fused threshold—a pointer-aware cuTile rejection selects Unfused. That path allocates a batch-by-cluster distance matrix plus KVP output, but this case only charges int + KVP per row. Because predict_core() passes zero batch limits, the next allocation can exceed the 512 MiB target by a large factor. What do you think about budgeting the maximum of the cuTile scratch and the use_legacy_fused() scratch rather than assuming every fallback is CUTLASS?

dantegd · 2026-09-02T14:46:45Z

+{
+  if (!runtime_cuda13_or_newer()) { return false; }
+  if (has_embedded_cubin_for_arch(cc_major, cc_minor)) { return true; }
+  return tileir_fallback_available(driver_version);


I wonder if we shoukd gate the TileIR fallback on cuTile-supported compute capabilities? RAPIDS CUDA 13 still supports Turing/SM75, while cuTile requires compute capability 8.x through 12.x. On SM75 with a new enough driver, this returns true, so KMeans selects FusedCutile and attempts to load TileIR instead of keeping the supported legacy CUDA path. What do you think about checking the supported architecture range before reporting availability?

dantegd · 2026-09-02T14:47:46Z

+{
+  cudaLibrary_t library{};
+  RAFT_CUDA_TRY(
+    cudaLibraryLoadData(&library, image.data, nullptr, nullptr, 0, nullptr, nullptr, 0));


Should we preserve the nullable try_get_launcher() contract for expected module compatibility and JIT failures?

If I'm not mistaken, RAFT_CUDA_TRY here turns errors such as cudaErrorNoKernelImageForDevice, cudaErrorUnsupportedPtxVersion, and cudaErrorJitCompilerNotFound into exceptions, so can_launch_fused_1nn_tile() never returns false and KMeans cannot take its legacy fallback.

What do you think about treating expected compatibility/JIT errors as a failed build, unloading the library if kernel lookup fails, while still propagating genuine resource failures such as allocation errors?

dantegd · 2026-09-02T14:50:35Z

-                        raft::make_const_mdspan(minClusterAndDistance.view()),
-                        raft::make_device_vector_view<LabelT, IdxT>(labels, n_rows),
-                        raft::compose_op<raft::cast_op<LabelT>, raft::key_op>());
+      auto nearest_dist =


Could we keep native KVP storage when cuTile isn’t selected? On CUDA 12 this always allocates nearest_dist—and nearest_idx when LabelT differs, then minClusterAndDistanceCompute() allocates another full KVP because it was given SoA outputs.

The old path used the KVP output directly, so this adds avoidable per row scratch to every legacy L2/cosine balanced prediction; the hidden KVP also uses the current RMM resource rather than mr.

What do you think about deciding the path first and using minClusterAndDistanceComputeKvp() for legacy paths, as process_batch() does?

dantegd · 2026-09-02T14:51:46Z

+                                 weight_per_cluster,
+                                 batch_workspace,
+                                 /*reset_sums=*/false);
+    auto weighted_dist = raft::make_device_vector<DataT, IndexT>(handle, n_samples);


We can we avoid materializing weighted_dist on the cuTile path, no?

The KVP branch just below feeds a transform iterator over distance × weight into computeClusterCostFromIterator(), and the same approach works with nearest_dist here. As written, every cuTile Lloyd batch carries an extra n_samples * sizeof(DataT) allocation. What do you think about using the same iterator pattern in both branches?

dantegd · 2026-09-02T14:52:49Z

+        ndim=2,
+        index_dtype=index_dtype,
+        stride_lower_bound_incl=(0, None),
+        alias_groups=(),


Question: are callers allowed to use overlapping read-only dataset and centroid views? I couldn’t find a non-aliasing precondition on the KMeans APIs, while an empty alias_groups sequence tells cuTile that this array cannot alias any other parameter. If that overlap is supported, passing the same storage for X and centroids invokes the exported kernel outside its declared signature. Would common alias groups for the read only inputs or an explicit eligibility check and documented precondition make the contract clearer?v

They are not, but when is this a realistic path? This is akin to adding __restrict__ on each input pointer, right? When would centroids and X be the exact same pointer?

# Conflicts: # python/libcuvs/pyproject.toml

add example

41b0713

github-project-automation Bot added this to Unstructured Data Processing Jun 17, 2026

divyegala added 4 commits June 24, 2026 16:34

initial integration

b10c02c

attempt to fix tile linkage

1b934dd

working test, remove example

c7f7cbd

style check

86c9311

divyegala changed the title ~~cuTile Python to CPP embedding example~~ Introduce fused GEMM + 1-NN primitive using cuTile Jun 24, 2026

divyegala added 5 commits June 25, 2026 21:16

start integrating other metrics

e352629

if constexpr exit

d6560fc

passing KMeans tests

6742853

undo kvp, add constrains and alignments to tile export

db6b385

add reproducible benchmark scripts

0a3da06

ahendriksen reviewed Jun 30, 2026

View reviewed changes

Comment thread cpp/tests/CMakeLists.txt Outdated

benchmark sweep

9da31ae

josephine-wolf-oberholtzer added this to Unstructured Data Processing Jul 1, 2026

divyegala added 9 commits July 2, 2026 17:46

index dtype

35291ec

new reduce

f48ec5f

try int64 vector constant

842863e

new tile shape for sm90

9fb768d

using prebuilt cubins

88ddda8

use 13.4 tile

099324a

relaxed cubin, sm100 dispatch, new reduction strat

fe77c61

split tile for relaxed and strict sm120 abi

7e086b1

remove i64, batch i32 to save binary size

82d14a6

viclafargue mentioned this pull request Jul 22, 2026

[BUG] KMeans crashes for larger Datasets NVIDIA/cuml#8401

Closed

tfeher mentioned this pull request Jul 22, 2026

Use batched unfused 1-NN for KMeans on Blackwell #2346

Merged

divyegala added 3 commits July 28, 2026 16:00

working through

b0efee6

new tiles and reduction shapes for sm120

9b2d2cb

new kernel from huy

54125c0

huuanhhuyn reviewed Aug 26, 2026

View reviewed changes

divyegala added 7 commits August 26, 2026 21:45

sm86 for sm89

4d7d01c

Merge remote-tracking branch 'upstream/main' into cutile-python-to-cpp

0e7863c

sm80/86 shape divisible by

09d2838

int64 fix

200aad6

Merge remote-tracking branch 'upstream/main' into cutile-python-to-cpp

e4a775c

fix tests

146fcce

Merge remote-tracking branch 'upstream/main' into cutile-python-to-cpp

c148569

dantegd requested changes Aug 27, 2026

View reviewed changes

divyegala added 3 commits August 28, 2026 02:33

address review

87080b8

fix fp16 norms

ab012a6

Merge branch 'main' into cutile-python-to-cpp

ef5dced

shaunakkapur mentioned this pull request Aug 31, 2026

[CI] Lucene: GPU rewrite path bypasses Lucene's exact-search fallback for selective filters #2522

Open

divyegala added 5 commits September 1, 2026 01:28

clamp only the winner

7aa42b8

Merge remote-tracking branch 'origin/cutile-python-to-cpp' into cutil…

fca0b71

…e-python-to-cpp

regression and test fixes

45c66a8

self review

b79f4bf

preserve fallback fidelity better

ff9dae1

dantegd requested changes Sep 2, 2026

View reviewed changes

divyegala added 11 commits September 2, 2026 20:49

working through path unification

cafd051

Merge remote-tracking branch 'upstream/main' into cutile-python-to-cpp

74a0562

# Conflicts: # python/libcuvs/pyproject.toml

refactor backend choice to primitive

638a2ec

embeddings

8e2a9b7

style check

1fcea21

c build

de3d89a

merge

db2d6ff

style

e1cfd30

package

ec1c7d7

Merge branch 'cutile-cubin-embedding' into cutile-python-to-cpp

087abfe

style check

c0bcb1d

Uh oh!

Conversation

divyegala commented Jun 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

copy-pr-bot Bot commented Jun 17, 2026

Uh oh!

Uh oh!

huuanhhuyn left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

dantegd left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

divyegala commented Jun 17, 2026 •

edited

Loading

huuanhhuyn left a comment •

edited

Loading