Skip to content

Fix the IVF-PQ list codepacking loops to stride over the whole grid - #2750

Merged
rapids-bot[bot] merged 2 commits into
NVIDIA:mainfrom
bdice:ivf-pq-codepacking-grid-stride
Oct 8, 2026
Merged

rapids-bot[bot] merged 2 commits into
NVIDIA:mainfrom
bdice:ivf-pq-codepacking-grid-stride

Conversation

@bdice

@bdice bdice commented Oct 6, 2026

Copy link
Copy Markdown
Contributor

write_list, write_list_flat and run_on_list in ivf_pq_codepacking.cuh started at the global row index but advanced by one block's worth of rows:

uint32_t stride = subwarp_align::div(blockDim.x);                             // rows per block
uint32_t ix     = subwarp_align::div(threadIdx.x + blockDim.x * blockIdx.x);  // global start
for (; ix < len; ix += stride) { ... }

Their launchers already size the grid to cover every row (encode_list_data launches one block per 8 rows, or 16 for 4-bit codes; pack/unpack/reconstruct launch one block per 256 rows), so block b processed every row from its first one to the end of the list. Rows were encoded or packed up to n_rows / rows_per_block times, and the subwarps of block 0 walked the entire list, each row a serial loop over all pq_dim subspaces. For the ~128-row lists in NEIGHBORS_ANN_IVF_PQ_TEST that is up to 16 passes per row; at pq_dim = 3072 one encode_list_data_interleaved_kernel launch took ~127–146 ms.

Affected paths:

  • encode_list_data, used by ivf_pq::helpers::codepacker::extend_list, for any list longer than one block (8–16 rows).
  • pack_list_data, unpack_list_data, {un}pack_contiguous_list_data and reconstruct_list_data for lists longer than 256 rows. The C API (cuvsIvfPqIndexUnpackContiguousListData) and Python index.lists() go through these; unpacking a 100k-row list did ~2·10⁷ row-unpacks instead of 10⁵.

The main build / extend encoder (process_and_fill_codes_kernel) does not use these loops and is unaffected.

This PR makes the three loops grid-stride loops by multiplying the stride by gridDim.x. Launch configurations are unchanged.

  • Codes are identical. Each row is still handled by the same per-row routine (encode_vectors, or the pack/unpack/reconstruct action) with the same lane layout, codebook scan order and reduction order, and is written to the same location. The only difference is that each row is now processed exactly once instead of up to 16 times.
  • This also removes a possible lost update: with pq_bits < 8, the FLAT-layout encoder and unpack_contiguous do bitfield read-modify-writes on shared bytes, and two blocks processing the same row could collide.

The commit also updates the file's SPDX copyright header to the canonical form required by the copyright pre-commit hook.

Testing

All IVF-PQ test executables and the full ctest suite pass. The existing tests already check that packing round-trips byte-exactly and that reconstruct → extend reproduces the vectors.

Measurements

Single-process wall time on an RTX 6000 Ada (48 GB), otherwise idle (no ctest parallelism). The old and new libcuvs.so were run alternately, 2 or more repetitions each, with the order reversed between repetitions; mean (min–max). The two builds differ only by this change.

executable before after change
NEIGHBORS_ANN_IVF_PQ_TEST 158.6 s (157.9–159.4) 138.5 s (138.2–138.7) -12.7%

The other executables measured the same way (IVF-Flat, IVF-SQ, CAGRA, multi-GPU, tiered index, dynamic batching, all-neighbors) were unchanged within noise. All tests passed in every run.

write_list, write_list_flat and run_on_list started at the global row index but advanced by one
block's worth of rows, so block b processed every row from its first one to the end of the list.
Each row was encoded (or packed/unpacked) up to 16 times, sequentially in the first blocks; with
pq_dim 3072 that made an encode launch take ~130 ms. Advance by the whole grid instead, so every
row is processed exactly once. The codes are unchanged.
@bdice bdice added bug Something isn't working non-breaking Introduces a non-breaking change labels Oct 6, 2026
@copy-pr-bot

copy-pr-bot Bot commented Oct 6, 2026

Copy link
Copy Markdown

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.

@bdice
bdice marked this pull request as ready for review October 7, 2026 16:05
@bdice
bdice requested a review from a team as a code owner October 7, 2026 16:05
@coderabbitai

coderabbitai Bot commented Oct 7, 2026 •

Copy link
Copy Markdown

Review in Change Stack →

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: ed91b39b-f0d4-4e3b-ae2a-7881047c7509
📥 Commits

Reviewing files that changed from the base of the PR and between 7f21ade and 154ba69.

📒 Files selected for processing (1)
  • cpp/src/neighbors/ivf_pq/ivf_pq_codepacking.cuh

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


📝 Summary

Summary by CodeRabbit

  • Bug Fixes
    • Corrected how vectors are assigned across GPU threads when processing lists, helping ensure vectors are handled consistently across the full grid.

Walkthrough

The IVF-PQ codepacking loops now advance vector indices using grid-wide strides. The copyright notice now includes affiliates.

Changes

IVF-PQ codepacking

Layer / File(s) Summary
Grid-wide vector iteration
cpp/src/neighbors/ivf_pq/ivf_pq_codepacking.cuh
run_on_list advances each thread’s index by the grid stride. write_list and write_list_flat advance each subwarp’s index by the grid-wide subwarp stride. The copyright notice includes affiliates.

Priority: ➖ Normal

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

Change: Bug fix

Suggested reviewers: tarang-jain

Merge Risk: ⚪ Minimal · up to 154ba

The grid-stride changes address repeated row processing, and the inspected packed-byte paths do not show a remaining merge-blocking issue.

🚥 Pre-merge checks | ✅ 5
✅ Passed checks (5 passed)
Check name Status Explanation
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.
Title check ✅ Passed The title clearly and concisely describes the main change: correcting IVF-PQ codepacking loops to use the full grid stride.
Description check ✅ Passed The description is directly related to the changeset and explains the affected loops, root cause, behavior change, testing, and performance results.
✨ 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.

@divyegala

Copy link
Copy Markdown
Contributor

/merge

@rapids-bot
rapids-bot Bot merged commit e288281 into NVIDIA:main Oct 8, 2026
213 of 217 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

bug Something isn't working non-breaking Introduces a non-breaking change

Projects

Status: Done

Development

Successfully merging this pull request may close these issues.

2 participants