Retrieving top-k results for k >= 257 produces corrupted outputs

Describe the bug

Retrieving top-k results with ivf_flat for k >= 257 produces corrupted outputs. I also encountered the same issue when using ivf_pq.

Steps/Code to reproduce the bug

I followed the tutorial at this link.

When setting k = 257, the outputs are corrupted. I then plotted the distance values:

import matplotlib.pyplot as plt
plt.plot(distances[0, :])
plt.show()

Expected behavior

For reference, this is the result when k = 256:

Environment details

Installed packages:

pip install cuvs-cu12==24.8.0 cupy-cuda12x==12.2.0 --extra-index-url=https://pypi.nvidia.com

Additional context

Interestingly, performing a search with k > 256 followed by a refinement step with k <= 256 works correctly. However, if I attempt to refine with k > 256, the issue reoccurs.

Begin with notebooks/ivf_flat_example.ipynb and reproduce the issue using cuvs-cu12==24.8.0 and the reported CuPy version, comparing k=256 with k=257 for ivf_flat and ivf_pq. Then inspect the search and refinement entry points for handling k values above 256. Done means both direct searches and refinement produce correct outputs for k > 256.