[huuanhhuyn]

Try if we can select optimal row batch across GPUs based on number of SMs and number of CTAs per SM. Goal is to have full waves per batch launch.

[coderabbitai]

No actionable comments were generated in the recent review.

Summary

Summary by CodeRabbit

• Performance
  • Batch sizing for candidate datasets larger than the device’s L2 cache now adapts to the selected kernel’s occupancy and device capacity. Datasets that fit in L2 retain the maximum batch size. This behavior is applied separately to each kernel launch, while batched processing and 64-bit index handling remain unchanged.

The fused 1-NN tile launcher now selects a row batch for each ABI launch. It uses the maximum batch when candidate data fits in L2 cache and an occupancy-based batch otherwise. Each ABI launch processes all batches and widens int64 indices as needed.

Changes: Distance batch sizing

Layer / File(s)	Summary
Occupancy-based batch sizing cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu	The code calculates a two-wave batch from kernel occupancy, using active-cluster occupancy for clustered kernels and active-block occupancy across the device otherwise. For candidate data larger than L2 cache, the launcher uses this batch, capped at the maximum. For data that fits in L2 cache, it uses the maximum batch. Each ABI launch processes its batches and widens int64 indices.

Estimated code review effort: 2 (Simple) | ~10 minutes

Change: Feature

Suggested reviewers: divyegala

Merge Risk: Minimal · up to d9c87

The change sizes fused 1-NN row batches from kernel occupancy when candidate data exceeds L2. No correctness or safety defect was found, and a mistaken occupancy estimate would only affect performance. Add the missing category and breaking labels so the label check passes. Benchmarking on target GPUs would confirm the intended full-wave behaviour.

Security architecture risk: Minimal · up to d9c87

The change adjusts how much GPU work is processed per batch while preserving existing input checks and output limits. No introduced or worsened risk to security, isolation, or failure containment was established.

Pre-merge checks: 4

Check name	Status	Explanation
Title check		The title names the main change: selecting the 1-NN row batch size from GPU occupancy (SM count and CTAs per SM) to get full waves per launch. The raw summary confirms the batch size is now derived fr…
Description check		The description states the goal of selecting an optimal row batch from SM count and CTAs per SM so each batch launch has full waves. This matches the changeset, which derives batch sizes from occupanc…
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: Title check

Explanation

The title names the main change: selecting the 1-NN row batch size from GPU occupancy (SM count and CTAs per SM) to get full waves per launch. The raw summary confirms the batch size is now derived from kernel occupancy in fused_1nn_tile.cu. The '[Experimental]' prefix and the 'Tail Effect / Wave Quantization' label are extra context but do not make the title unrelated. The title is 68 characters, within a reasonable length.

Full details: Description check

Explanation

The description states the goal of selecting an optimal row batch from SM count and CTAs per SM so each batch launch has full waves. This matches the changeset, which derives batch sizes from occupancy in fused_1nn_tile.cu. Most of the description is the unedited contribution template, but the opening paragraph is related to the changes.

Finishing Touches

Generate unit tests (beta)

• Create a new PR

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

Recent review info

Run configuration

• Source mode: Auto
• Repository: NVIDIA/cuvs
• Branch: main
• File: .coderabbit.yaml
• Commit: 471e6d7
• Layers: Repository YAML
• Review profile: CHILL
• Plan: Enterprise
• Run ID: f5f03486-5bf5-48b7-b609-8af3556d3823

Commits

Reviewing files that changed from the base of the PR and between 8cb5d5a and d9c8790.

Files selected for processing: 1

• cpp/src/distance/detail/fused_distance_nn/cutile/fused_1nn_tile.cu

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