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.

Dataset Arm Merge ms Δ Recall@12 Δ Rep spread

wiki-1m baseline 545.6 — 0.991683 — 1.1e-4

wiki-1m roundrobin 546.6 +0.18% 0.991875 +0.000192 4.1e-5

wiki-1m resplit 857.7 +57.2% 0.992083 +0.000400 2.5e-5

openai-2m baseline 2271.5 — 0.958383 — 7.5e-5

openai-2m roundrobin 2272.1 +0.03% 0.958758 +0.000375 8.4e-5

openai-2m resplit 3292.8 +45.0% 0.958158 −0.000225 1.2e-4

yfcc-10m baseline 2744.5 — 0.981367 — 7.5e-5

yfcc-10m roundrobin 2734.6 −0.36% 0.983583 +0.002216 3.4e-5

yfcc-10m resplit 3915.5 +42.7% 0.982783 +0.001416 4.2e-5

Good catch here, the interaction between concatenating the datasets and stable sorting here didn't occur to me. Experiments show that:

6% of leaves are single origin in the current design on wiki-1M, so this is happening, although even with round-robin you get like 5-10% of those still just because some single origin leaves happen for organic geometric reasons

round-robin slightly bumps recall with an impact on build time that seems to be within noise

further geometric splitting until you have leaves all at or below the max leaf size is expensive and doesn't seem amazing for recall.

Also trying not splitting them further and just treating them as big leaves, which I don't think will work but want to rule out before committing to round robin.

Dataset Arm Merge ms Δ Recall@12 vs baseline vs roundrobin Spread

wiki-1m baseline 514.5 — 0.991775 — +0.000008 2.5e-5

wiki-1m roundrobin 516.2 +0.35% 0.991767 −0.000008 — 1.0e-4

wiki-1m big512 542.1 +5.37% 0.992825 +0.001050 +0.001058 8e-6

wiki-1m big1024 554.0 +7.68% 0.993075 +0.001300 +0.001308 1.7e-5

openai-2m baseline 2203.4 — 0.958417 — −0.000308 6.7e-5

openai-2m roundrobin 2206.6 +0.15% 0.958725 +0.000308 — 1.6e-5

openai-2m big512 2273.0 +3.16% 0.960408 +0.001991 +0.001683 1.0e-4

openai-2m big1024 2287.2 +3.80% 0.960958 +0.002541 +0.002233 4.2e-5

yfcc-10m baseline 2419.1 — 0.981433 — −0.002125 5.9e-5

yfcc-10m roundrobin 2430.8 +0.49% 0.983558 +0.002125 — 1.2e-4

yfcc-10m big512 2528.8 +4.54% 0.984625 +0.003192 +0.001067 1.7e-5

yfcc-10m big1024 2594.4 +7.25% 0.985058 +0.003625 +0.001500 6.7e-5

Treating leaves of size up to 512 (default max leaf size is 256) as normal leaves that don't need round robin gives slightly higher recall but has some added complexity and takes slightly longer to do. For now, I'll stick to round robin for simplicity but let me know if you think this is a worthwhile tradeoff.

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.

I compared against always doing non-batched GEMM, which was substantially slower, and batching with individual sizes, which was moderately faster but iirc added more LoC than felt worth it.

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?

This seems to perform ~the same but adds a small amount of complexity to do the caching.

Dataset Arm Merge ms (median) Range Δ

wiki-1m baseline 517.0 506.0–518.8 —

wiki-1m cached 515.7 509.5–518.7 −0.25%

openai-2m baseline 2214.8 2199.9–2216.9 —

openai-2m cached 2220.0 2208.8–2221.4 +0.23%

yfcc-10m baseline 2422.0 2405.9–2429.0 —

yfcc-10m cached 2427.8 2417.4–2432.1 +0.24%

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.

landrumb · 2026-07-29T15:26:23Z

The Java API is still lagging, but is a bit more to convert. Could add to this PR or defer, up to @dantegd.

landrumb · 2026-07-29T20:30:33Z

/ok to test

Brings the branch up to date with main (f72199e), which resolves the check-c-abi failure: the checker was reporting cuvsResourcesSetWorkspacePool, cuvsRMMAsyncMemoryResourceEnable, cuvsCagraSearchMultiPartition and cuvsSelectK as removed functions, when in fact they were added upstream after this branch's base and were simply missing here. Only cpp/src/neighbors/cagra.cuh conflicted, and only positionally: both sides append a new function immediately after the existing merge() overload. Kept both -- this branch's merge() taking merge_params, and main's multi-partition search() overload. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

