Optimize hnsw::from_cagra<GPU>

[achirkin]

Reduce the CAGRA-for-HNSW build times by:

• avoiding unnecessary copies of the data between cagra::build and hnsw::from_cagra in the benchmarks
• avoiding unnecessary temporary data buffers in hnsw::from_cagra
• reducing random reads via forcing 1-1 mapping between the internal indices and external labels during HNSW import

As a side-effect, this PR also fixes the bug where hnsw::from_cagra segfaults in benchmarks if the dataset is passed in device memory (and incorrectly wrapped in a host_matrix_view).

In addition, this PR adds a bit more verbose NVTX reporting of different stages during the CAGRA/HNSW index build.

[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.

[achirkin]

/ok to test

[achirkin]

Example timeline

Hardware: NVIDIA H200 NVL + Intel(R) Xeon(R) Gold 6444Y (16 cores / 32 threads)
Dataset: DEEP-100M

      "build_param": {
        "M": 16,
        "hierarchy": "gpu",
        "graph_degree": 32,
        "intermediate_graph_degree": 48,
        "graph_build_algo": "IVF_PQ",
        "ivf_pq_build_pq_dim": 32,
        "ivf_pq_build_pq_bits": 8,
        "ivf_pq_build_nlist": 10000,
        "ivf_pq_build_niter": 10,
        "ivf_pq_build_ratio": 10,
        "ivf_pq_build_codebook_kind": "subspace",
        "ivf_pq_search_nprobe": 20,
        "ivf_pq_search_internalDistanceDtype": "half",
        "ivf_pq_search_smemLutDtype": "half",
        "ivf_pq_search_refine_ratio": 1
      }

branch-25.06 / input data on the host

branch-25.06 / input data on the device

(segmentation fault)

PR-826 / input data on the host

PR-826 / input data on the device

Notes

1. Open the images in a separate tab to see the numbers
2. There's a cudaMemCpy2dAsync in the branch-25.06 at the end of the CAGRA build. This is attaching the dataset to the CAGRA index; it's not used subsequently if the dataset is on the host
3. The most savings in hnsw::from_cagra come from the initialization phase, which is congested on atomic counter and slowed down due to copying data in a slightly randomized order (due to the concurrent updates of the atomic counter).
4. The data-on-device version is faster by ~5% than the data-on-host version

[cjnolet]

@achirkin it looks like we have some python failures from these changes (this makes me happy we now have python tests for the benchmarks).

[achirkin]

@cjnolet lol these tests haven't made me happy while I was trying to setup my environment to run them :) and they just complain about a missing useless column I removed for HNSW algorithm (GPU time).

[cjnolet]

/merge

[cjnolet]

@achirkin does this issue address #762?