[WIP] Fast CAGRA Index Merge by landrumb · Pull Request #2352 · NVIDIA/cuvs

landrumb · 2026-07-22T21:16:58Z

This PR implements the Fastener graph merge operation. This PR supports merging for float, half, int8, and uint8 dtypes, and euclidean distances.

There was originally specialization to use int8 GEMM for the integer types, but I ran into portability issues on Ada and it turns out that using the same f32 path for both is simpler and not substantially slower. Currently investigating switching the unified path to TF32 to use tensor cores.

The core logic resides in cagra_merge_scaffold.cuh.

This PR adds 7.64 MiB to libcuvs.so, a 2.94% increase.

(H100)

copy-pr-bot · 2026-07-22T21:17:01Z

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.

copy-pr-bot · 2026-07-22T21:29:09Z

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.

dantegd · 2026-07-22T21:32:37Z

/ok to test 2c54392

landrumb · 2026-07-23T19:24:32Z

/ok to test

landrumb · 2026-07-24T00:31:03Z

/ok to test

landrumb · 2026-07-24T19:09:05Z

/ok to test

landrumb · 2026-07-27T16:00:43Z

/ok to test

landrumb · 2026-07-27T16:05:26Z

/ok to test

landrumb · 2026-07-27T19:35:39Z

/ok to test

dantegd

Overall, I really like the direction and PR! Reusing the partition graphs, adding a lightweight cross-partition scaffold, and then running the existing CAGRA optimization is such a great way to make merge substantially cheaper than a rebuild. The split between preflight, consolidation, scaffold construction, and graph assembly also makes the implementation fairly easy to follow.

One API concern: the C API still calls the existing C++ overload without merge_params, so it needs updating, right?

dantegd · 2026-07-28T21:14:24Z

