[BUG] cagra::search on vpq_f16_index triggers OOB for every search_algo (26.10 nightly)

[cpegeric]

Describe the bug

cuvs::neighbors::cagra::search on a cagra::vpq_f16_index<half, uint32_t> that was constructed by hand from (1) a graph produced by cagra::build(..., host_padded_dataset_view, attach_dataset_on_build=false) and (2) a device_vpq_dataset produced by preprocessing::quantize::pq::make_vpq_dataset(..., host_matrix_view) — the workflow the pq.hpp docstring documents as the CAGRA-Q pattern — triggers an out-of-bounds device read in the search kernel. Reproduces on all three search_algo variants (MULTI_CTA, SINGLE_CTA, MULTI_KERNEL).

Steps/Code to reproduce bug

Attached: repro_cuvs_vpq_search_oob.cu — single file, ~180 lines, only depends on libcuvs / libraft / librmm. 10K rows × 128 dim, pq_dim=32, pq_bits=8, graph_degree=32.

# Build (any recent nvcc; only public cuvs headers used)
nvcc -std=c++20 -O2 --extended-lambda --expt-relaxed-constexpr \
     -I$CONDA_PREFIX/include -I$CONDA_PREFIX/include/rapids \
     -I$CONDA_PREFIX/include/raft -I$CONDA_PREFIX/include/cuvs \
     -L$CONDA_PREFIX/lib -lcuvs -lcuvs_c -lrmm -lstdc++ \
     -DLIBCUDACXX_ENABLE_EXPERIMENTAL_MEMORY_RESOURCE \
     -DRAFT_SYSTEM_LITTLE_ENDIAN=1 \
     -gencode arch=compute_89,code=sm_89 \
     repro_cuvs_vpq_search_oob.cu -o repro

# Run (immediate crash on MULTI_CTA search)
LD_LIBRARY_PATH=$CONDA_PREFIX/lib ./repro

Bare run output:

graph built: size=10000 dim=0 graph_degree=32
vpq trained
vpq_idx: size=10000 dim=128 graph_degree=32
--- searching with algo=MULTI_CTA ---
terminate called after throwing an instance of 'raft::cuda_error'
  what():  CUDA error encountered at:
    file=.../raft/core/interruptible.hpp line=294:
    call='query_result',
    Reason=cudaErrorIllegalAddress:an illegal memory access was encountered

Under compute-sanitizer --tool memcheck the failing read looks like:

========= Invalid __global__ read of size 4 bytes
=========     at search_multi_cta+0xd1f0
=========     by thread (0,0,0) in block (1,29,0)
=========     Access to 0x795d9e2d2d00 is out of bounds
=========     and is 240466944 bytes after the nearest allocation
=========     at 0x795d8fd7f100 of size 1 bytes
=========     Host Frame: cuvs::neighbors::cagra::detail::multi_cta_search::select_and_run<
=========                   __half, unsigned int, float, unsigned int,
=========                   cuvs::neighbors::filtering::none_sample_filter>(...)
=========     Host Frame: cuvs::neighbors::cagra::search<
=========                   __half, unsigned int,
=========                   dataset_view<vpq_dataset_container, __half, long, ...>,
=========                   unsigned int>(...)

The nearest allocation ... of size 1 bytes looks like a placeholder descriptor being passed as a real buffer.

Setting search_params::algo = SINGLE_CTA or MULTI_KERNEL produces the same class of cudaErrorIllegalAddress — different search kernel, same shape.

Expected behavior

cagra::search on a vpq_f16_index returns the top-K neighbours without a memory violation, matching the CAGRA-Q pattern documented at cuvs/preprocessing/quantize/pq.hpp:280:

// `idx` is a `cagra::index<float, uint32_t>` with graph built on dense rows.
// `padded` is a `device_padded_dataset_view<float, int64_t>` view of those same rows.
cuvs::neighbors::vpq_params vpq_params{};
auto vpq = cuvs::preprocessing::quantize::pq::make_vpq_dataset(res, vpq_params, padded.view());
idx.update_device_dataset_same_layout(res, vpq.as_dataset_view());

