Skip to content

Use itopk_size=128 in the core CAGRA multi-partition test sweep #2839

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
bdice wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Use itopk_size=128 in the core CAGRA multi-partition test sweep #2839

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 8 additions & 8 deletions cpp/tests/neighbors/ann_cagra.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2485,15 +2485,15 @@ inline std::vector<AnnCagraMpInputs> generate_mp_inputs()
split,
graph_build_algo::NN_DESCENT,
algo,

@bdice bdice Oct 12, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Question for reviewers: does doubling itopk here lose coverage of a small-itopk path in multi-partition search? As far as I can tell, multi-partition search is still exercised at itopk=64 by:

  • the dimensionality block: AUTO (which resolves to SINGLE_CTA, per the comment at L2464-L2467), 4 partitions, dim 8 and 1024, L2 (L2518-L2532);
  • the max_queries-chunk block: both SINGLE_CTA and MULTI_CTA (2 CTAs per query at 64), 4 partitions, L2 (L2534-L2551);
  • the C API tests CagraC.BuildSearchMultiPartition* (ann_cagra_c.cu L1675), which use the default itopk_size = 64 (cagra.cpp L2404-L2409). These use tiny 2-row partitions, though, so they only check plumbing, not recall.

At itopk=64, the C++ suite no longer covers InnerProduct or Cosine, or the SKEWED split, in multi-partition search; those now run only at 128. If that matters, I can add back one itopk=64 case per metric at num_partitions=8, where recall at 64 is about 1.0.

/*itopk_size*/ 64,
// With a single 10000-row partition, itopk_size=64 leaves
// the recall of the fp16/int8/uint8 datasets at ~0.98, with
// a tail below 0.97 since the NN-descent graph differs from
// build to build. 128 puts every case near 1.0 (minimum
// ~0.994 over thousands of builds), so the single-index
// 0.985 convention applies.
/*itopk_size*/ 128,
Comment on lines +2488 to +2494

@bdice bdice Oct 12, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is the budget change. makeSearchParams() (L2197-L2204) passes itopk_size unchanged to the multi-partition cagra::search overload (cagra.cuh L589-L620), which forwards to detail::search_multi_partition (L351-L366). That function builds one search plan from the shared search_params, sized for the largest partition (L389-L398, L652-L665). It runs that plan for every (query, partition) pair and merges the per-partition results with a batched select_k (SINGLE_CTA L721-L733, MULTI_CTA L734-L746). So itopk_size is a per-partition budget:

  • With num_partitions=8 and an even split, each 1250-row partition gets its own itopk=64 search, and recall is about 1.0.
  • With num_partitions=1, the whole 10000-row index gets one itopk=64 search. That is effectively an ordinary single-index search: a plain cagra::search on the same index gave identical recall in 800/800 runs. On fp16, int8 and uint8 data, its recall sits around 0.98.

Inside the plan, adjust_search_params() (search_plan.cuh L200-L246) derives the default max_iterations from itopk_size / search_width. For MULTI_CTA, num_cta_per_query = ceildiv(itopk_size, 32) (L257-L260). Going from 64 to 128 therefore roughly doubles the SINGLE_CTA iteration cap and doubles MULTI_CTA's CTAs per query from 2 to 4, in addition to widening the candidate buffer.

Question for reviewers: is raising the budget to 128 better than keeping 64 and lowering the floor? I chose 128 because, with the floor also going up to 0.985, the test becomes stricter rather than looser. Keeping 64 would need a floor below the 0.966 we measured at that budget, and a floor that low would also let a few-percent regression in the merge or decode logic through.

metric,

@bdice bdice Oct 12, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why this only showed up on GB300, and for L2: the "fixed" dataset is not the same on every GPU. SetUp() seeds RngState(1234) (L2434-L2442), and InitDataset draws fp16/fp32 data through raft::random::uniformInt (L193-L236). For fp16 at dim=64 that gives 8 distinct values per dimension (resolution 4), so the data is coarse and has many exact distance ties. raft's RNG launches 4 * getMultiProcessorCount() blocks (rng_impl.cuh L65-L75), and each thread seeds its generator from its global thread id and strides by the grid size (rng_device.cuh L675-L694). So seed 1234 produces a different dataset on each SM count. The GB300 dataset happens to be harder for L2 (mean 0.9786 vs 0.9845 on A100). Over a dataset-seed sweep, both GPUs averaged 0.984. At itopk=64, InnerProduct and Cosine sat at 0.993 to 0.996, so L2 was the metric near the floor.

Question for reviewers: every recall test that seeds InitDataset this way depends on SM count, so a fixed seed only pins the data per GPU model. Is that worth a separate issue (for example, generating test data with a fixed launch grid or on the host), or is the per-GPU variation acceptable as extra coverage?

// Lower than the single-index 0.985 convention. These
// low-redundancy single-partition / skewed configs dip
// further under filtering, where removing rows thins the
// survivors' graph connectivity; observed recall floors
// around 0.974, so 0.97 leaves headroom for run-to-run
// jitter while still catching real regressions.
/*min_recall*/ 0.97});
/*min_recall*/ 0.985});

@bdice bdice Oct 12, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This brings the floor back to the single-index convention. The single-index CAGRA inputs use 0.985 (e.g. L1911-L1932 and the 10000-row L1984-L2006). This fixture uses the same graph degrees on purpose so that the thresholds are comparable (L2179-L2182).

The deleted comment says 0.97 was chosen from observed floors around 0.974. With eps = 0.003 the check is actual_recall < min_recall - eps (ann_utils.cuh L214), so with 1000 neighbors the test fails at 0.966 or below. At itopk=64 the GB300 Search/0 minimum was already 0.967 (run 38177641942), so there was almost no headroom. At itopk=128 the minimum was 0.993 for num_partitions=1 and 0.999 for num_partitions=8 (run 38184518531), so 0.985 leaves room for build-to-build variation and still catches regressions that 0.97 would have missed.

The max_queries-chunk block below still uses 0.97 (L2534-L2551). The unchanged blocks had a minimum of 0.990 in the verification run, so the chunk block could also move to 0.985, but I left the non-core blocks out of this PR.

}
}
}
Expand Down
Loading