Improve CAGRA multi-CTA heuristic: max_iterations

[sherylll]

For multi-CTA mode, current max_iterations is assigned a large number. However, at low itopk, extra iterations only add latency and does not do much to recall. So the intuition is to use a lower max_iterations for smaller itopk.

minimum_depth = 8 * (search_quality - 1), since we don't provide a knob for search_quality for the moment, hardcoding this to 16. search_quality = 5 would mean the previous default 32.

When the problem is harder, or when more results are needed, it is still a good idea to turn up the max_iterations.

Quality = 3 is the new default, and quality = 5 is the old default.

This is a subset of #2502, where the change to width or hashmap size seem to have different effect on different GPU generations, which still remains to be investigated. But the conclusion on max_iterations should hold.

[copy-pr-bot]

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.

[achirkin]

/ok to test fe21bf4

[copy-pr-bot]

/ok to test fe21bf4

@achirkin, there was an error processing your request: E2

See the following link for more information: https://docs.gha-runners.nvidia.com/cpr/e/2/

[achirkin]

/ok to test fe21bf4

[achirkin]

Thanks, the speedup looks promising! A small request below

[achirkin]

Could you please add small comments explaining the logic for setting the max_iterations like this on these two lines?

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

• Configuration used: Repository: NVIDIA/cuvs/.coderabbit.yaml
• Review profile: CHILL
• Plan: Advanced
• Run ID: 6209de7e-a713-4bf7-8252-945f96545592

📥 Commits

Reviewing files that changed from the base of the PR and between efc7d30 and 5cb5bfb.

📒 Files selected for processing (2)

• cpp/src/neighbors/cagra.cuh
• cpp/src/neighbors/detail/cagra/search_plan.cuh

Included review availability: This review used your included allowance. Your plan provides up to 4 included reviews per hour; 3 remain after this review.

---

📝 Summary

Summary by CodeRabbit

• Improvements
  • Searches with bitset filters now estimate the filtering rate automatically when no rate is specified, including across multi-partition indexes.
  • Multi-CTA searches now adjust their initial search depth based on search settings and filtering, which can improve search behavior across different workloads.

Walkthrough

CAGRA now computes filtering rates from bitsets when no explicit rate is set. Automatic MULTI_CTA iteration limits also use search width, top-k size, and filtering rate.

Changes

CAGRA search tuning

Layer / File(s)	Summary
Bitset filtering rate and search integration cpp/src/neighbors/cagra.cuh	A helper computes and clamps the fraction of rows removed by bitsets. Single-index and multi-partition searches use this rate when filtering_rate is negative.
Automatic MULTI_CTA iteration limit cpp/src/neighbors/detail/cagra/search_plan.cuh	The initial automatic iteration limit now depends on search width, rounded itopk_size, topk, and filtering rate. Existing reachable-node and minimum-iteration adjustments remain.

Priority: ⬇️ Low

Estimated code review effort: 3 (Moderate) | ~20 minutes

Change: Bug fix

Suggested reviewers: achirkin

Merge Risk: ⚪ Minimal · up to 5cb5b

No search-impacting issue was established; this change is mergeable after normal checks.

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly describes the main change: improving the CAGRA multi-CTA max_iterations heuristic.
Description check	✅ Passed	The description explains the rationale for reducing max_iterations at smaller itopk values and provides benchmark context relevant to the changeset.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check. Docstring coverage is scoped to functions touched by this diff. Analyzed 0 functions across 0…
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create a new PR

• Autopilot · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

---

_{Comment @coderabbitai help to get the list of available commands.}