The reproducer takes the equivalent "construct a new vpq_f16_index from the graph + vpq view" shape because it's the natural path when attach_dataset_on_build = false was used at build time.

Environment details:

• Environment location: Bare-metal, Linux 6.17.0-1013-aws (Ubuntu 24.04), NVIDIA L40S (compute 8.9), driver 590.48.01, 491 GB host RAM.
• Method of cuVS install: conda (rapidsai-nightly).
• Relevant packages:

  libcuvs         26.10.00a206  cuda13_260819055025_2140532c  rapidsai-nightly
  libcuvs-headers 26.10.00a206  cuda13_260819055025_2140532c  rapidsai-nightly
  libraft         26.10.00a26   cuda13_260819053610_9aa17e57  rapidsai-nightly
  cuda-version    13.3          hcbadf70_3                    conda-forge
  

• Source tree HEAD matches the shipped .so: 2140532c ("Use librtcx for JIT-LTO", 2026-08-18).

Additional context

• Reproduces on every search_algo value (MULTI_CTA, SINGLE_CTA, MULTI_KERNEL), so it does not appear to be one specialised kernel — likely a mistake in the dataset-descriptor plumbing that all three search implementations consume.
• Traced from the pq.hpp:280 CAGRA-Q docstring example — the intended workflow.
• Superficially similar to issue #1802 but different kernel (multi_cta_search vs single_cta_search) and different dataset descriptor (vpq_dataset_container vs plain float). Fixing one may or may not cover this; worth checking whether the underlying dataset_descriptor_base_t construction path is shared.
• Build path itself is fine — cagra::build returns a host_padded_index with sane size() and graph(), and make_vpq_dataset returns a device_vpq_dataset with correct n_rows/dim/stride. The hand-stitched vpq_f16_index also reports correct size()/dim()/graph_degree(). Only the search kernel touches invalid memory.

[cpegeric]

// ----------------------------------------------------------------------------
// Minimal reproducer for cuvs 26.10-nightly OOB when searching a
// vpq_f16_index that was hand-stitched from a separately-built cagra graph
// and a separately-built VPQ dataset.
//
// This is a public-API-only pipeline: no MatrixOne code, no rmm helpers
// beyond what raft pulls in.
//
// The pipeline mirrors the CAGRA-Q workflow that the pq.hpp docstring
// documents (build a dense graph, then attach VPQ codes for search):
//
//   1. build cagra graph on dense f16 rows via host_padded_dataset_view,
//      with `attach_dataset_on_build = false` (host_padded_index has no
//      device dataset — cuvs streams host batches internally).
//
//   2. train VPQ codes over the same dense host rows via
//      cuvs::preprocessing::quantize::pq::make_vpq_dataset(res, params, view).
//      Passing a raft::host_matrix_view intentionally: cuvs's detail path
//      sees a host pointer, subsamples for kmeans training, and streams the
//      encode in bounded batches (per pq.cu train_from_rows docstring).
//
//   3. hand-construct a `vpq_f16_index<half, uint32_t>` from
//      (metric, vpq.as_dataset_view(), graph_idx.graph()). The cagra::index
//      ctor documented at cagra.hpp:600 copies the graph into its own
//      device buffer via update_graph.
//
//   4. cuvs::cagra::search on the stitched vpq_f16_index → OOB in
//      search_multi_cta on the DEFAULT (AUTO) search algo. Setting
//      SINGLE_CTA moves the OOB into search_single_cta::search_kernel;
//      MULTI_KERNEL fails the same way.
//
// Sanitizer output on our repro (10K rows, dim 64, dim large enough for
// kmeans + PQ; k = 10 queries):
//
//   ========= Invalid __global__ read of size 4 bytes
//   =========     at search_multi_cta+0xd1f0
//   =========     by thread (0,0,0) in block (1,29,0)
//   =========     Access to 0x795d9e2d2d00 is out of bounds
//   =========     and is 240466944 bytes after the nearest allocation
//   =========     at 0x795d8fd7f100 of size 1 bytes
//
// The "size 1 bytes" nearest-allocation is the smoking gun for a bogus
// descriptor being passed to the search kernel — the dataset_descriptor
// materialized from the hand-stitched vpq_f16_index doesn't match what
// search_multi_cta expects.
//
// Environment:
//   - libcuvs 26.10.00a206 (nightly, cuda13 build 260819055025_2140532c)
//   - libraft 26.10.00a26
//   - CUDA 13.3
//   - Hardware: L40S (compute 8.9). Not device-specific — the OOB fires
//     on any Ampere+ card.
//
// Build:
//   nvcc -std=c++20 -O2 --extended-lambda \
//        -I$CONDA_PREFIX/include -I$CONDA_PREFIX/include/rapids \
//        -I$CONDA_PREFIX/include/raft -I$CONDA_PREFIX/include/cuvs \
//        -L$CONDA_PREFIX/lib -lcuvs -lcuvs_c -lstdc++ \
//        repro_cuvs_vpq_search_oob.cu -o repro
//
// Reproduce:
//   LD_LIBRARY_PATH=$CONDA_PREFIX/lib \
//     compute-sanitizer --tool memcheck --print-limit 3 ./repro
// ----------------------------------------------------------------------------