landrumb · 2026-07-30T21:34:31Z

Currently failing because the wheel builds are too big. Not sure there's much I can do about this.

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

landrumb added 2 commits July 29, 2026 14:03

switch to TF32 matmul and remove obsolete template

3b83294

updating C API

5e3f436

landrumb and others added 3 commits July 29, 2026 20:17

round-robin leaf splitting + more tests

bf21963

Device OOM fallback for AUTO

bf331a2

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

e788608

landrumb added 2 commits July 30, 2026 14:37

inferred tile_rows

488ab76

removed needless sync

a44ca42

landrumb marked this pull request as ready for review July 30, 2026 15:14

landrumb requested review from a team as code owners July 30, 2026 15:14

landrumb changed the title ~~[WIP] Fast CAGRA Index Merge~~ Fast CAGRA Index Merge Jul 30, 2026

landrumb and others added 6 commits July 30, 2026 15:39

bumping default max leaders to power of 2

b9b7da4

determinism test

ef0cf5e

missed files on default increase

298bfcf

removed obsolete template

8891aae

style miss

38c59f7

Dataset	Arm	Merge ms	Δ	Recall@12	Δ	Rep spread
wiki-1m	baseline	545.6	—	0.991683	—	1.1e-4
wiki-1m	roundrobin	546.6	+0.18%	0.991875	+0.000192	4.1e-5
wiki-1m	resplit	857.7	+57.2%	0.992083	+0.000400	2.5e-5
openai-2m	baseline	2271.5	—	0.958383	—	7.5e-5
openai-2m	roundrobin	2272.1	+0.03%	0.958758	+0.000375	8.4e-5
openai-2m	resplit	3292.8	+45.0%	0.958158	−0.000225	1.2e-4
yfcc-10m	baseline	2744.5	—	0.981367	—	7.5e-5
yfcc-10m	roundrobin	2734.6	−0.36%	0.983583	+0.002216	3.4e-5
yfcc-10m	resplit	3915.5	+42.7%	0.982783	+0.001416	4.2e-5

Dataset	Arm	Merge ms	Δ	Recall@12	vs baseline	vs roundrobin	Spread
wiki-1m	baseline	514.5	—	0.991775	—	+0.000008	2.5e-5
wiki-1m	roundrobin	516.2	+0.35%	0.991767	−0.000008	—	1.0e-4
wiki-1m	big512	542.1	+5.37%	0.992825	+0.001050	+0.001058	8e-6
wiki-1m	big1024	554.0	+7.68%	0.993075	+0.001300	+0.001308	1.7e-5
openai-2m	baseline	2203.4	—	0.958417	—	−0.000308	6.7e-5
openai-2m	roundrobin	2206.6	+0.15%	0.958725	+0.000308	—	1.6e-5
openai-2m	big512	2273.0	+3.16%	0.960408	+0.001991	+0.001683	1.0e-4
openai-2m	big1024	2287.2	+3.80%	0.960958	+0.002541	+0.002233	4.2e-5
yfcc-10m	baseline	2419.1	—	0.981433	—	−0.002125	5.9e-5
yfcc-10m	roundrobin	2430.8	+0.49%	0.983558	+0.002125	—	1.2e-4
yfcc-10m	big512	2528.8	+4.54%	0.984625	+0.003192	+0.001067	1.7e-5
yfcc-10m	big1024	2594.4	+7.25%	0.985058	+0.003625	+0.001500	6.7e-5

Dataset	Arm	Merge ms (median)	Range	Δ
wiki-1m	baseline	517.0	506.0–518.8	—
wiki-1m	cached	515.7	509.5–518.7	−0.25%
openai-2m	baseline	2214.8	2199.9–2216.9	—
openai-2m	cached	2220.0	2208.8–2221.4	+0.23%
yfcc-10m	baseline	2422.0	2405.9–2429.0	—
yfcc-10m	cached	2427.8	2417.4–2432.1	+0.24%

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!

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!

landrumb Jul 30, 2026 • 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!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

landrumb commented Jul 29, 2026

Uh oh!

landrumb commented Jul 29, 2026

Uh oh!

landrumb commented Jul 30, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

landrumb commented Jul 22, 2026 •

edited

Loading

landrumb Jul 30, 2026 •

edited

Loading