Refactor NND BBQ kernels

[sherylll]

• Extract wmma INT8 implementation into a helper file
• Simplify SIMT implementation: remove the unrolled inner product variants to keep code lean, and add proper warp guard; net effect is positive, except for the 1x1 cases

pair	fill	before (ms)	after (ms)	Δ
1×1	half	82.2	106.6	+30%
1×1	light	80.0	96.2	+20%
2t×2t	half	187.0	187.3	±0
2t×2t	light	182.7	159.1	−13%
1×4t	half	223.3	186.8	−16%
1×4t	light	215.6	160.4	−26%
2t×4t	half	340.3	342.9	+1%
2t×4t	light	332.3	289.3	−13%

On wmma implementation: both u8 and fp8 wmma are slower on B300 than on earlier generations, but the matmul is memory bound, so the overall slowdown is small; will move to tcgen05 if this becomes relevant

[copy-pr-bot]

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.

[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: Enterprise
• Run ID: 584cf2dd-a3ee-41e1-890b-dceadd59c0d1

📥 Commits

Reviewing files that changed from the base of the PR and between 56d6a0d and 843ddc9.

📒 Files selected for processing: 4

• cpp/src/neighbors/detail/cagra/graph_shared.cu
• cpp/src/neighbors/detail/nn_descent.cuh
• cpp/src/neighbors/detail/promoted_i4_warp_mma.cuh
• cpp/src/preprocessing/quantize/detail/bbq_distance.cuh

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

📝 Summary

Summary by CodeRabbit

• Performance
  • Updated BBQ-based nearest-neighbor distance calculations for packed 4-bit data using GPU tensor-core operations.
• Bug Fixes
  • Improved handling of uneven search lists by checking row and column bounds before calculating or storing distances.
  • Updated distance calculations to support differing encoded layouts where supported, while requiring matching layouts for packed 7/8-bit data.

The change updates BBQ code inner-product dispatch and CAGRA row-distance calls. It also changes NN-descent SIMT join indexing and replaces the int4 tensor-core join implementation with a warp-level MMA path.

Changes: BBQ Distance and NN-Descent Joins

Layer / File(s)	Summary
BBQ code inner-product and metric contract cpp/src/preprocessing/quantize/detail/bbq_distance.cuh, cpp/src/neighbors/detail/cagra/graph_shared.cu	The inner-product API now accepts code pointers, layouts, and document row length. Query layout selects packed or bit-plane computation. bbq_row_distance passes the encoded rows to the updated API.
SIMT join cell processing cpp/src/neighbors/detail/nn_descent.cuh	Both SIMT phases process one row-column cell per accumulator. They check row and column indices against the applicable list sizes before computing or converting distances.
Int4 warp-MMA join path cpp/src/neighbors/detail/promoted_i4_warp_mma.cuh, cpp/src/neighbors/detail/nn_descent.cuh	The new header provides packed-int4 warp MMA with unsigned-byte MMA tiles. The join stores raw accumulators before converting valid cells to distances.

Priority: ➖ Normal

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

Change: Refactor

Suggested reviewers: lowener, cjnolet

Merge Risk: ⚪ Minimal · up to 843dd

The changed layout arithmetic, join-cell coverage, and int4 distance conversion preserve their contracts. The mixed-dimension quantizer hazard predates this change; no new actionable merge risk remains beyond normal checks.

Pre-merge checks: 4 · 1

Failed checks: 1 warning

Check name	Status	Explanation	Resolution
Docstring Coverage		Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 1 functions across 1 files. (3 skipped: 3 …	Write docstrings for the functions missing them to satisfy the coverage threshold.

Passed checks: 4 passed

Check name	Status	Explanation
Title check		The title clearly and concisely identifies the main change: refactoring the NND BBQ kernels.
Description check		The description is directly related to the changes and includes the implementation scope, performance results, and known regressions.
Linked Issues check		Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check		Check skipped because no linked issues were found for this pull request.

Full details: Docstring Coverage

Explanation

Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 1 functions across 1 files. (3 skipped: 3 unsupported.)

• Fix all pre-merge checks with AI

✨ Finishing Touches: 💡 1

🛠️ Fix failing CI checks: 💡

• Commit to this branch
• Create a new PR

🧪 Generate unit tests (beta)

• Create a new PR

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

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