#include <cuvs/neighbors/cagra.hpp>
#include <cuvs/neighbors/common.hpp>
#include <cuvs/preprocessing/quantize/pq.hpp>

#include <raft/core/device_mdarray.hpp>
#include <raft/core/device_resources.hpp>
#include <raft/core/host_mdarray.hpp>
#include <raft/core/mdspan.hpp>
#include <raft/core/resource/cuda_stream.hpp>

#include <cuda_fp16.h>
#include <cuda_runtime.h>

#include <cstdint>
#include <cstdio>
#include <cstring>
#include <memory>
#include <random>
#include <vector>

namespace {

constexpr int64_t  kRows    = 10000;
constexpr int64_t  kDim     = 128;   // any value that PQ can chunk (multiple of pq_dim)
constexpr int64_t  kQueries = 16;
constexpr uint32_t kTopK    = 10;

constexpr uint32_t kIntGDeg  = 64;
constexpr uint32_t kGraphDeg = 32;
constexpr uint32_t kPqBits   = 8;
constexpr uint32_t kPqDim    = 32;   // must divide kDim

#define CUDA_CHECK(call)                                                              \
    do {                                                                              \
        cudaError_t _e = (call);                                                      \
        if (_e != cudaSuccess) {                                                      \
            std::fprintf(stderr, "CUDA error %s at %s:%d\n",                          \
                         cudaGetErrorString(_e), __FILE__, __LINE__);                 \
            std::exit(1);                                                             \
        }                                                                             \
    } while (0)

void fill_synthetic(half* dst, int64_t n, int64_t dim, uint32_t seed)
{
    std::mt19937 rng(seed);
    std::uniform_real_distribution<float> uni(-1.f, 1.f);
    std::normal_distribution<float>       jitter(0.f, 0.05f);
    constexpr int64_t kClusters = 32;
    std::vector<float> centers(static_cast<size_t>(kClusters) * dim);
    for (auto& v : centers) v = uni(rng);
    for (int64_t i = 0; i < n; ++i) {
        const float* c = centers.data() + (i % kClusters) * dim;
        for (int64_t j = 0; j < dim; ++j)
            dst[i * dim + j] = __float2half(c[j] + jitter(rng));
    }
}

}  // namespace