+    RAFT_EXPECTS(range.start == covered && range.end > range.start &&
+                   range.end <= static_cast<int64_t>(partitions.memberships.size()),
+                 "Fastener partition ranges must compactly cover all memberships");
+    for (int64_t start = range.start; start < range.end; start += leaf_size) {


I think there’s a connectivity edge case here. Since the sorts are stable, rows within a partition can remain grouped by their source index. Slicing an oversized partition into consecutive chunks can then produce leaves containing rows from only one input, and the leaf kernel skips every same-origin pair.

For example, identical inputs whose row counts are multiples of leaf_size can end up with every chunk aligned to a single origin, leaving the original graphs disconnected. Could we either keep splitting or deterministically mix origins before slicing? A regression test for that case would also be helpful.

dantegd · 2026-07-28T21:15:28Z

+    raft::common::nvtx::range<cuvs::common::nvtx::domain::cuvs> scope("cagra::merge/append/sort");
+    cagra::detail::graph::sort_knn_graph_device_inplace(
+      handle, params.metric, raft::make_const_mdspan(dataset.view()), merged_graph.view());
+    merged_graph = merge_scaffold::cap_sorted_graph(


Are we capping this a little too early? If an input already has graph_degree local neighbors that are closer than its scaffold neighbors, this nearest prefix cap can remove every cross index edge before graph::optimize() gets a chance to consider it.

It also leaves the optimizer with the same input and output degree, so there’s no candidate headroom for robust pruning. Was wondering if maybe we keep up to max(intermediate_graph_degree, graph_degree) candidates here, bounded by the available candidate width, and let the optimizer produce the final graph_degree?

dantegd · 2026-07-28T21:17:42Z

+    RAFT_FAIL("FASTENER cagra::merge is unsupported: %s", preflight.reason.c_str());
+  }
+
+  return merge_fastener(handle, params, merge_params, indices, preflight);


Should AUTO fall back to rebuild if Fastener runs out of memory? Preflight can validate shapes and limits, but it can’t know whether the consolidated dataset, temporary GEMM buffers, and graphs will all fit at runtime.

The previous merge path could fall back to host memory, so changing the existing overload to AUTO means some merges that previously succeeded may now fail. I think explicit FASTENER should still fail, but AUTO retrying the rebuild path would preserve the existing behavior. What do you think?

dantegd · 2026-07-28T21:18:06Z

+  if (params.graph_degree == 0 || static_cast<uint64_t>(params.graph_degree) >= rows) {
+    return reject("graph_degree must be positive and smaller than the combined row count");
+  }
+  if (static_cast<uint64_t>(params.graph_degree) > max_input_degree + scaffold_degree) {


I think one limit is missing from preflight: the graph passed to sort_knn_graph_device_inplace() can’t exceed kMaxSortDegree (1024). For example, an input degree of 1024 plus any scaffold currently passes preflight and then fails later in the sorter. In AUTO mode, that also means we miss the rebuild fallback.

Could we include max_input_degree + scaffold_degree <= kMaxSortDegree in the eligibility check?

dantegd · 2026-07-28T21:18:34Z

+}  // namespace
+
+typedef AnnCagraIndexMergeTest<float, float, std::uint32_t> AnnCagraFastenerMergeRecallTest;
+TEST_P(AnnCagraFastenerMergeRecallTest, DefaultFloatMergeSearchRecallAgainstBruteForce)


Could we add recall coverage for at least one non-float dtype and one merge with more than two inputs? The current dtype tests show that the output graph is structurally valid, but only the small two-part float case verifies that the merged graph is actually useful for search. That feels important if AUTO is going to become the default.

dantegd · 2026-07-28T21:24:53Z

+  auto dim    = dataset.extent(1);
+
+  // Cut every parent in the bucket into fixed-height tiles
+  int tile_rows = context.assignment_tile_rows;


Could we derive tile_rows from the dimension and padded leader count instead of always using 2048?

Preflight only checks the leaf-GEMM shape, while assignment needs roughly (tile_rows + padded_leaders) * dim + tile_rows * padded_leaders float elements. Some configurations therefore pass preflight and later fail the workspace assertion here if I'm not mistaken. Could we reduce the tile height per bucket and reject only when the leader matrix plus even one point row cannot fit?

dantegd · 2026-07-28T21:26:33Z

+                  dataset,
+                  parents,
+                  buckets[bucket_index],
+                  1 << bucket_index,


Curious, have you measured how much work is going to padded leaders here? A parent just over a power-of-two boundary could nearly doubles the gathered leader storage, dot-product GEMM, distance materialization, and select_k width.

dantegd · 2026-07-28T21:26:55Z

+    RAFT_CUDA_TRY(cudaGetLastError());
+
+    int leader_blocks = strided_grid_size(static_cast<int64_t>(batch_size * leader_elements));
+    manyway_gather_tile_leaders_kernel<<<leader_blocks, THREADS_PER_BLOCK, 0, stream>>>(


It looks like every tile of the same parent gathers the same leader matrix again. For large parents split across many 2048-row tiles, that can add substantial conversion and memory traffic, especially when the leader count is close to the tile height. Could we gather each parent’s leaders once and reuse them across its tiles, or did benchmarking show that the added bookkeeping costs more?

dantegd · 2026-07-28T21:27:39Z

+  int64_t total  = static_cast<int64_t>(batch_size) * tile_rows * padded_leaders;
+  for (; linear < total; linear += stride) {
+    int leader     = static_cast<int>(linear % padded_leaders);
+    int row        = static_cast<int>((linear / padded_leaders) % tile_rows);


This is a small thing, but wanted to ask: for dimensions that aren’t naturally 16-byte aligned, this allocates and copies a second full consolidated dataset while all input datasets are still alive. Would it be practical to consolidate directly into aligned storage and pass a row stride through the scaffold/sort kernels? That could materially lower peak memory for int8/uint8 and other unaligned dimensions.

dantegd · 2026-07-28T21:28:10Z

+      static_cast<uint32_t>(offsets[part]));
+    RAFT_CUDA_TRY(cudaGetLastError());
+  }
+  raft::resource::sync_stream(res);


Is this synchronization needed? The helper has no local temporary buffers consumed by the copy kernels, and the caller immediately sorts the returned graph on the same resource stream. Stream ordering should already preserve that dependency

landrumb added 7 commits July 16, 2026 18:15

factoring out in-place edge sort

637e0fc

core merge logic

7972d3a

remove VMM dataset consolidation from fastener merge

bc6ec7b

pedantic docstring change

e2ce96f

assert/validation changes

a9149db

test improvements

cf78e32

assert/message changes & aligned signatures

209485c

github-project-automation Bot added this to Unstructured Data Processing Jul 22, 2026

landrumb and others added 2 commits July 22, 2026 21:26

pre-commit hooks

da96f26

Merge branch 'main' into landrumb/cagra-fastener-merge

2c54392

aamijar assigned landrumb Jul 22, 2026

aamijar added non-breaking Introduces a non-breaking change feature request New feature or request labels Jul 22, 2026

aamijar added this to Unstructured Data Processing Jul 22, 2026

aamijar moved this to In Progress in Unstructured Data Processing Jul 22, 2026

landrumb added 2 commits July 22, 2026 22:16

clearer check rationale

9a3cabf

test arg mismatch

af25e86

landrumb and others added 3 commits July 23, 2026 21:05

refactor to use existing select_k kernel

2eb64fa

fixed attached dataset alignment issue

6f65f87

Merge branch 'main' into landrumb/cagra-fastener-merge

2e6e148

CI fixes

c8fd540

aligned test with earlier ci fix

951b0e2

Merge branch 'main' into landrumb/cagra-fastener-merge

678f2ea

adjusting test with excessive assumptions

731bb30

dantegd requested changes Jul 28, 2026

View reviewed changes

Uh oh!

Conversation

landrumb commented Jul 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

copy-pr-bot Bot commented Jul 22, 2026

Uh oh!

copy-pr-bot Bot commented Jul 22, 2026

Uh oh!

dantegd commented Jul 22, 2026

Uh oh!

landrumb commented Jul 23, 2026

Uh oh!

landrumb commented Jul 24, 2026

Uh oh!

landrumb commented Jul 24, 2026

Uh oh!

landrumb commented Jul 27, 2026

Uh oh!

landrumb commented Jul 27, 2026

Uh oh!

landrumb commented Jul 27, 2026

Uh oh!

dantegd left a comment

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!

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!

3 participants

landrumb commented Jul 22, 2026 •

edited

Loading