int main()
{
    CUDA_CHECK(cudaSetDevice(0));
    raft::device_resources res;
    auto stream = raft::resource::get_cuda_stream(res);

    // ---- 1. dense f16 rows, plain host ------------------------------------
    std::vector<half> host(static_cast<size_t>(kRows) * kDim);
    fill_synthetic(host.data(), kRows, kDim, /*seed=*/42);
    auto host_view = raft::make_host_matrix_view<const half, int64_t>(
        host.data(), kRows, kDim);

    // ---- 2. build cagra graph on host rows (graph-only) -------------------
    cuvs::neighbors::cagra::index_params ip;
    ip.metric                    = cuvs::distance::DistanceType::L2Expanded;
    ip.intermediate_graph_degree = kIntGDeg;
    ip.graph_degree              = kGraphDeg;
    ip.attach_dataset_on_build   = false;  // ignored for host builds
    ip.graph_build_params        = cuvs::neighbors::cagra::graph_build_params::ivf_pq_params(
        raft::matrix_extent<int64_t>(kRows, kDim));

    auto host_pv = cuvs::neighbors::make_host_padded_dataset_view(host_view);
    auto graph_idx = cuvs::neighbors::cagra::build(res, ip, host_pv);
    std::printf("graph built: size=%ld dim=%u graph_degree=%u\n",
                static_cast<long>(graph_idx.size()),
                graph_idx.dim(), graph_idx.graph_degree());

    // ---- 3. train VPQ codes on the same host rows -------------------------
    cuvs::neighbors::vpq_params vp;
    vp.pq_bits = kPqBits;
    vp.pq_dim  = kPqDim;
    auto vpq_owner = std::make_unique<cuvs::neighbors::device_vpq_dataset<half, int64_t>>(
        cuvs::preprocessing::quantize::pq::make_vpq_dataset(res, vp, host_view));
    std::printf("vpq trained\n");

    // ---- 4. hand-stitch vpq_f16_index (graph copied by update_graph in ctor)
    auto vpq_idx = cuvs::neighbors::cagra::vpq_f16_index<half, uint32_t>(
        res,
        cuvs::distance::DistanceType::L2Expanded,
        vpq_owner->as_dataset_view(),
        raft::make_const_mdspan(graph_idx.graph()));
    std::printf("vpq_idx: size=%ld dim=%u graph_degree=%u\n",
                static_cast<long>(vpq_idx.size()),
                vpq_idx.dim(), vpq_idx.graph_degree());
    // graph_idx is redundant now — the vpq_idx ctor copied its graph.
    { auto drop = std::move(graph_idx); }

    // ---- 5. search — OOBs under sanitizer regardless of algo --------------
    std::vector<half> q_host(static_cast<size_t>(kQueries) * kDim);
    fill_synthetic(q_host.data(), kQueries, kDim, /*seed=*/1);
    // upload queries
    half* q_dev = nullptr;
    CUDA_CHECK(cudaMalloc(&q_dev, q_host.size() * sizeof(half)));
    CUDA_CHECK(cudaMemcpy(q_dev, q_host.data(), q_host.size() * sizeof(half),
                          cudaMemcpyHostToDevice));
    auto qv = raft::make_device_matrix_view<const half, int64_t>(q_dev, kQueries, kDim);

    auto nbr  = raft::make_device_matrix<uint32_t, int64_t>(res, kQueries, kTopK);
    auto dist = raft::make_device_matrix<float, int64_t>(res, kQueries, kTopK);

    for (auto algo : {cuvs::neighbors::cagra::search_algo::MULTI_CTA,
                      cuvs::neighbors::cagra::search_algo::SINGLE_CTA,
                      cuvs::neighbors::cagra::search_algo::MULTI_KERNEL}) {
        const char* name = (algo == cuvs::neighbors::cagra::search_algo::MULTI_CTA)   ? "MULTI_CTA"
                         : (algo == cuvs::neighbors::cagra::search_algo::SINGLE_CTA)  ? "SINGLE_CTA"
                         : "MULTI_KERNEL";
        std::printf("--- searching with algo=%s ---\n", name);
        std::fflush(stdout);

        cuvs::neighbors::cagra::search_params sp;
        sp.itopk_size = 64;
        sp.algo       = algo;

        cuvs::neighbors::cagra::search(res, sp, vpq_idx, qv, nbr.view(), dist.view());
        res.sync_stream();
        std::printf("    algo=%s: no CUDA error reported (sanitizer may still flag OOB)\n", name);
    }
    CUDA_CHECK(cudaFree(q_dev));
    return 0;
}