From 8e24da10ba2c4e1aacbbb9c7fb2699a0bc7295b9 Mon Sep 17 00:00:00 2001 From: Devin AI Date: Thu, 9 Jul 2026 06:30:45 +0000 Subject: [PATCH 01/24] feat: GPU_HNSW index using faiss::gpu::GpuIndexHNSW GPU_HNSW (and GPU_HNSW_SQ) index type via vendored faiss GpuIndexHNSW: IndexNode wiring, index-type registration, faiss_gpu_hnsw CUDA target, and recall/cosine/topk/serialize tests. Clean feature-only commit on current main; dead cooperative-distance code excluded. Signed-off-by: Devin AI --- cmake/libs/libdiskann.cmake | 4 +- cmake/libs/libfaiss.cmake | 23 + include/knowhere/comp/index_param.h | 1 + include/knowhere/index/index_table.h | 2 + src/index/hnsw/faiss_hnsw.cc | 244 ++++++++ tests/ut/test_gpu_search.cc | 381 ++++++++++++ thirdparty/faiss/faiss/gpu/CMakeLists.txt | 7 + thirdparty/faiss/faiss/gpu/GpuIndexHNSW.cu | 276 +++++++++ thirdparty/faiss/faiss/gpu/GpuIndexHNSW.h | 135 +++++ .../faiss/faiss/gpu/impl/GpuHnswBuild.cuh | 341 +++++++++++ .../faiss/faiss/gpu/impl/GpuHnswSearch.cuh | 167 ++++++ .../faiss/gpu/impl/GpuHnswSearchKernel.cuh | 558 ++++++++++++++++++ .../faiss/faiss/gpu/impl/GpuHnswTypes.cu | 200 +++++++ .../faiss/faiss/gpu/impl/GpuHnswTypes.h | 162 +++++ 14 files changed, 2499 insertions(+), 2 deletions(-) create mode 100644 thirdparty/faiss/faiss/gpu/GpuIndexHNSW.cu create mode 100644 thirdparty/faiss/faiss/gpu/GpuIndexHNSW.h create mode 100644 thirdparty/faiss/faiss/gpu/impl/GpuHnswBuild.cuh create mode 100644 thirdparty/faiss/faiss/gpu/impl/GpuHnswSearch.cuh create mode 100644 thirdparty/faiss/faiss/gpu/impl/GpuHnswSearchKernel.cuh create mode 100644 thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.cu create mode 100644 thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.h diff --git a/cmake/libs/libdiskann.cmake b/cmake/libs/libdiskann.cmake index 9c6ce5eb3..ca3b64f46 100644 --- a/cmake/libs/libdiskann.cmake +++ b/cmake/libs/libdiskann.cmake @@ -1,6 +1,6 @@ add_definitions(-DKNOWHERE_WITH_DISKANN) -find_package(Boost REQUIRED COMPONENTS program_options) -include_directories(${Boost_INCLUDE_DIR}) +find_package(Boost REQUIRED COMPONENTS program_options CONFIG) +include_directories(${Boost_INCLUDE_DIRS}) find_package(aio REQUIRED) include_directories(${AIO_INCLUDE}) find_package(fmt REQUIRED) diff --git a/cmake/libs/libfaiss.cmake b/cmake/libs/libfaiss.cmake index 056b530e6..828a47817 100644 --- a/cmake/libs/libfaiss.cmake +++ b/cmake/libs/libfaiss.cmake @@ -593,3 +593,26 @@ if(__PPC64) knowhere_utils) target_compile_definitions(faiss PRIVATE FINTEGER=int) endif() + +# GPU HNSW CUDA sources — compiled when WITH_CUVS is enabled +if(WITH_CUVS) + set(FAISS_GPU_HNSW_SRCS + thirdparty/faiss/faiss/gpu/GpuIndexHNSW.cu + thirdparty/faiss/faiss/gpu/GpuIndex.cu + thirdparty/faiss/faiss/gpu/GpuResources.cpp + thirdparty/faiss/faiss/gpu/StandardGpuResources.cpp + thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.cu + thirdparty/faiss/faiss/gpu/impl/IndexUtils.cu + thirdparty/faiss/faiss/gpu/utils/DeviceUtils.cu + thirdparty/faiss/faiss/gpu/utils/StackDeviceMemory.cpp + thirdparty/faiss/faiss/gpu/utils/Timer.cpp + ) + add_library(faiss_gpu_hnsw OBJECT ${FAISS_GPU_HNSW_SRCS}) + target_include_directories(faiss_gpu_hnsw PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/thirdparty/faiss + ${Boost_INCLUDE_DIRS} + ) + target_compile_definitions(faiss_gpu_hnsw PRIVATE FINTEGER=int) + target_link_libraries(faiss_gpu_hnsw PRIVATE CUDA::cudart) + target_link_libraries(faiss PUBLIC faiss_gpu_hnsw) +endif() diff --git a/include/knowhere/comp/index_param.h b/include/knowhere/comp/index_param.h index 9066b3641..ab26efe52 100644 --- a/include/knowhere/comp/index_param.h +++ b/include/knowhere/comp/index_param.h @@ -53,6 +53,7 @@ constexpr const char* INDEX_GPU_BRUTEFORCE = "GPU_BRUTE_FORCE"; constexpr const char* INDEX_GPU_IVFFLAT = "GPU_IVF_FLAT"; constexpr const char* INDEX_GPU_IVFPQ = "GPU_IVF_PQ"; constexpr const char* INDEX_GPU_CAGRA = "GPU_CAGRA"; +constexpr const char* INDEX_GPU_HNSW = "GPU_HNSW"; constexpr const char* INDEX_HNSW = "HNSW"; constexpr const char* INDEX_HNSW_SQ = "HNSW_SQ"; diff --git a/include/knowhere/index/index_table.h b/include/knowhere/index/index_table.h index d807305e9..88e12ce73 100644 --- a/include/knowhere/index/index_table.h +++ b/include/knowhere/index/index_table.h @@ -82,6 +82,8 @@ static std::set> legal_knowhere_index = { {IndexEnum::INDEX_GPU_CAGRA, VecType::VECTOR_FLOAT16}, {IndexEnum::INDEX_GPU_CAGRA, VecType::VECTOR_INT8}, {IndexEnum::INDEX_GPU_CAGRA, VecType::VECTOR_BINARY}, + {IndexEnum::INDEX_GPU_HNSW, VecType::VECTOR_FLOAT}, + {IndexEnum::INDEX_GPU_HNSW, VecType::VECTOR_INT8}, // hnsw {IndexEnum::INDEX_HNSW, VecType::VECTOR_FLOAT}, diff --git a/src/index/hnsw/faiss_hnsw.cc b/src/index/hnsw/faiss_hnsw.cc index 63101aa21..153ae6daf 100644 --- a/src/index/hnsw/faiss_hnsw.cc +++ b/src/index/hnsw/faiss_hnsw.cc @@ -3286,4 +3286,248 @@ KNOWHERE_SIMPLE_REGISTER_DENSE_FLOAT_ALL_GLOBAL(HNSW_PRQ, BaseFaissRegularIndexH KNOWHERE_SIMPLE_REGISTER_DENSE_INT_GLOBAL(HNSW_PRQ, BaseFaissRegularIndexHNSWPRQNodeTemplate, knowhere::feature::MMAP | knowhere::feature::MV | knowhere::feature::EMB_LIST) +#ifdef KNOWHERE_WITH_CUVS +} // namespace knowhere — temporarily close to include GPU headers at file scope +// ── GPU HNSW ───────────────────────────────────────────────────────────────── +#include +#include +namespace knowhere { // reopen namespace knowhere + +// Process-global StandardGpuResources shared by all GpuHnswIndexNode instances. +// Avoids per-segment 256 MiB pinned memory, cuBLAS handle, and CUDA stream +// allocations that accumulate to tens of GiB with many segments. +static std::shared_ptr& +GetSharedGpuResources() { + static std::once_flag flag; + static std::shared_ptr instance; + std::call_once(flag, [] { + instance = std::make_shared(); + instance->setTempMemory(0); + instance->setPinnedMemory(0); + }); + return instance; +} + +// Serialize GpuIndexHNSW construction across segments. +// StandardGpuResourcesImpl::initializeForDevice is not thread-safe; +// concurrent constructors race on the allocs_ map assertion. +static std::mutex& +GetGpuConstructionMutex() { + static std::mutex mtx; + return mtx; +} + +// Single GPU HNSW index node that handles all CPU storage formats (F32, SQ8, +// FP16, BF16) transparently. Uses faiss::gpu::GpuIndexHNSW for GPU search. +// Accepts CPU-serialized HNSW or HNSW_SQ binaries at load time. +class GpuHnswIndexNode : public BaseFaissRegularIndexHNSWNode { + public: + GpuHnswIndexNode(const int32_t& version, const Object& object) + : BaseFaissRegularIndexHNSWNode(version, object, DataFormatEnum::fp32) { + } + + static std::unique_ptr + StaticCreateConfig() { + return std::make_unique(); + } + + static bool + StaticHasRawData(const knowhere::BaseConfig& config, const IndexVersion& version) { + return true; + } + + static expected + StaticEstimateLoadResource(const uint64_t file_size_in_bytes, const int64_t num_rows, const int64_t dim, + const knowhere::BaseConfig& config, const IndexVersion& version) { + // GPU HNSW stores vectors and graph in VRAM; the CPU copy is freed + // after upload in Deserialize(). Report zero CPU memory cost so the + // Milvus segment loader does not over-commit host RAM reservations. + return Resource{.memoryCost = 0, .diskCost = 0}; + } + + std::unique_ptr + CreateConfig() const override { + return StaticCreateConfig(); + } + + std::string + Type() const override { + return IndexEnum::INDEX_GPU_HNSW; + } + + protected: + Status + TrainInternal(const DataSetPtr /*dataset*/, const Config& /*cfg*/) override { + return Status::not_implemented; + } + + public: + Status + Deserialize(const BinarySet& binset, std::shared_ptr cfg) override { + std::unique_lock lock(gpu_mutex_); + gpu_index_.reset(); + + // Accept CPU-built HNSW (F32) or HNSW_SQ (quantized) binaries. + Status status; + if (!binset.Contains(IndexEnum::INDEX_GPU_HNSW)) { + BinarySet aliased = binset; + for (const char* key : {IndexEnum::INDEX_HNSW_SQ, IndexEnum::INDEX_HNSW}) { + if (binset.Contains(key)) { + aliased.Append(IndexEnum::INDEX_GPU_HNSW, binset.GetByName(key)); + break; + } + } + status = BaseFaissRegularIndexHNSWNode::Deserialize(aliased, cfg); + } else { + status = BaseFaissRegularIndexHNSWNode::Deserialize(binset, cfg); + } + if (status != Status::success) { + return status; + } + + // Eager GPU upload via faiss::gpu::GpuIndexHNSW. + const auto* faiss_idx = GetFaissHnswIndex(); + if (faiss_idx) { + try { + // Detect metric from the FAISS index type rather than config, + // because Deserialize may be called without metric_type in the config + // (e.g. empty json defaults metric_type to L2). + bool is_cosine = + dynamic_cast(faiss_idx) != nullptr; + bool use_ip = is_cosine || (faiss_idx->metric_type == ::faiss::METRIC_INNER_PRODUCT); + + { + std::lock_guard gpu_ctor_lock(GetGpuConstructionMutex()); + gpu_resources_ = GetSharedGpuResources(); + gpu_index_ = std::make_unique(gpu_resources_.get(), faiss_idx->d, + faiss_idx->metric_type); + } + gpu_index_->copyFromWithMetric(faiss_idx, use_ip, is_cosine); + // Release CPU copy — vectors and graph are now on GPU. + indexes[0].reset(); + } catch (const std::exception& e) { + fprintf(stderr, "[gpu_hnsw] eager GPU upload failed: %s\n", e.what()); + gpu_index_.reset(); + } + } + return Status::success; + } + + expected + Search(const DataSetPtr dataset, std::unique_ptr cfg, const BitsetView& bitset, + milvus::OpContext* op_context) const override { + if (!bitset.empty() && bitset.count() > 0) { + return expected::Err(Status::invalid_args, "GPU_HNSW does not support filtered search"); + } + + // Fast path: gpu_index_ is set during Deserialize and never cleared. + if (!gpu_index_) { + std::unique_lock lock(gpu_mutex_); + if (!gpu_index_) { + const auto* faiss_idx = GetFaissHnswIndex(); + if (!faiss_idx) { + return expected::Err(Status::empty_index, "index not loaded"); + } + try { + const auto& hnsw_cfg = static_cast(*cfg); + bool is_cosine = IsMetricType(hnsw_cfg.metric_type.value(), metric::COSINE); + bool use_ip = IsMetricType(hnsw_cfg.metric_type.value(), metric::IP) || is_cosine; + + { + std::lock_guard gpu_ctor_lock(GetGpuConstructionMutex()); + gpu_resources_ = GetSharedGpuResources(); + gpu_index_ = std::make_unique(gpu_resources_.get(), faiss_idx->d, + faiss_idx->metric_type); + } + gpu_index_->copyFromWithMetric(faiss_idx, use_ip, is_cosine); + const_cast&>(indexes[0]).reset(); + } catch (const std::exception& e) { + return expected::Err(Status::cuvs_inner_error, + std::string("failed to build GPU HNSW index: ") + e.what()); + } + } + } + + const auto& hnsw_cfg = static_cast(*cfg); + auto k = hnsw_cfg.k.value(); + auto nq = dataset->GetRows(); + auto dim = dataset->GetDim(); + auto ef = hnsw_cfg.ef.value_or(200); + const auto* h_queries_raw = reinterpret_cast(dataset->GetTensor()); + + // For COSINE metric, normalize queries to unit length. + const float* h_queries = h_queries_raw; + std::unique_ptr normalized_queries; + if (IsMetricType(hnsw_cfg.metric_type.value(), metric::COSINE)) { + normalized_queries = std::make_unique(nq * dim); + for (int64_t i = 0; i < nq; i++) { + const float* src = h_queries_raw + i * dim; + float* dst = normalized_queries.get() + i * dim; + float sq_norm = 0.0f; + for (int64_t d = 0; d < dim; d++) sq_norm += src[d] * src[d]; + float inv = (sq_norm > 0.0f) ? (1.0f / std::sqrt(sq_norm)) : 1.0f; + for (int64_t d = 0; d < dim; d++) dst[d] = src[d] * inv; + } + h_queries = normalized_queries.get(); + } + + auto h_ids = std::make_unique(nq * k); + auto h_dist = std::make_unique(nq * k); + + try { + faiss::gpu::GpuHnswSearchParams gsp; + gsp.ef = ef; + gpu_index_->searchHost(nq, h_queries, k, h_dist.get(), h_ids.get(), gsp); + } catch (const std::exception& e) { + LOG_KNOWHERE_ERROR_ << "GPU_HNSW search failed: " << e.what(); + return expected::Err(Status::cuvs_inner_error, + std::string("GPU HNSW search failed: ") + e.what()); + } + + // Negate back to positive for IP and COSINE. + if (IsMetricType(hnsw_cfg.metric_type.value(), metric::IP) || + IsMetricType(hnsw_cfg.metric_type.value(), metric::COSINE)) { + for (int64_t i = 0; i < static_cast(nq * k); i++) { + h_dist[i] = -h_dist[i]; + } + } + + return GenResultDataSet(nq, k, h_ids.release(), h_dist.release()); + } + + ~GpuHnswIndexNode() override = default; + + private: + const ::faiss::cppcontrib::knowhere::IndexHNSW* + GetFaissHnswIndex() const { + if (indexes.empty() || !indexes[0]) + return nullptr; + return dynamic_cast(indexes[0].get()); + } + + mutable std::mutex gpu_mutex_; + mutable std::shared_ptr gpu_resources_; + mutable std::unique_ptr gpu_index_; +}; + +// Register GPU_HNSW in the static config map at process startup. +__attribute__((constructor)) static void +register_gpu_hnsw_static_config() { + IndexStaticFaced::Instance().RegisterStaticFunc(IndexEnum::INDEX_GPU_HNSW); + IndexStaticFaced::Instance().RegisterStaticFunc(IndexEnum::INDEX_GPU_HNSW); +} + +KNOWHERE_REGISTER_GLOBAL( + GPU_HNSW, + [](const int32_t& version, const Object& object) { return Index::Create(version, object); }, fp32, + true, feature::GPU_ANN_FLOAT_INDEX); + +KNOWHERE_REGISTER_GLOBAL( + GPU_HNSW, + [](const int32_t& version, const Object& object) { + return Index>::Create(std::make_unique(version, object)); + }, + int8, true, (feature::INT8 | feature::GPU)); +#endif // KNOWHERE_WITH_CUVS + } // namespace knowhere diff --git a/tests/ut/test_gpu_search.cc b/tests/ut/test_gpu_search.cc index 4f7f21ed7..136131912 100644 --- a/tests/ut/test_gpu_search.cc +++ b/tests/ut/test_gpu_search.cc @@ -443,5 +443,386 @@ TEST_CASE("Test All GPU Index", "[search]") { CHECK(GetRelativeLoss(gt_dist[i], dist[i]) < 0.1f); } } + + // GPU_HNSW tests: build on CPU as HNSW, serialize, deserialize as GPU_HNSW, search on GPU. + SECTION("Test GPU HNSW Search (CPU build -> GPU search)") { + // Build a CPU HNSW index + knowhere::Json hnsw_json; + hnsw_json[knowhere::meta::DIM] = dim; + hnsw_json[knowhere::meta::METRIC_TYPE] = knowhere::metric::L2; + hnsw_json[knowhere::meta::TOPK] = 1; + hnsw_json[knowhere::indexparam::HNSW_M] = 16; + hnsw_json[knowhere::indexparam::EFCONSTRUCTION] = 200; + hnsw_json[knowhere::indexparam::EF] = 200; + + auto train_ds = GenDataSet(nb, dim, seed); + auto query_ds = GenDataSet(nq, dim, seed + 2); + + auto cpu_idx = + knowhere::IndexFactory::Instance().Create(knowhere::IndexEnum::INDEX_HNSW, version).value(); + auto res = cpu_idx.Build(train_ds, hnsw_json); + REQUIRE(res == knowhere::Status::success); + + // Serialize the CPU index + knowhere::BinarySet bs; + cpu_idx.Serialize(bs); + + // Deserialize as GPU_HNSW + auto gpu_idx = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_GPU_HNSW, version) + .value(); + auto deser_res = gpu_idx.Deserialize(bs); + REQUIRE(deser_res == knowhere::Status::success); + + // Search on GPU + auto results = gpu_idx.Search(query_ds, hnsw_json, nullptr); + REQUIRE(results.has_value()); + + // Compare against brute force + auto gt = knowhere::BruteForce::Search(train_ds, query_ds, hnsw_json, nullptr); + REQUIRE(gt.has_value()); + float recall = GetKNNRecall(*gt.value(), *results.value()); + REQUIRE(recall >= 0.9f); + } + + SECTION("Test GPU HNSW Search Cosine Metric") { + knowhere::Json hnsw_json; + hnsw_json[knowhere::meta::DIM] = dim; + hnsw_json[knowhere::meta::METRIC_TYPE] = knowhere::metric::COSINE; + hnsw_json[knowhere::meta::TOPK] = 1; + hnsw_json[knowhere::indexparam::HNSW_M] = 16; + hnsw_json[knowhere::indexparam::EFCONSTRUCTION] = 200; + hnsw_json[knowhere::indexparam::EF] = 200; + + auto train_ds = GenDataSet(nb, dim, seed); + auto query_ds = GenDataSet(nq, dim, seed + 1); + + auto cpu_idx = + knowhere::IndexFactory::Instance().Create(knowhere::IndexEnum::INDEX_HNSW, version).value(); + auto res = cpu_idx.Build(train_ds, hnsw_json); + REQUIRE(res == knowhere::Status::success); + + knowhere::BinarySet bs; + cpu_idx.Serialize(bs); + + auto gpu_idx = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_GPU_HNSW, version) + .value(); + auto deser_res = gpu_idx.Deserialize(bs); + REQUIRE(deser_res == knowhere::Status::success); + + auto results = gpu_idx.Search(query_ds, hnsw_json, nullptr); + REQUIRE(results.has_value()); + + auto gt = knowhere::BruteForce::Search(train_ds, query_ds, hnsw_json, nullptr); + REQUIRE(gt.has_value()); + float recall = GetKNNRecall(*gt.value(), *results.value()); + REQUIRE(recall >= 0.65f); + } + + SECTION("Test GPU HNSW Search TopK") { + knowhere::Json hnsw_json; + hnsw_json[knowhere::meta::DIM] = dim; + hnsw_json[knowhere::meta::METRIC_TYPE] = knowhere::metric::L2; + hnsw_json[knowhere::meta::TOPK] = 1; + hnsw_json[knowhere::indexparam::HNSW_M] = 16; + hnsw_json[knowhere::indexparam::EFCONSTRUCTION] = 200; + hnsw_json[knowhere::indexparam::EF] = 200; + + auto train_ds = GenDataSet(nb, dim, seed); + auto query_ds = GenDataSet(nq, dim, seed + 2); + + auto cpu_idx = + knowhere::IndexFactory::Instance().Create(knowhere::IndexEnum::INDEX_HNSW, version).value(); + auto res = cpu_idx.Build(train_ds, hnsw_json); + REQUIRE(res == knowhere::Status::success); + + knowhere::BinarySet bs; + cpu_idx.Serialize(bs); + + auto gpu_idx = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_GPU_HNSW, version) + .value(); + auto deser_res = gpu_idx.Deserialize(bs); + REQUIRE(deser_res == knowhere::Status::success); + + const auto topk_values = { + std::make_tuple(5, 0.85f), + std::make_tuple(25, 0.85f), + std::make_tuple(100, 0.85f), + }; + + for (const auto& [topk, threshold] : topk_values) { + hnsw_json[knowhere::meta::TOPK] = topk; + auto results = gpu_idx.Search(query_ds, hnsw_json, nullptr); + REQUIRE(results.has_value()); + auto gt = knowhere::BruteForce::Search(train_ds, query_ds, hnsw_json, nullptr); + float recall = GetKNNRecall(*gt.value(), *results.value()); + REQUIRE(recall >= threshold); + } + } + + SECTION("Test GPU HNSW Serialize/Deserialize Round Trip") { + knowhere::Json hnsw_json; + hnsw_json[knowhere::meta::DIM] = dim; + hnsw_json[knowhere::meta::METRIC_TYPE] = knowhere::metric::L2; + hnsw_json[knowhere::meta::TOPK] = 1; + hnsw_json[knowhere::indexparam::HNSW_M] = 16; + hnsw_json[knowhere::indexparam::EFCONSTRUCTION] = 200; + hnsw_json[knowhere::indexparam::EF] = 200; + + auto train_ds = GenDataSet(nb, dim, seed); + + auto cpu_idx = + knowhere::IndexFactory::Instance().Create(knowhere::IndexEnum::INDEX_HNSW, version).value(); + auto res = cpu_idx.Build(train_ds, hnsw_json); + REQUIRE(res == knowhere::Status::success); + + knowhere::BinarySet bs; + cpu_idx.Serialize(bs); + + auto gpu_idx = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_GPU_HNSW, version) + .value(); + auto deser_res = gpu_idx.Deserialize(bs); + REQUIRE(deser_res == knowhere::Status::success); + + // Self-search: each vector should find itself as nearest neighbor + auto results = gpu_idx.Search(train_ds, hnsw_json, nullptr); + REQUIRE(results.has_value()); + auto ids = results.value()->GetIds(); + int correct = 0; + for (int i = 0; i < nq; ++i) { + if (ids[i] == i) + correct++; + } + float self_recall = static_cast(correct) / nq; + REQUIRE(self_recall >= 0.95f); + } + + SECTION("Test GPU HNSW SQ8 Deserialization") { + // Build a CPU HNSW_SQ index, then load as GPU_HNSW + knowhere::Json hnsw_json; + hnsw_json[knowhere::meta::DIM] = dim; + hnsw_json[knowhere::meta::METRIC_TYPE] = knowhere::metric::L2; + hnsw_json[knowhere::meta::TOPK] = 1; + hnsw_json[knowhere::indexparam::HNSW_M] = 16; + hnsw_json[knowhere::indexparam::EFCONSTRUCTION] = 200; + hnsw_json[knowhere::indexparam::EF] = 200; + + auto train_ds = GenDataSet(nb, dim, seed); + auto query_ds = GenDataSet(nq, dim, seed + 2); + + auto cpu_idx = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_HNSW_SQ, version) + .value(); + auto res = cpu_idx.Build(train_ds, hnsw_json); + REQUIRE(res == knowhere::Status::success); + + knowhere::BinarySet bs; + cpu_idx.Serialize(bs); + + // GPU_HNSW should accept HNSW_SQ binaries + auto gpu_idx = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_GPU_HNSW, version) + .value(); + auto deser_res = gpu_idx.Deserialize(bs); + REQUIRE(deser_res == knowhere::Status::success); + + auto results = gpu_idx.Search(query_ds, hnsw_json, nullptr); + REQUIRE(results.has_value()); + + auto gt = knowhere::BruteForce::Search(train_ds, query_ds, hnsw_json, nullptr); + REQUIRE(gt.has_value()); + float recall = GetKNNRecall(*gt.value(), *results.value()); + // SQ8 has lower recall than flat due to quantization + REQUIRE(recall >= 0.7f); + } +} + +TEST_CASE("Test CPU vs GPU HNSW Comparison", "[gpu_hnsw_compare]") { + using Catch::Approx; + + int64_t nb = 10000, nq = 100; + int64_t dim = 128; + int64_t seed = 42; + + auto version = GenTestVersionList(); + + auto run_comparison = [&](const std::string& metric, float min_recall, float max_dist_drift) { + CAPTURE(metric, nb, nq, dim); + + knowhere::Json hnsw_json; + hnsw_json[knowhere::meta::DIM] = dim; + hnsw_json[knowhere::meta::METRIC_TYPE] = metric; + hnsw_json[knowhere::meta::TOPK] = 10; + hnsw_json[knowhere::indexparam::HNSW_M] = 16; + hnsw_json[knowhere::indexparam::EFCONSTRUCTION] = 200; + hnsw_json[knowhere::indexparam::EF] = 200; + + auto train_ds = GenDataSet(nb, dim, seed); + auto query_ds = GenDataSet(nq, dim, seed + 2); + + // --- Build CPU HNSW --- + auto cpu_idx = + knowhere::IndexFactory::Instance().Create(knowhere::IndexEnum::INDEX_HNSW, version).value(); + auto res = cpu_idx.Build(train_ds, hnsw_json); + REQUIRE(res == knowhere::Status::success); + + // Serialize for GPU deserialization + knowhere::BinarySet bs; + cpu_idx.Serialize(bs); + + // --- CPU Search with timing --- + auto cpu_start = std::chrono::high_resolution_clock::now(); + auto cpu_results = cpu_idx.Search(query_ds, hnsw_json, nullptr); + auto cpu_end = std::chrono::high_resolution_clock::now(); + REQUIRE(cpu_results.has_value()); + double cpu_ms = std::chrono::duration(cpu_end - cpu_start).count(); + + // --- Build GPU HNSW from serialized CPU index --- + auto gpu_idx = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_GPU_HNSW, version) + .value(); + auto deser_res = gpu_idx.Deserialize(bs); + REQUIRE(deser_res == knowhere::Status::success); + + // Warm-up search (first GPU call has kernel launch overhead) + auto warmup = gpu_idx.Search(query_ds, hnsw_json, nullptr); + REQUIRE(warmup.has_value()); + + // --- GPU Search with timing --- + auto gpu_start = std::chrono::high_resolution_clock::now(); + auto gpu_results = gpu_idx.Search(query_ds, hnsw_json, nullptr); + auto gpu_end = std::chrono::high_resolution_clock::now(); + REQUIRE(gpu_results.has_value()); + double gpu_ms = std::chrono::duration(gpu_end - gpu_start).count(); + + // --- Brute force ground truth --- + auto gt = knowhere::BruteForce::Search(train_ds, query_ds, hnsw_json, nullptr); + REQUIRE(gt.has_value()); + + // --- Recall comparison --- + float cpu_recall = GetKNNRecall(*gt.value(), *cpu_results.value()); + float gpu_recall = GetKNNRecall(*gt.value(), *gpu_results.value()); + + // --- Distance accuracy: compare GPU distances against CPU distances --- + auto cpu_dist = cpu_results.value()->GetDistance(); + auto gpu_dist = gpu_results.value()->GetDistance(); + auto k = hnsw_json[knowhere::meta::TOPK].get(); + + double total_rel_error = 0.0; + int valid_pairs = 0; + for (int64_t i = 0; i < nq * k; ++i) { + if (std::abs(cpu_dist[i]) > 1e-6f) { + total_rel_error += std::abs((gpu_dist[i] - cpu_dist[i]) / cpu_dist[i]); + valid_pairs++; + } + } + double avg_rel_dist_error = (valid_pairs > 0) ? (total_rel_error / valid_pairs) : 0.0; + + // --- ID overlap: how many of the same results are returned --- + auto cpu_ids = cpu_results.value()->GetIds(); + auto gpu_ids = gpu_results.value()->GetIds(); + int id_overlap = 0; + for (int64_t q = 0; q < nq; ++q) { + std::set cpu_set(cpu_ids + q * k, cpu_ids + (q + 1) * k); + for (int64_t j = 0; j < k; ++j) { + if (cpu_set.count(gpu_ids[q * k + j])) { + id_overlap++; + } + } + } + float id_overlap_ratio = static_cast(id_overlap) / (nq * k); + + // --- Print comparison report --- + fprintf(stderr, "\n=== CPU vs GPU HNSW Comparison (%s, nb=%ld, nq=%ld, dim=%ld, k=%ld) ===\n", metric.c_str(), + (long)nb, (long)nq, (long)dim, (long)k); + fprintf(stderr, " CPU recall@%ld: %.4f\n", (long)k, cpu_recall); + fprintf(stderr, " GPU recall@%ld: %.4f\n", (long)k, gpu_recall); + fprintf(stderr, " Recall delta (GPU - CPU): %+.4f\n", gpu_recall - cpu_recall); + fprintf(stderr, " CPU search time: %.2f ms\n", cpu_ms); + fprintf(stderr, " GPU search time: %.2f ms (includes H2D/D2H transfers)\n", gpu_ms); + fprintf(stderr, " Speedup: %.2fx\n", cpu_ms / gpu_ms); + fprintf(stderr, " Avg relative distance error: %.6f\n", avg_rel_dist_error); + fprintf(stderr, " ID overlap (CPU vs GPU): %.4f (%d/%ld)\n", id_overlap_ratio, id_overlap, (long)(nq * k)); + fprintf(stderr, "===\n\n"); + + // --- Assertions --- + // GPU recall should be close to CPU recall (within a small tolerance) + REQUIRE(gpu_recall >= min_recall); + // GPU should return at least 80% of the same IDs as CPU + REQUIRE(id_overlap_ratio >= 0.8f); + // Average relative distance error should be small + REQUIRE(avg_rel_dist_error <= max_dist_drift); + }; + + SECTION("L2 metric comparison") { + run_comparison(knowhere::metric::L2, 0.85f, 0.05); + } + + SECTION("IP metric comparison") { + run_comparison(knowhere::metric::IP, 0.85f, 0.05); + } + + SECTION("COSINE metric comparison") { + run_comparison(knowhere::metric::COSINE, 0.60f, 0.10); + } + + SECTION("Varying ef comparison") { + knowhere::Json hnsw_json; + hnsw_json[knowhere::meta::DIM] = dim; + hnsw_json[knowhere::meta::METRIC_TYPE] = knowhere::metric::L2; + hnsw_json[knowhere::meta::TOPK] = 10; + hnsw_json[knowhere::indexparam::HNSW_M] = 16; + hnsw_json[knowhere::indexparam::EFCONSTRUCTION] = 200; + + auto train_ds = GenDataSet(nb, dim, seed); + auto query_ds = GenDataSet(nq, dim, seed + 2); + + auto cpu_idx = + knowhere::IndexFactory::Instance().Create(knowhere::IndexEnum::INDEX_HNSW, version).value(); + auto res = cpu_idx.Build(train_ds, hnsw_json); + REQUIRE(res == knowhere::Status::success); + + knowhere::BinarySet bs; + cpu_idx.Serialize(bs); + + auto gpu_idx = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_GPU_HNSW, version) + .value(); + auto deser_res = gpu_idx.Deserialize(bs); + REQUIRE(deser_res == knowhere::Status::success); + + auto gt = knowhere::BruteForce::Search(train_ds, query_ds, hnsw_json, nullptr); + REQUIRE(gt.has_value()); + + fprintf(stderr, "\n=== ef sweep (L2, nb=%ld, nq=%ld, k=10) ===\n", (long)nb, (long)nq); + fprintf(stderr, " %6s %10s %10s %10s %10s\n", "ef", "cpu_recall", "gpu_recall", "cpu_ms", "gpu_ms"); + + for (int ef : {16, 32, 64, 128, 256, 512}) { + hnsw_json[knowhere::indexparam::EF] = ef; + + auto t0 = std::chrono::high_resolution_clock::now(); + auto cpu_res = cpu_idx.Search(query_ds, hnsw_json, nullptr); + auto t1 = std::chrono::high_resolution_clock::now(); + auto gpu_res = gpu_idx.Search(query_ds, hnsw_json, nullptr); + auto t2 = std::chrono::high_resolution_clock::now(); + + REQUIRE(cpu_res.has_value()); + REQUIRE(gpu_res.has_value()); + + float cpu_recall = GetKNNRecall(*gt.value(), *cpu_res.value()); + float gpu_recall = GetKNNRecall(*gt.value(), *gpu_res.value()); + double cpu_ms = std::chrono::duration(t1 - t0).count(); + double gpu_ms = std::chrono::duration(t2 - t1).count(); + + fprintf(stderr, " %6d %10.4f %10.4f %10.2f %10.2f\n", ef, cpu_recall, gpu_recall, cpu_ms, gpu_ms); + + // GPU recall should be reasonable at all ef values + REQUIRE(gpu_recall >= 0.5f); + } + fprintf(stderr, "===\n\n"); + } } #endif diff --git a/thirdparty/faiss/faiss/gpu/CMakeLists.txt b/thirdparty/faiss/faiss/gpu/CMakeLists.txt index 50740450a..59c7e7eab 100644 --- a/thirdparty/faiss/faiss/gpu/CMakeLists.txt +++ b/thirdparty/faiss/faiss/gpu/CMakeLists.txt @@ -29,9 +29,11 @@ set(FAISS_GPU_SRC GpuIndexIVF.cu GpuIndexIVFFlat.cu GpuIndexIVFPQ.cu + GpuIndexHNSW.cu GpuIndexIVFScalarQuantizer.cu GpuResources.cpp StandardGpuResources.cpp + impl/GpuHnswTypes.cu impl/BinaryDistance.cu impl/BinaryFlatIndex.cu impl/BroadcastSum.cu @@ -96,10 +98,15 @@ set(FAISS_GPU_HEADERS GpuIndexIVF.h GpuIndexIVFFlat.h GpuIndexIVFPQ.h + GpuIndexHNSW.h GpuIndexIVFScalarQuantizer.h GpuIndicesOptions.h GpuResources.h StandardGpuResources.h + impl/GpuHnswBuild.cuh + impl/GpuHnswSearch.cuh + impl/GpuHnswSearchKernel.cuh + impl/GpuHnswTypes.h impl/BinaryDistance.cuh impl/BinaryFlatIndex.cuh impl/BroadcastSum.cuh diff --git a/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.cu b/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.cu new file mode 100644 index 000000000..e92db6d49 --- /dev/null +++ b/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.cu @@ -0,0 +1,276 @@ +// @lint-ignore-every LICENSELINT +/** + * Copyright (c) Facebook, Inc. and its affiliates. + * + * This source code is licensed under the MIT license found in the + * LICENSE file in the root directory of this source tree. + */ +/* + * Copyright (c) 2026, 6sense Insights Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +namespace faiss { +namespace gpu { + +GpuIndexHNSW::GpuIndexHNSW( + GpuResourcesProvider* provider, + int dims, + faiss::MetricType metric, + GpuIndexHNSWConfig config) + : GpuIndex(provider->getResources(), dims, metric, 0.0f, config), + hnswConfig_(config) { + this->is_trained = false; +} + +GpuIndexHNSW::~GpuIndexHNSW() = default; + +void GpuIndexHNSW::copyFrom( + const faiss::cppcontrib::knowhere::IndexHNSW* index) { + FAISS_THROW_IF_NOT_MSG(index, "index must not be null"); + FAISS_THROW_IF_NOT_MSG(index->ntotal > 0, "index must not be empty"); + + DeviceScope scope(config_.device); + + this->d = index->d; + this->metric_type = index->metric_type; + this->ntotal = index->ntotal; + + // Detect cosine from index type (IndexHNSWFlatCosine / IndexHNSWSQCosine + // implement HasInverseL2Norms). + bool is_cosine = + dynamic_cast( + index) != nullptr; + bool use_ip = + is_cosine || (index->metric_type == faiss::METRIC_INNER_PRODUCT); + + if (dynamic_cast(index->storage)) { + deviceIndex_ = from_faiss_hnsw_sq(*index, use_ip, is_cosine); + } else { + deviceIndex_ = from_faiss_hnsw_flat(*index, use_ip, is_cosine); + } + + this->is_trained = true; +} + +void GpuIndexHNSW::copyFromWithMetric( + const faiss::cppcontrib::knowhere::IndexHNSW* index, + bool use_ip, + bool is_cosine) { + FAISS_THROW_IF_NOT_MSG(index, "index must not be null"); + FAISS_THROW_IF_NOT_MSG(index->ntotal > 0, "index must not be empty"); + + DeviceScope scope(config_.device); + + this->d = index->d; + this->metric_type = index->metric_type; + this->ntotal = index->ntotal; + + if (dynamic_cast(index->storage)) { + deviceIndex_ = from_faiss_hnsw_sq(*index, use_ip, is_cosine); + } else { + deviceIndex_ = from_faiss_hnsw_flat(*index, use_ip, is_cosine); + } + + this->is_trained = true; +} + +void GpuIndexHNSW::reset() { + deviceIndex_.reset(); + this->ntotal = 0; + this->is_trained = false; +} + +void GpuIndexHNSW::setSearchParams(const GpuHnswSearchParams& params) const { + std::lock_guard lock(searchParamsMutex_); + directSearchParams_ = params; + hasDirectSearchParams_ = true; +} + +bool GpuIndexHNSW::addImplRequiresIDs_() const { + return false; +} + +void GpuIndexHNSW::addImpl_(idx_t, const float*, const idx_t*) { + FAISS_THROW_MSG( + "GpuIndexHNSW does not support add(). " + "Build on CPU with IndexHNSW, then call copyFrom()."); +} + +void GpuIndexHNSW::searchImpl_( + idx_t n, + const float* x, + int k, + float* distances, + idx_t* labels, + const SearchParameters* search_params) const { + FAISS_THROW_IF_NOT_MSG( + this->is_trained && deviceIndex_, + "Index not loaded. Call copyFrom() first."); + FAISS_THROW_IF_NOT_MSG(n > 0, "n must be > 0"); + + auto& idx = *deviceIndex_; + + GpuHnswSearchParams sp; + bool got_params = false; + + // Prefer direct params set via setSearchParams() — avoids dynamic_cast. + { + std::lock_guard lock(searchParamsMutex_); + if (hasDirectSearchParams_) { + sp = directSearchParams_; + hasDirectSearchParams_ = false; + got_params = true; + } + } + + // Fallback: try dynamic_cast from SearchParameters. + if (!got_params && search_params) { + auto* params = + dynamic_cast(search_params); + if (params) { + sp.ef = params->ef; + sp.search_width = params->search_width; + sp.max_iterations = params->max_iterations; + sp.thread_block_size = params->thread_block_size; + sp.overflow_factor = params->overflow_factor; + got_params = true; + } + } + + ScratchPoolGuard guard(*idx.scratch_pool); + auto* slot = guard.get(); + auto& sc = slot->scratch; + cudaStream_t stream = slot->stream; + + int nq = static_cast(n); + int dim = static_cast(idx.dim); + int overflow_ef = sp.overflow_factor * sp.ef; + sc.ensure(nq, k, dim, static_cast(idx.n_rows), overflow_ef); + + // D2D: query vectors (GpuIndex::search passes device pointers) + GPU_HNSW_CUDA_CHECK(cudaMemcpyAsync( + sc.d_queries, + x, + static_cast(nq) * dim * sizeof(float), + cudaMemcpyDeviceToDevice, + stream)); + + gpu_hnsw_search(stream, sp, idx, sc, nq, k); + + // D2D: distances (output is a device pointer from GpuIndex::search) + GPU_HNSW_CUDA_CHECK(cudaMemcpyAsync( + distances, + sc.d_distances, + static_cast(nq) * k * sizeof(float), + cudaMemcpyDeviceToDevice, + stream)); + + // Labels: D2H stage (uint64_t→idx_t conversion), then H2D back + auto tmp = std::make_unique(nq * k); + GPU_HNSW_CUDA_CHECK(cudaMemcpyAsync( + tmp.get(), + sc.d_neighbors, + static_cast(nq) * k * sizeof(uint64_t), + cudaMemcpyDeviceToHost, + stream)); + GPU_HNSW_CUDA_CHECK(cudaStreamSynchronize(stream)); + + auto h_labels = std::make_unique(nq * k); + for (int i = 0; i < nq * k; i++) { + h_labels[i] = (tmp[i] == UINT64_MAX) ? -1 : static_cast(tmp[i]); + } + + GPU_HNSW_CUDA_CHECK(cudaMemcpyAsync( + labels, + h_labels.get(), + static_cast(nq) * k * sizeof(idx_t), + cudaMemcpyHostToDevice, + stream)); + GPU_HNSW_CUDA_CHECK(cudaStreamSynchronize(stream)); +} + +void GpuIndexHNSW::searchHost( + idx_t n, + const float* x_host, + int k, + float* distances_host, + idx_t* labels_host, + const GpuHnswSearchParams& sp) const { + FAISS_THROW_IF_NOT_MSG( + this->is_trained && deviceIndex_, + "Index not loaded. Call copyFrom() first."); + FAISS_THROW_IF_NOT_MSG(n > 0, "n must be > 0"); + + GPU_HNSW_CUDA_CHECK(cudaSetDevice(config_.device)); + DeviceScope scope(config_.device); + auto& idx = *deviceIndex_; + + ScratchPoolGuard guard(*idx.scratch_pool); + auto* slot = guard.get(); + auto& sc = slot->scratch; + cudaStream_t stream = slot->stream; + + int nq = static_cast(n); + int dim = static_cast(idx.dim); + int overflow_ef = sp.overflow_factor * sp.ef; + sc.ensure(nq, k, dim, static_cast(idx.n_rows), overflow_ef); + + GPU_HNSW_CUDA_CHECK(cudaMemcpyAsync( + sc.d_queries, + x_host, + static_cast(nq) * dim * sizeof(float), + cudaMemcpyDefault, + stream)); + + gpu_hnsw_search(stream, sp, idx, sc, nq, k); + + GPU_HNSW_CUDA_CHECK(cudaMemcpyAsync( + distances_host, + sc.d_distances, + static_cast(nq) * k * sizeof(float), + cudaMemcpyDeviceToHost, + stream)); + + auto tmp = std::make_unique(nq * k); + GPU_HNSW_CUDA_CHECK(cudaMemcpyAsync( + tmp.get(), + sc.d_neighbors, + static_cast(nq) * k * sizeof(uint64_t), + cudaMemcpyDeviceToHost, + stream)); + + GPU_HNSW_CUDA_CHECK(cudaStreamSynchronize(stream)); + + for (int i = 0; i < nq * k; i++) { + labels_host[i] = + (tmp[i] == UINT64_MAX) ? -1 : static_cast(tmp[i]); + } +} + +} // namespace gpu +} // namespace faiss diff --git a/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.h b/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.h new file mode 100644 index 000000000..a6fde1701 --- /dev/null +++ b/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.h @@ -0,0 +1,135 @@ +// @lint-ignore-every LICENSELINT +/** + * Copyright (c) Facebook, Inc. and its affiliates. + * + * This source code is licensed under the MIT license found in the + * LICENSE file in the root directory of this source tree. + */ +/* + * Copyright (c) 2026, 6sense Insights Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include + +#include +#include + +namespace faiss { +namespace cppcontrib { +namespace knowhere { +struct IndexHNSW; +} // namespace knowhere +} // namespace cppcontrib +namespace gpu { + +struct GpuHnswDeviceIndex; + +struct GpuIndexHNSWConfig : public GpuIndexConfig {}; + +struct SearchParametersGpuHNSW : SearchParameters { + /// Search ef — number of candidates maintained during search. + /// Higher values improve recall at the cost of speed. + int ef = 200; + + /// Number of candidates to expand per iteration. + int search_width = 4; + + /// Maximum search iterations (0 = auto). + int max_iterations = 0; + + /// Thread block size (0 = auto, default 128). + int thread_block_size = 0; + + /// Overflow queue factor: overflow_ef = overflow_factor * ef. + int overflow_factor = 2; +}; + +/// GPU implementation of HNSW search. +/// +/// This index type does NOT build an HNSW graph on GPU — it takes a +/// CPU-built faiss::IndexHNSW (Flat or SQ storage), converts the graph +/// to a GPU-friendly dense format, and runs the search on GPU using an +/// Overflow Candidate Queue (OCQ) beam search kernel. +/// +/// Supports L2, inner product, and cosine metrics. +/// Supports float32 and int8 (QT_8bit_direct_signed) data. +struct GpuIndexHNSW : public GpuIndex { + public: + GpuIndexHNSW( + GpuResourcesProvider* provider, + int dims, + faiss::MetricType metric = faiss::METRIC_L2, + GpuIndexHNSWConfig config = GpuIndexHNSWConfig()); + + ~GpuIndexHNSW() override; + + /// Load an HNSW index from CPU to GPU. + /// The CPU index must have been built and trained already. + /// Supports IndexHNSWFlat (float32) and IndexHNSWSQ + /// (QT_8bit_direct_signed for INT8, or dequantized for other SQ types). + void copyFrom(const faiss::cppcontrib::knowhere::IndexHNSW* index); + + /// Load with explicit metric specification. + void copyFromWithMetric( + const faiss::cppcontrib::knowhere::IndexHNSW* index, + bool use_ip, + bool is_cosine); + + void reset() override; + + /// Set search parameters directly, bypassing SearchParameters. + /// Thread-safe: uses atomic/mutex internally. + void setSearchParams(const GpuHnswSearchParams& params) const; + + /// Search with host pointers directly, bypassing GpuIndex::search. + /// All input/output pointers must be host memory. + /// This avoids the GpuIndex::search_ex temp allocation chain + /// which can cause SIGSEGV from pointer lifetime issues. + void searchHost( + idx_t n, + const float* x_host, + int k, + float* distances_host, + idx_t* labels_host, + const GpuHnswSearchParams& params) const; + + protected: + bool addImplRequiresIDs_() const override; + + void addImpl_(idx_t n, const float* x, const idx_t* ids) override; + + void searchImpl_( + idx_t n, + const float* x, + int k, + float* distances, + idx_t* labels, + const SearchParameters* search_params) const override; + + private: + GpuIndexHNSWConfig hnswConfig_; + + std::unique_ptr deviceIndex_; + + mutable std::mutex searchParamsMutex_; + mutable GpuHnswSearchParams directSearchParams_; + mutable bool hasDirectSearchParams_ = false; +}; + +} // namespace gpu +} // namespace faiss diff --git a/thirdparty/faiss/faiss/gpu/impl/GpuHnswBuild.cuh b/thirdparty/faiss/faiss/gpu/impl/GpuHnswBuild.cuh new file mode 100644 index 000000000..8c4ea63fc --- /dev/null +++ b/thirdparty/faiss/faiss/gpu/impl/GpuHnswBuild.cuh @@ -0,0 +1,341 @@ +// @lint-ignore-every LICENSELINT +/** + * Copyright (c) Facebook, Inc. and its affiliates. + * + * This source code is licensed under the MIT license found in the + * LICENSE file in the root directory of this source tree. + */ +/* + * Copyright (c) 2026, 6sense Insights Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Converts a FAISS HNSW index (CSR graph) + vectors into a +// GpuHnswDeviceIndex on device memory. +// +// This version is adapted for Knowhere's FAISS fork which uses +// faiss::cppcontrib::knowhere::IndexHNSW / HNSW types. + +#pragma once + +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +#include +#include + +#define GPU_HNSW_BUILD_CUDA_CHECK(expr) \ + do { \ + cudaError_t _e = (expr); \ + if (_e != cudaSuccess) { \ + throw std::runtime_error( \ + std::string("CUDA error: ") + \ + cudaGetErrorString(_e) + " at " + __FILE__ + ":" + \ + std::to_string(__LINE__)); \ + } \ + } while (0) + +namespace faiss { +namespace gpu { + +/// Extract HNSW graph layers from a Knowhere HNSW struct. +/// Template parameter HnswT can be faiss::cppcontrib::knowhere::HNSW +/// or faiss::HNSW — any type with neighbor_range(), nb_neighbors(), +/// neighbors, levels, entry_point, max_level. +template +inline void extract_hnsw_layers( + const HnswT& hnsw, + int64_t n_rows, + std::vector& h_upper_layers, + std::vector& h_layer0_flat, + uint32_t& entry_point, + int& M, + int& max_degree0, + int& num_layers) { + const int maxM0 = hnsw.nb_neighbors(0); + const int maxM = hnsw.nb_neighbors(1); + const int max_lv = hnsw.max_level; + + entry_point = static_cast(hnsw.entry_point); + M = maxM; + max_degree0 = maxM0; + num_layers = max_lv + 1; + + // Layer 0: dense [n_rows x maxM0] + h_layer0_flat.assign(n_rows * maxM0, UINT32_MAX); + for (int64_t i = 0; i < n_rows; i++) { + size_t begin, end; + hnsw.neighbor_range(i, 0, &begin, &end); + uint32_t count = static_cast(end - begin); + for (uint32_t j = 0; j < count; j++) { + auto nb = hnsw.neighbors[begin + j]; + if (nb >= 0) + h_layer0_flat[i * maxM0 + j] = static_cast(nb); + } + } + + // Upper layers (1..max_level): sparse [num_nodes_at_L x maxM] + h_upper_layers.resize(max_lv); + for (int layer = 1; layer <= max_lv; layer++) { + auto& ul = h_upper_layers[layer - 1]; + ul.max_degree = static_cast(maxM); + + std::vector node_ids; + for (int64_t i = 0; i < n_rows; i++) { + if (hnsw.levels[i] > layer) + node_ids.push_back(static_cast(i)); + } + ul.num_nodes = static_cast(node_ids.size()); + + std::vector h_neighbors(ul.num_nodes * maxM, UINT32_MAX); + std::vector h_node_ids = node_ids; + + for (uint32_t idx = 0; idx < ul.num_nodes; idx++) { + int64_t i = node_ids[idx]; + size_t begin, end; + hnsw.neighbor_range(i, layer, &begin, &end); + uint32_t count = static_cast(end - begin); + for (uint32_t j = 0; j < count; j++) { + auto nb = hnsw.neighbors[begin + j]; + if (nb >= 0) + h_neighbors[idx * maxM + j] = static_cast(nb); + } + } + + GPU_HNSW_BUILD_CUDA_CHECK( + cudaMalloc(&ul.d_node_ids, ul.num_nodes * sizeof(uint32_t))); + GPU_HNSW_BUILD_CUDA_CHECK(cudaMemcpy( + ul.d_node_ids, + h_node_ids.data(), + ul.num_nodes * sizeof(uint32_t), + cudaMemcpyHostToDevice)); + + GPU_HNSW_BUILD_CUDA_CHECK(cudaMalloc( + &ul.d_neighbors, + ul.num_nodes * maxM * sizeof(uint32_t))); + GPU_HNSW_BUILD_CUDA_CHECK(cudaMemcpy( + ul.d_neighbors, + h_neighbors.data(), + ul.num_nodes * maxM * sizeof(uint32_t), + cudaMemcpyHostToDevice)); + } +} + +inline void normalize_vectors( + std::vector& h_vectors, + int64_t n_rows, + int64_t dim) { + for (int64_t i = 0; i < n_rows; i++) { + float* v = h_vectors.data() + i * dim; + float sq_norm = 0.0f; + for (int64_t d = 0; d < dim; d++) + sq_norm += v[d] * v[d]; + if (sq_norm > 0.0f) { + float inv = 1.0f / std::sqrt(sq_norm); + for (int64_t d = 0; d < dim; d++) + v[d] *= inv; + } + } +} + +template +inline void upload_graph_to_gpu( + GpuHnswDeviceIndex& idx, + const HnswT& hnsw, + int64_t n_rows) { + std::vector h_layer0_flat; + extract_hnsw_layers( + hnsw, + n_rows, + idx.upper_layers, + h_layer0_flat, + idx.entry_point, + idx.M, + idx.max_degree0, + idx.num_layers); + + size_t graph0_bytes = + static_cast(n_rows) * idx.max_degree0 * sizeof(uint32_t); + GPU_HNSW_BUILD_CUDA_CHECK(cudaMalloc(&idx.d_layer0_graph, graph0_bytes)); + GPU_HNSW_BUILD_CUDA_CHECK(cudaMemcpy( + idx.d_layer0_graph, + h_layer0_flat.data(), + graph0_bytes, + cudaMemcpyHostToDevice)); + + int num_upper = static_cast(idx.upper_layers.size()); + idx.num_upper_layers_built = num_upper; + if (num_upper > 0) { + using kernel_ptrs = hnsw_kernel::upper_layer_ptrs; + std::vector h_ptrs(num_upper); + for (int i = 0; i < num_upper; i++) { + const auto& ul = idx.upper_layers[i]; + h_ptrs[i] = { + ul.d_node_ids, ul.d_neighbors, ul.num_nodes, ul.max_degree}; + } + size_t ptrs_bytes = num_upper * sizeof(kernel_ptrs); + GPU_HNSW_BUILD_CUDA_CHECK( + cudaMalloc(&idx.d_upper_layer_ptrs, ptrs_bytes)); + GPU_HNSW_BUILD_CUDA_CHECK(cudaMemcpy( + idx.d_upper_layer_ptrs, + h_ptrs.data(), + ptrs_bytes, + cudaMemcpyHostToDevice)); + } + +} + +inline void upload_fp32_dataset( + GpuHnswDeviceIndex& idx, + std::vector& h_vectors, + int64_t n_rows, + bool is_cosine) { + int64_t dim = idx.dim; + if (is_cosine) + normalize_vectors(h_vectors, n_rows, dim); + + size_t dataset_bytes = static_cast(n_rows) * dim * sizeof(float); + GPU_HNSW_BUILD_CUDA_CHECK(cudaMalloc(&idx.d_dataset, dataset_bytes)); + GPU_HNSW_BUILD_CUDA_CHECK(cudaMemcpy( + idx.d_dataset, + h_vectors.data(), + dataset_bytes, + cudaMemcpyHostToDevice)); + idx.dataset_int8 = false; +} + +inline void upload_int8_dataset( + GpuHnswDeviceIndex& idx, + const uint8_t* codes, + int64_t n_rows, + bool is_cosine) { + int64_t dim = idx.dim; + size_t dataset_bytes = static_cast(n_rows) * dim; + + std::vector signed_codes(dataset_bytes); + for (size_t i = 0; i < dataset_bytes; i++) { + signed_codes[i] = + static_cast(static_cast(codes[i]) - 128); + } + + GPU_HNSW_BUILD_CUDA_CHECK(cudaMalloc(&idx.d_dataset, dataset_bytes)); + GPU_HNSW_BUILD_CUDA_CHECK(cudaMemcpy( + idx.d_dataset, + signed_codes.data(), + dataset_bytes, + cudaMemcpyHostToDevice)); + idx.dataset_int8 = true; + + if (is_cosine) { + std::vector h_inv_norms(n_rows); + for (int64_t i = 0; i < n_rows; i++) { + const int8_t* row = signed_codes.data() + i * dim; + float sq_norm = 0.0f; + for (int64_t d = 0; d < dim; d++) { + float v = static_cast(row[d]); + sq_norm += v * v; + } + h_inv_norms[i] = + (sq_norm > 0.0f) ? (1.0f / std::sqrt(sq_norm)) : 0.0f; + } + size_t norms_bytes = static_cast(n_rows) * sizeof(float); + GPU_HNSW_BUILD_CUDA_CHECK( + cudaMalloc(&idx.d_inv_norms, norms_bytes)); + GPU_HNSW_BUILD_CUDA_CHECK(cudaMemcpy( + idx.d_inv_norms, + h_inv_norms.data(), + norms_bytes, + cudaMemcpyHostToDevice)); + } +} + +/// Build from Knowhere's HNSW index with SQ storage. +inline std::unique_ptr from_faiss_hnsw_sq( + const faiss::cppcontrib::knowhere::IndexHNSW& hnsw_index, + bool use_ip, + bool is_cosine = false) { + const auto* sq_storage = + dynamic_cast( + hnsw_index.storage); + if (!sq_storage) + throw std::runtime_error( + "gpu_hnsw: storage is not IndexScalarQuantizer"); + + int64_t n_rows = hnsw_index.ntotal; + int64_t dim = hnsw_index.d; + + auto idx = std::make_unique(); + idx->n_rows = n_rows; + idx->dim = dim; + idx->use_ip = use_ip; + idx->scratch_pool = std::make_unique(4, 0); + + bool is_direct_signed = (sq_storage->sq.qtype == + faiss::ScalarQuantizer::QT_8bit_direct_signed); + + if (is_direct_signed) { + upload_int8_dataset(*idx, sq_storage->codes.data(), n_rows, is_cosine); + } else { + std::vector h_vectors(n_rows * dim); + sq_storage->sa_decode( + n_rows, sq_storage->codes.data(), h_vectors.data()); + upload_fp32_dataset(*idx, h_vectors, n_rows, is_cosine); + } + + upload_graph_to_gpu(*idx, hnsw_index.hnsw, n_rows); + return idx; +} + +/// Build from Knowhere's HNSW index with Flat storage. +inline std::unique_ptr from_faiss_hnsw_flat( + const faiss::cppcontrib::knowhere::IndexHNSW& hnsw_index, + bool use_ip, + bool is_cosine = false) { + const auto* flat_storage = + dynamic_cast(hnsw_index.storage); + if (!flat_storage) + throw std::runtime_error("gpu_hnsw: storage is not IndexFlat"); + + int64_t n_rows = hnsw_index.ntotal; + int64_t dim = hnsw_index.d; + + // Use reconstruct_n instead of get_xb() — get_xb() can return a + // device pointer in GPU querynode context, causing SIGSEGV when + // accessed from CPU. + std::vector h_vectors(n_rows * dim); + flat_storage->reconstruct_n(0, n_rows, h_vectors.data()); + + auto idx = std::make_unique(); + idx->n_rows = n_rows; + idx->dim = dim; + idx->use_ip = use_ip; + idx->scratch_pool = std::make_unique(4, 0); + + upload_fp32_dataset(*idx, h_vectors, n_rows, is_cosine); + upload_graph_to_gpu(*idx, hnsw_index.hnsw, n_rows); + return idx; +} + +} // namespace gpu +} // namespace faiss diff --git a/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearch.cuh b/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearch.cuh new file mode 100644 index 000000000..126c88674 --- /dev/null +++ b/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearch.cuh @@ -0,0 +1,167 @@ +// @lint-ignore-every LICENSELINT +/** + * Copyright (c) Facebook, Inc. and its affiliates. + * + * This source code is licensed under the MIT license found in the + * LICENSE file in the root directory of this source tree. + */ +/* + * Copyright (c) 2026, 6sense Insights Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include + +#include +#include +#include +#include +#include + +#include +#include + +#define GPU_HNSW_CUDA_CHECK(expr) \ + do { \ + cudaError_t _e = (expr); \ + if (_e != cudaSuccess) { \ + throw std::runtime_error( \ + std::string("CUDA error: ") + \ + cudaGetErrorString(_e) + " at " + __FILE__ + ":" + \ + std::to_string(__LINE__)); \ + } \ + } while (0) + +namespace faiss { +namespace gpu { + +inline void gpu_hnsw_search( + cudaStream_t stream, + const GpuHnswSearchParams& params, + const GpuHnswDeviceIndex& idx, + GpuHnswSearchScratch& sc, + int num_queries, + int k) { + + int ef = params.ef; + int sw = params.search_width; + int overflow_ef = params.overflow_factor * ef; + int max_iter = params.max_iterations > 0 + ? params.max_iterations + : 2 * ef / sw + 10; + int dim = static_cast(idx.dim); + int num_upper_layers = idx.num_upper_layers_built; + + auto launch_kernels = [&]( + const DataT* d_data, + const float* d_inv_norms) { + if (num_upper_layers > 0) { + auto* d_layer_ptrs = static_cast( + idx.d_upper_layer_ptrs); + + int warps_per_block = 4; + int threads_per_block = warps_per_block * 32; + int num_blocks = + (num_queries + warps_per_block - 1) / warps_per_block; + + hnsw_kernel::upper_layer_search_kernel + <<>>( + sc.d_queries, + d_data, + d_inv_norms, + d_layer_ptrs, + sc.d_entry_points, + idx.entry_point, + num_queries, + dim, + num_upper_layers, + idx.use_ip); + GPU_HNSW_CUDA_CHECK(cudaGetLastError()); + } else { + std::vector h_eps(num_queries, idx.entry_point); + GPU_HNSW_CUDA_CHECK(cudaMemcpyAsync( + sc.d_entry_points, + h_eps.data(), + num_queries * sizeof(uint32_t), + cudaMemcpyHostToDevice, + stream)); + } + + int block_size = + params.thread_block_size > 0 ? params.thread_block_size : 128; + + { + int smem_overhead = sw * idx.max_degree0 * 8 + sw * 4 + 12; + int max_ef = (49152 - smem_overhead) / 12; + if (ef > max_ef) { + ef = max_ef; + } + } + + size_t smem_size = hnsw_kernel::calc_layer0_smem_size( + ef, sw, idx.max_degree0); + int N_int = static_cast(idx.n_rows); + size_t bitmap_bytes = hnsw_kernel::calc_visited_bitmap_size( + num_queries, N_int); + + GPU_HNSW_CUDA_CHECK( + cudaMemsetAsync(sc.d_visited_bitmaps, 0, bitmap_bytes, stream)); + if (overflow_ef > 0) { + GPU_HNSW_CUDA_CHECK(cudaMemsetAsync( + sc.d_overflow_count, + 0, + static_cast(num_queries) * sizeof(int), + stream)); + } + + hnsw_kernel::layer0_beam_search_kernel + <<>>( + sc.d_queries, + d_data, + d_inv_norms, + idx.d_layer0_graph, + sc.d_entry_points, + sc.d_visited_bitmaps, + sc.d_neighbors, + sc.d_distances, + num_queries, + N_int, + dim, + idx.max_degree0, + k, + ef, + sw, + max_iter, + idx.use_ip, + overflow_ef, + sc.d_overflow_ids, + sc.d_overflow_dists, + sc.d_overflow_expanded, + sc.d_overflow_count); + GPU_HNSW_CUDA_CHECK(cudaGetLastError()); + }; + + if (idx.dataset_int8) { + launch_kernels( + static_cast(idx.d_dataset), idx.d_inv_norms); + } else { + launch_kernels( + static_cast(idx.d_dataset), idx.d_inv_norms); + } +} + +} // namespace gpu +} // namespace faiss diff --git a/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearchKernel.cuh b/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearchKernel.cuh new file mode 100644 index 000000000..84c5ba659 --- /dev/null +++ b/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearchKernel.cuh @@ -0,0 +1,558 @@ +// @lint-ignore-every LICENSELINT +/** + * Copyright (c) Facebook, Inc. and its affiliates. + * + * This source code is licensed under the MIT license found in the + * LICENSE file in the root directory of this source tree. + */ +/* + * Copyright (c) 2026, 6sense Insights Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include + +#include +#include + +namespace faiss { +namespace gpu { +namespace hnsw_kernel { + +// ============================================================================ +// Distance computation helpers (templated on dataset element type) +// ============================================================================ + +__device__ __forceinline__ float load_elem(const float* ptr, int idx) { + return __ldg(&ptr[idx]); +} +__device__ __forceinline__ float load_elem(const half* ptr, int idx) { + return __half2float(__ldg(&ptr[idx])); +} +__device__ __forceinline__ float load_elem(const int8_t* ptr, int idx) { + return static_cast(__ldg(&ptr[idx])); +} + +template +__device__ __forceinline__ float thread_l2_distance( + const float* __restrict__ query, + const DataT* __restrict__ vec, + int dim) { + float sum = 0.0f; +#pragma unroll 8 + for (int d = 0; d < dim; d++) { + float diff = query[d] - load_elem(vec, d); + sum += diff * diff; + } + return sum; +} + +template +__device__ __forceinline__ float thread_ip_distance( + const float* __restrict__ query, + const DataT* __restrict__ vec, + int dim) { + float sum = 0.0f; +#pragma unroll 8 + for (int d = 0; d < dim; d++) { + sum += query[d] * load_elem(vec, d); + } + return -sum; +} + +// ============================================================================ +// Phase 1: Upper-layer greedy search +// ============================================================================ + +struct upper_layer_ptrs { + const uint32_t* d_node_ids; + const uint32_t* d_neighbors; + uint32_t num_nodes; + uint32_t max_degree; +}; + +__device__ __forceinline__ uint32_t +binary_search_node(const uint32_t* d_node_ids, uint32_t n, uint32_t global_id) { + uint32_t lo = 0, hi = n; + while (lo < hi) { + uint32_t mid = (lo + hi) / 2; + if (__ldg(&d_node_ids[mid]) < global_id) { + lo = mid + 1; + } else { + hi = mid; + } + } + if (lo < n && __ldg(&d_node_ids[lo]) == global_id) + return lo; + return UINT32_MAX; +} + +template +__global__ void upper_layer_search_kernel( + const float* __restrict__ d_queries, + const DataT* __restrict__ d_dataset, + const float* __restrict__ d_inv_norms, + const upper_layer_ptrs* __restrict__ d_layer_ptrs, + uint32_t* __restrict__ d_entry_points, + uint32_t global_entry_point, + int num_queries, + int dim, + int num_upper_layers, + bool use_inner_product) { + int warp_id = (blockIdx.x * blockDim.x + threadIdx.x) / 32; + int lane = threadIdx.x % 32; + if (warp_id >= num_queries) + return; + + const float* query = d_queries + static_cast(warp_id) * dim; + uint32_t current = global_entry_point; + + float best_dist; + if (lane == 0) { + if (use_inner_product) { + best_dist = thread_ip_distance( + query, + d_dataset + static_cast(current) * dim, + dim); + if (d_inv_norms) + best_dist *= __ldg(&d_inv_norms[current]); + } else { + best_dist = thread_l2_distance( + query, + d_dataset + static_cast(current) * dim, + dim); + } + } + best_dist = __shfl_sync(0xffffffff, best_dist, 0); + + for (int li = num_upper_layers - 1; li >= 0; li--) { + const upper_layer_ptrs& lp = d_layer_ptrs[li]; + bool improved = true; + while (improved) { + improved = false; + uint32_t local_idx = + binary_search_node(lp.d_node_ids, lp.num_nodes, current); + if (local_idx == UINT32_MAX) + break; + + uint32_t best_nbr = UINT32_MAX; + float best_nbr_dist = best_dist; + + for (uint32_t j = lane; j < lp.max_degree; j += 32) { + uint32_t nbr = lp.d_neighbors + [static_cast(local_idx) * + lp.max_degree + + j]; + float dist = FLT_MAX; + if (nbr != UINT32_MAX) { + const DataT* nbr_vec = + d_dataset + static_cast(nbr) * dim; + if (use_inner_product) { + dist = thread_ip_distance(query, nbr_vec, dim); + if (d_inv_norms) + dist *= d_inv_norms[nbr]; + } else { + dist = thread_l2_distance(query, nbr_vec, dim); + } + } + if (dist < best_nbr_dist) { + best_nbr_dist = dist; + best_nbr = nbr; + } + } + + for (int offset = 16; offset > 0; offset >>= 1) { + float other_dist = + __shfl_down_sync(0xffffffff, best_nbr_dist, offset); + uint32_t other_id = + __shfl_down_sync(0xffffffff, best_nbr, offset); + if (other_dist < best_nbr_dist) { + best_nbr_dist = other_dist; + best_nbr = other_id; + } + } + best_nbr_dist = __shfl_sync(0xffffffff, best_nbr_dist, 0); + best_nbr = __shfl_sync(0xffffffff, best_nbr, 0); + + if (best_nbr != UINT32_MAX && best_nbr_dist < best_dist) { + best_dist = best_nbr_dist; + current = best_nbr; + improved = true; + } + } + } + + if (lane == 0) { + d_entry_points[warp_id] = current; + } +} + +// ============================================================================ +// Phase 2: Layer-0 beam search kernel with Overflow Candidate Queue (OCQ) +// ============================================================================ + +__device__ __forceinline__ bool bitmap_visit( + uint32_t* bitmap, + uint32_t node_id) { + uint32_t word = node_id >> 5; + uint32_t bit = 1u << (node_id & 31); + uint32_t old = atomicOr(&bitmap[word], bit); + return (old & bit) == 0; +} + +__device__ __forceinline__ void overflow_insert( + uint32_t* ovf_ids, + float* ovf_dists, + uint32_t* ovf_exp, + int& ovf_rc, + int overflow_ef, + uint32_t id, + float dist, + uint32_t expanded) { + if (ovf_rc >= overflow_ef && dist >= ovf_dists[ovf_rc - 1]) + return; + + int lo = 0, hi = ovf_rc; + while (lo < hi) { + int mid = (lo + hi) / 2; + if (ovf_dists[mid] < dist) + lo = mid + 1; + else + hi = mid; + } + + int insert_end = ovf_rc < overflow_ef ? ovf_rc : overflow_ef - 1; + for (int i = insert_end; i > lo; i--) { + ovf_ids[i] = ovf_ids[i - 1]; + ovf_dists[i] = ovf_dists[i - 1]; + ovf_exp[i] = ovf_exp[i - 1]; + } + ovf_ids[lo] = id; + ovf_dists[lo] = dist; + ovf_exp[lo] = expanded; + if (ovf_rc < overflow_ef) + ovf_rc++; +} + +template +__global__ void layer0_beam_search_kernel( + const float* __restrict__ d_queries, + const DataT* __restrict__ d_dataset, + const float* __restrict__ d_inv_norms, + const uint32_t* __restrict__ d_layer0_graph, + const uint32_t* __restrict__ d_entry_points, + uint32_t* __restrict__ d_visited_bitmaps, + uint64_t* __restrict__ d_neighbors, + float* __restrict__ d_distances, + int num_queries, + int N, + int dim, + int max_degree0, + int k, + int ef, + int search_width, + int max_iterations, + bool use_inner_product, + int overflow_ef, + uint32_t* __restrict__ d_overflow_ids, + float* __restrict__ d_overflow_dists, + uint32_t* __restrict__ d_overflow_expanded, + int* __restrict__ d_overflow_count) { + int query_idx = blockIdx.x; + if (query_idx >= num_queries) + return; + + const float* query = d_queries + static_cast(query_idx) * dim; + + extern __shared__ char smem[]; + + int max_staging = search_width * max_degree0; + + uint32_t* result_ids = reinterpret_cast(smem); + float* result_dists = reinterpret_cast(result_ids + ef); + uint32_t* is_expanded = reinterpret_cast(result_dists + ef); + uint32_t* staging_ids = is_expanded + ef; + float* staging_dists = reinterpret_cast(staging_ids + max_staging); + uint32_t* parent_ids = + reinterpret_cast(staging_dists + max_staging); + int* meta = reinterpret_cast(parent_ids + search_width); + + int bitmap_words = (N + 31) / 32; + uint32_t* visited_bmap = + d_visited_bitmaps + static_cast(query_idx) * bitmap_words; + + uint32_t* ovf_ids = overflow_ef > 0 + ? d_overflow_ids + static_cast(query_idx) * overflow_ef + : nullptr; + float* ovf_dists = overflow_ef > 0 + ? d_overflow_dists + static_cast(query_idx) * overflow_ef + : nullptr; + uint32_t* ovf_exp = overflow_ef > 0 + ? d_overflow_expanded + + static_cast(query_idx) * overflow_ef + : nullptr; + + for (int i = threadIdx.x; i < ef; i += blockDim.x) { + result_ids[i] = UINT32_MAX; + result_dists[i] = FLT_MAX; + is_expanded[i] = 0; + } + if (threadIdx.x == 0) { + meta[0] = 0; + meta[1] = 0; + meta[2] = 0; + if (overflow_ef > 0) + d_overflow_count[query_idx] = 0; + } + __syncthreads(); + + // --- Seed with entry point --- + uint32_t ep = d_entry_points[query_idx]; + if (threadIdx.x == 0) { + float ep_dist; + if (use_inner_product) { + ep_dist = thread_ip_distance( + query, d_dataset + static_cast(ep) * dim, dim); + if (d_inv_norms) + ep_dist *= __ldg(&d_inv_norms[ep]); + } else { + ep_dist = thread_l2_distance( + query, d_dataset + static_cast(ep) * dim, dim); + } + result_ids[0] = ep; + result_dists[0] = ep_dist; + is_expanded[0] = 0; + meta[0] = 1; + bitmap_visit(visited_bmap, ep); + } + __syncthreads(); + + // --- Seed with entry point's neighbors --- + if (threadIdx.x == 0) + meta[1] = 0; + __syncthreads(); + + for (int j = threadIdx.x; j < max_degree0; j += blockDim.x) { + uint32_t nbr = + d_layer0_graph[static_cast(ep) * max_degree0 + j]; + if (nbr == UINT32_MAX || nbr >= static_cast(N)) + continue; + if (!bitmap_visit(visited_bmap, nbr)) + continue; + + const DataT* nbr_vec = d_dataset + static_cast(nbr) * dim; + float dist; + if (use_inner_product) { + dist = thread_ip_distance(query, nbr_vec, dim); + if (d_inv_norms) + dist *= d_inv_norms[nbr]; + } else { + dist = thread_l2_distance(query, nbr_vec, dim); + } + + int slot = atomicAdd(&meta[1], 1); + if (slot < max_staging) { + staging_ids[slot] = nbr; + staging_dists[slot] = dist; + } + } + __syncthreads(); + + if (threadIdx.x == 0) { + int staging_count = min(meta[1], max_staging); + int rc = meta[0]; + + for (int s = 0; s < staging_count; s++) { + uint32_t sid = staging_ids[s]; + float sdist = staging_dists[s]; + if (rc >= ef && sdist >= result_dists[rc - 1]) + continue; + + int lo = 0, hi = rc; + while (lo < hi) { + int mid = (lo + hi) / 2; + if (result_dists[mid] < sdist) + lo = mid + 1; + else + hi = mid; + } + int insert_end = rc < ef ? rc : ef - 1; + for (int i = insert_end; i > lo; i--) { + result_ids[i] = result_ids[i - 1]; + result_dists[i] = result_dists[i - 1]; + is_expanded[i] = is_expanded[i - 1]; + } + result_ids[lo] = sid; + result_dists[lo] = sdist; + is_expanded[lo] = 0; + if (rc < ef) + rc++; + } + + for (int i = 0; i < rc; i++) { + if (result_ids[i] == ep) { + is_expanded[i] = 1; + break; + } + } + meta[0] = rc; + } + __syncthreads(); + + // --- Unified main loop --- + for (int iter = 0; iter < max_iterations; iter++) { + if (threadIdx.x == 0) { + int num_parents = 0; + int rc = meta[0]; + + for (int i = 0; i < rc && num_parents < search_width; i++) { + if (!is_expanded[i]) { + parent_ids[num_parents++] = result_ids[i]; + is_expanded[i] = 1; + } + } + + if (num_parents == 0 && overflow_ef > 0) { + int ovf_rc = d_overflow_count[query_idx]; + for (int i = 0; i < ovf_rc && num_parents < search_width; + i++) { + if (!ovf_exp[i]) { + parent_ids[num_parents++] = ovf_ids[i]; + ovf_exp[i] = 1; + } + } + } + meta[2] = num_parents; + } + __syncthreads(); + + int num_parents = meta[2]; + if (num_parents == 0) + break; + + if (threadIdx.x == 0) + meta[1] = 0; + __syncthreads(); + + int total_work = num_parents * max_degree0; + for (int wi = threadIdx.x; wi < total_work; wi += blockDim.x) { + int parent_idx = wi / max_degree0; + int nbr_slot = wi % max_degree0; + + uint32_t parent = parent_ids[parent_idx]; + uint32_t nbr = + d_layer0_graph + [static_cast(parent) * max_degree0 + + nbr_slot]; + if (nbr == UINT32_MAX || nbr >= static_cast(N)) + continue; + if (!bitmap_visit(visited_bmap, nbr)) + continue; + + const DataT* nbr_vec = d_dataset + static_cast(nbr) * dim; + float dist; + if (use_inner_product) { + dist = thread_ip_distance(query, nbr_vec, dim); + if (d_inv_norms) + dist *= d_inv_norms[nbr]; + } else { + dist = thread_l2_distance(query, nbr_vec, dim); + } + + int slot = atomicAdd(&meta[1], 1); + if (slot < max_staging) { + staging_ids[slot] = nbr; + staging_dists[slot] = dist; + } + } + __syncthreads(); + + if (threadIdx.x == 0) { + int staging_count = min(meta[1], max_staging); + int rc = meta[0]; + + for (int s = 0; s < staging_count; s++) { + uint32_t sid = staging_ids[s]; + float sdist = staging_dists[s]; + if (rc >= ef && sdist >= result_dists[rc - 1]) + continue; + + int lo = 0, hi = rc; + while (lo < hi) { + int mid = (lo + hi) / 2; + if (result_dists[mid] < sdist) + lo = mid + 1; + else + hi = mid; + } + int insert_end = rc < ef ? rc : ef - 1; + for (int i = insert_end; i > lo; i--) { + result_ids[i] = result_ids[i - 1]; + result_dists[i] = result_dists[i - 1]; + is_expanded[i] = is_expanded[i - 1]; + } + result_ids[lo] = sid; + result_dists[lo] = sdist; + is_expanded[lo] = 0; + if (rc < ef) + rc++; + } + + meta[0] = rc; + } + __syncthreads(); + + // Stagnation detected when num_parents==0 (all expanded) → break above + } + + // --- Copy top-k results to global memory --- + int rc = meta[0]; + for (int i = threadIdx.x; i < k; i += blockDim.x) { + if (i < rc) { + d_neighbors[static_cast(query_idx) * k + i] = + static_cast(result_ids[i]); + d_distances[static_cast(query_idx) * k + i] = + result_dists[i]; + } else { + d_neighbors[static_cast(query_idx) * k + i] = UINT64_MAX; + d_distances[static_cast(query_idx) * k + i] = FLT_MAX; + } + } +} + +inline size_t calc_layer0_smem_size(int ef, int search_width, int max_degree0) { + int max_staging = search_width * max_degree0; + + size_t size = 0; + size += ef * sizeof(uint32_t); + size += ef * sizeof(float); + size += ef * sizeof(uint32_t); + size += max_staging * sizeof(uint32_t); + size += max_staging * sizeof(float); + size += search_width * sizeof(uint32_t); + size += 3 * sizeof(int); + return size; +} + +inline size_t calc_visited_bitmap_size(int num_queries, int N) { + int bitmap_words = (N + 31) / 32; + return static_cast(num_queries) * bitmap_words * sizeof(uint32_t); +} + +} // namespace hnsw_kernel +} // namespace gpu +} // namespace faiss diff --git a/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.cu b/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.cu new file mode 100644 index 000000000..1cd98bbdb --- /dev/null +++ b/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.cu @@ -0,0 +1,200 @@ +// @lint-ignore-every LICENSELINT +/** + * Copyright (c) Facebook, Inc. and its affiliates. + * + * This source code is licensed under the MIT license found in the + * LICENSE file in the root directory of this source tree. + */ +/* + * Copyright (c) 2026, 6sense Insights Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include + +#include +#include + +namespace faiss { +namespace gpu { + +namespace { + +inline void check_cuda(cudaError_t err, const char* file, int line) { + if (err != cudaSuccess) { + throw std::runtime_error( + std::string("CUDA error in scratch alloc: ") + + cudaGetErrorString(err) + " at " + file + ":" + + std::to_string(line)); + } +} + +#define SCRATCH_CUDA_CHECK(expr) check_cuda((expr), __FILE__, __LINE__) + +} // namespace + +void GpuHnswSearchScratch::ensure( + int nq, + int k, + int dim, + int N, + int overflow_ef) { + size_t need_q = static_cast(nq) * dim * sizeof(float); + if (need_q > queries_bytes) { + if (d_queries) + cudaFree(d_queries); + SCRATCH_CUDA_CHECK(cudaMalloc(&d_queries, need_q)); + queries_bytes = need_q; + } + size_t need_n = static_cast(nq) * k * sizeof(uint64_t); + if (need_n > neighbors_bytes) { + if (d_neighbors) + cudaFree(d_neighbors); + SCRATCH_CUDA_CHECK(cudaMalloc(&d_neighbors, need_n)); + neighbors_bytes = need_n; + } + size_t need_d = static_cast(nq) * k * sizeof(float); + if (need_d > distances_bytes) { + if (d_distances) + cudaFree(d_distances); + SCRATCH_CUDA_CHECK(cudaMalloc(&d_distances, need_d)); + distances_bytes = need_d; + } + if (nq > entry_cap) { + if (d_entry_points) + cudaFree(d_entry_points); + SCRATCH_CUDA_CHECK(cudaMalloc( + &d_entry_points, + static_cast(nq) * sizeof(uint32_t))); + entry_cap = nq; + } + int bitmap_words = (N + 31) / 32; + size_t need_bm = + static_cast(nq) * bitmap_words * sizeof(uint32_t); + if (need_bm > bitmap_bytes) { + if (d_visited_bitmaps) + cudaFree(d_visited_bitmaps); + SCRATCH_CUDA_CHECK(cudaMalloc(&d_visited_bitmaps, need_bm)); + bitmap_bytes = need_bm; + } + if (overflow_ef > 0) { + size_t ovf_entries = static_cast(nq) * overflow_ef; + size_t need_ovf = + ovf_entries * + (sizeof(uint32_t) + sizeof(float) + sizeof(uint32_t)) + + static_cast(nq) * sizeof(int); + if (need_ovf > overflow_bytes) { + if (d_overflow_ids) + cudaFree(d_overflow_ids); + if (d_overflow_dists) + cudaFree(d_overflow_dists); + if (d_overflow_expanded) + cudaFree(d_overflow_expanded); + if (d_overflow_count) + cudaFree(d_overflow_count); + SCRATCH_CUDA_CHECK(cudaMalloc( + &d_overflow_ids, ovf_entries * sizeof(uint32_t))); + SCRATCH_CUDA_CHECK(cudaMalloc( + &d_overflow_dists, ovf_entries * sizeof(float))); + SCRATCH_CUDA_CHECK(cudaMalloc( + &d_overflow_expanded, ovf_entries * sizeof(uint32_t))); + SCRATCH_CUDA_CHECK(cudaMalloc( + &d_overflow_count, + static_cast(nq) * sizeof(int))); + overflow_bytes = need_ovf; + } + } +} + +GpuHnswSearchScratch::~GpuHnswSearchScratch() { + if (d_queries) + cudaFree(d_queries); + if (d_neighbors) + cudaFree(d_neighbors); + if (d_distances) + cudaFree(d_distances); + if (d_entry_points) + cudaFree(d_entry_points); + if (d_visited_bitmaps) + cudaFree(d_visited_bitmaps); + if (d_overflow_ids) + cudaFree(d_overflow_ids); + if (d_overflow_dists) + cudaFree(d_overflow_dists); + if (d_overflow_expanded) + cudaFree(d_overflow_expanded); + if (d_overflow_count) + cudaFree(d_overflow_count); +} + +GpuHnswScratchSlot::~GpuHnswScratchSlot() { + if (stream) + cudaStreamDestroy(stream); +} + +GpuHnswScratchPool::GpuHnswScratchPool(int pool_size, int device) + : pool_size_(pool_size), device_(device) {} + +void GpuHnswScratchPool::init_once() { + if (initialized_) + return; + slots_.reserve(pool_size_); + available_.reserve(pool_size_); + for (int i = 0; i < pool_size_; i++) { + auto slot = std::make_unique(); + SCRATCH_CUDA_CHECK(cudaSetDevice(device_)); + SCRATCH_CUDA_CHECK( + cudaStreamCreateWithFlags(&slot->stream, cudaStreamNonBlocking)); + available_.push_back(slot.get()); + slots_.push_back(std::move(slot)); + } + initialized_ = true; +} + +GpuHnswScratchSlot* GpuHnswScratchPool::acquire() { + std::unique_lock lock(mutex_); + init_once(); + cv_.wait(lock, [this] { return !available_.empty(); }); + auto* slot = available_.back(); + available_.pop_back(); + return slot; +} + +void GpuHnswScratchPool::release(GpuHnswScratchSlot* slot) { + { + std::lock_guard lock(mutex_); + available_.push_back(slot); + } + cv_.notify_one(); +} + +GpuHnswDeviceIndex::~GpuHnswDeviceIndex() { + if (d_dataset) + cudaFree(d_dataset); + if (d_inv_norms) + cudaFree(d_inv_norms); + if (d_layer0_graph) + cudaFree(d_layer0_graph); + for (auto& ul : upper_layers) { + if (ul.d_node_ids) + cudaFree(ul.d_node_ids); + if (ul.d_neighbors) + cudaFree(ul.d_neighbors); + } + if (d_upper_layer_ptrs) + cudaFree(d_upper_layer_ptrs); +} + +} // namespace gpu +} // namespace faiss diff --git a/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.h b/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.h new file mode 100644 index 000000000..86b781282 --- /dev/null +++ b/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.h @@ -0,0 +1,162 @@ +// @lint-ignore-every LICENSELINT +/** + * Copyright (c) Facebook, Inc. and its affiliates. + * + * This source code is licensed under the MIT license found in the + * LICENSE file in the root directory of this source tree. + */ +/* + * Copyright (c) 2026, 6sense Insights Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include + +#include +#include +#include +#include +#include + +namespace faiss { +namespace gpu { + +struct GpuHnswSearchParams { + int ef = 200; + int search_width = 4; + int max_iterations = 0; + int thread_block_size = 0; + int overflow_factor = 0; +}; + +struct GpuHnswDeviceUpperLayer { + uint32_t* d_node_ids = nullptr; + uint32_t* d_neighbors = nullptr; + uint32_t num_nodes = 0; + uint32_t max_degree = 0; +}; + +struct GpuHnswSearchScratch { + float* d_queries = nullptr; + uint64_t* d_neighbors = nullptr; + float* d_distances = nullptr; + uint32_t* d_entry_points = nullptr; + uint32_t* d_visited_bitmaps = nullptr; + + uint32_t* d_overflow_ids = nullptr; + float* d_overflow_dists = nullptr; + uint32_t* d_overflow_expanded = nullptr; + int* d_overflow_count = nullptr; + + size_t queries_bytes = 0; + size_t neighbors_bytes = 0; + size_t distances_bytes = 0; + int entry_cap = 0; + size_t bitmap_bytes = 0; + size_t overflow_bytes = 0; + + void ensure(int nq, int k, int dim, int N, int overflow_ef); + + ~GpuHnswSearchScratch(); + + GpuHnswSearchScratch() = default; + GpuHnswSearchScratch(const GpuHnswSearchScratch&) = delete; + GpuHnswSearchScratch& operator=(const GpuHnswSearchScratch&) = delete; +}; + +/// One slot in the scratch pool: a scratch buffer + its own CUDA stream. +struct GpuHnswScratchSlot { + GpuHnswSearchScratch scratch; + cudaStream_t stream = nullptr; + + ~GpuHnswScratchSlot(); + GpuHnswScratchSlot() = default; + GpuHnswScratchSlot(const GpuHnswScratchSlot&) = delete; + GpuHnswScratchSlot& operator=(const GpuHnswScratchSlot&) = delete; +}; + +/// Pool of scratch buffers allowing concurrent GPU searches on the same segment. +/// Each acquire() returns a slot with its own scratch + CUDA stream. +/// Callers block if all slots are in use. +class GpuHnswScratchPool { + public: + /// Create a pool. CUDA streams are allocated lazily on first acquire(). + explicit GpuHnswScratchPool(int pool_size = 4, int device = 0); + ~GpuHnswScratchPool() = default; + + GpuHnswScratchPool(const GpuHnswScratchPool&) = delete; + GpuHnswScratchPool& operator=(const GpuHnswScratchPool&) = delete; + + /// Acquire a scratch slot (blocks until one is available). + GpuHnswScratchSlot* acquire(); + /// Release a previously acquired scratch slot back to the pool. + void release(GpuHnswScratchSlot* slot); + + int pool_size() const { return pool_size_; } + + private: + void init_once(); + + std::mutex mutex_; + std::condition_variable cv_; + int pool_size_; + int device_; + bool initialized_ = false; + std::vector> slots_; + std::vector available_; +}; + +/// RAII guard: acquires a scratch slot on construction, releases on destruction. +class ScratchPoolGuard { + public: + ScratchPoolGuard(GpuHnswScratchPool& pool) + : pool_(pool), slot_(pool.acquire()) {} + ~ScratchPoolGuard() { pool_.release(slot_); } + GpuHnswScratchSlot* get() const { return slot_; } + + ScratchPoolGuard(const ScratchPoolGuard&) = delete; + ScratchPoolGuard& operator=(const ScratchPoolGuard&) = delete; + + private: + GpuHnswScratchPool& pool_; + GpuHnswScratchSlot* slot_; +}; + +struct GpuHnswDeviceIndex { + void* d_dataset = nullptr; + bool dataset_int8 = false; + float* d_inv_norms = nullptr; + uint32_t* d_layer0_graph = nullptr; + std::vector upper_layers; + + int64_t n_rows = 0; + int64_t dim = 0; + uint32_t entry_point = 0; + int num_layers = 0; + int M = 0; + int max_degree0 = 0; + bool use_ip = false; + + void* d_upper_layer_ptrs = nullptr; + int num_upper_layers_built = 0; + + mutable std::unique_ptr scratch_pool; + + ~GpuHnswDeviceIndex(); +}; + +} // namespace gpu +} // namespace faiss From b06fb3284b98fad43e3d58bd156cfc8c3524aba0 Mon Sep 17 00:00:00 2001 From: premal Date: Sat, 11 Jul 2026 19:38:37 +0000 Subject: [PATCH 02/24] fix: address Devin Review findings on GPU HNSW (vendored faiss) Mirror the 6si/faiss GPU HNSW review fixes into the vendored copy: - default overflow_factor to 2 in internal GpuHnswSearchParams - propagate config_.device into the scratch pool (no hardcoded device 0) - make setSearchParams sticky (no consume-once race) - query device shared-memory opt-in limit and opt the kernel into >48KiB dynamic shared memory instead of assuming a 48KiB cap - synchronize before the stack-local entry-point buffer is destroyed Signed-off-by: premal Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- thirdparty/faiss/faiss/gpu/GpuIndexHNSW.cu | 15 ++++++--- thirdparty/faiss/faiss/gpu/GpuIndexHNSW.h | 5 ++- .../faiss/faiss/gpu/impl/GpuHnswBuild.cuh | 10 +++--- .../faiss/faiss/gpu/impl/GpuHnswSearch.cuh | 33 ++++++++++++++++++- .../faiss/faiss/gpu/impl/GpuHnswTypes.h | 4 ++- 5 files changed, 55 insertions(+), 12 deletions(-) diff --git a/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.cu b/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.cu index e92db6d49..c389c10b6 100644 --- a/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.cu +++ b/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.cu @@ -69,9 +69,11 @@ void GpuIndexHNSW::copyFrom( is_cosine || (index->metric_type == faiss::METRIC_INNER_PRODUCT); if (dynamic_cast(index->storage)) { - deviceIndex_ = from_faiss_hnsw_sq(*index, use_ip, is_cosine); + deviceIndex_ = + from_faiss_hnsw_sq(*index, use_ip, is_cosine, config_.device); } else { - deviceIndex_ = from_faiss_hnsw_flat(*index, use_ip, is_cosine); + deviceIndex_ = + from_faiss_hnsw_flat(*index, use_ip, is_cosine, config_.device); } this->is_trained = true; @@ -91,9 +93,11 @@ void GpuIndexHNSW::copyFromWithMetric( this->ntotal = index->ntotal; if (dynamic_cast(index->storage)) { - deviceIndex_ = from_faiss_hnsw_sq(*index, use_ip, is_cosine); + deviceIndex_ = + from_faiss_hnsw_sq(*index, use_ip, is_cosine, config_.device); } else { - deviceIndex_ = from_faiss_hnsw_flat(*index, use_ip, is_cosine); + deviceIndex_ = + from_faiss_hnsw_flat(*index, use_ip, is_cosine, config_.device); } this->is_trained = true; @@ -139,11 +143,12 @@ void GpuIndexHNSW::searchImpl_( bool got_params = false; // Prefer direct params set via setSearchParams() — avoids dynamic_cast. + // These are sticky: they stay in effect for every subsequent search until + // overwritten, so a later search never silently falls back to defaults. { std::lock_guard lock(searchParamsMutex_); if (hasDirectSearchParams_) { sp = directSearchParams_; - hasDirectSearchParams_ = false; got_params = true; } } diff --git a/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.h b/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.h index a6fde1701..4437d8b1d 100644 --- a/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.h +++ b/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.h @@ -93,7 +93,10 @@ struct GpuIndexHNSW : public GpuIndex { void reset() override; /// Set search parameters directly, bypassing SearchParameters. - /// Thread-safe: uses atomic/mutex internally. + /// Mutex-guarded. The params are sticky: once set they apply to every + /// subsequent search() until overwritten by another setSearchParams() + /// call. Prefer passing SearchParametersGpuHNSW per-search (or using + /// searchHost) when different concurrent searches need different params. void setSearchParams(const GpuHnswSearchParams& params) const; /// Search with host pointers directly, bypassing GpuIndex::search. diff --git a/thirdparty/faiss/faiss/gpu/impl/GpuHnswBuild.cuh b/thirdparty/faiss/faiss/gpu/impl/GpuHnswBuild.cuh index 8c4ea63fc..5f4ca5613 100644 --- a/thirdparty/faiss/faiss/gpu/impl/GpuHnswBuild.cuh +++ b/thirdparty/faiss/faiss/gpu/impl/GpuHnswBuild.cuh @@ -274,7 +274,8 @@ inline void upload_int8_dataset( inline std::unique_ptr from_faiss_hnsw_sq( const faiss::cppcontrib::knowhere::IndexHNSW& hnsw_index, bool use_ip, - bool is_cosine = false) { + bool is_cosine = false, + int device = 0) { const auto* sq_storage = dynamic_cast( hnsw_index.storage); @@ -289,7 +290,7 @@ inline std::unique_ptr from_faiss_hnsw_sq( idx->n_rows = n_rows; idx->dim = dim; idx->use_ip = use_ip; - idx->scratch_pool = std::make_unique(4, 0); + idx->scratch_pool = std::make_unique(4, device); bool is_direct_signed = (sq_storage->sq.qtype == faiss::ScalarQuantizer::QT_8bit_direct_signed); @@ -311,7 +312,8 @@ inline std::unique_ptr from_faiss_hnsw_sq( inline std::unique_ptr from_faiss_hnsw_flat( const faiss::cppcontrib::knowhere::IndexHNSW& hnsw_index, bool use_ip, - bool is_cosine = false) { + bool is_cosine = false, + int device = 0) { const auto* flat_storage = dynamic_cast(hnsw_index.storage); if (!flat_storage) @@ -330,7 +332,7 @@ inline std::unique_ptr from_faiss_hnsw_flat( idx->n_rows = n_rows; idx->dim = dim; idx->use_ip = use_ip; - idx->scratch_pool = std::make_unique(4, 0); + idx->scratch_pool = std::make_unique(4, device); upload_fp32_dataset(*idx, h_vectors, n_rows, is_cosine); upload_graph_to_gpu(*idx, hnsw_index.hnsw, n_rows); diff --git a/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearch.cuh b/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearch.cuh index 126c88674..336e37a6f 100644 --- a/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearch.cuh +++ b/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearch.cuh @@ -98,14 +98,35 @@ inline void gpu_hnsw_search( num_queries * sizeof(uint32_t), cudaMemcpyHostToDevice, stream)); + // h_eps is stack-local; synchronize before it is destroyed so the + // copy never reads freed memory (safe even if the source buffer is + // ever switched to pinned host memory). + GPU_HNSW_CUDA_CHECK(cudaStreamSynchronize(stream)); } int block_size = params.thread_block_size > 0 ? params.thread_block_size : 128; + // Per-block dynamic shared-memory budget for this device. The default + // limit is 48 KiB, but Volta+ GPUs can opt into more via + // cudaFuncSetAttribute; query the real limit instead of assuming 48 KiB. + int smem_max = 49152; + { + int device = 0; + int optin = 0; + if (cudaGetDevice(&device) == cudaSuccess && + cudaDeviceGetAttribute( + &optin, + cudaDevAttrMaxSharedMemoryPerBlockOptin, + device) == cudaSuccess && + optin > smem_max) { + smem_max = optin; + } + } + { int smem_overhead = sw * idx.max_degree0 * 8 + sw * 4 + 12; - int max_ef = (49152 - smem_overhead) / 12; + int max_ef = (smem_max - smem_overhead) / 12; if (ef > max_ef) { ef = max_ef; } @@ -113,6 +134,16 @@ inline void gpu_hnsw_search( size_t smem_size = hnsw_kernel::calc_layer0_smem_size( ef, sw, idx.max_degree0); + + // Opt into >48 KiB dynamic shared memory when the device supports it; + // without this the kernel launch would fail for large ef. + if (smem_size > 49152) { + GPU_HNSW_CUDA_CHECK(cudaFuncSetAttribute( + hnsw_kernel::layer0_beam_search_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, + static_cast(smem_size))); + } + int N_int = static_cast(idx.n_rows); size_t bitmap_bytes = hnsw_kernel::calc_visited_bitmap_size( num_queries, N_int); diff --git a/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.h b/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.h index 86b781282..9f2db4152 100644 --- a/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.h +++ b/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.h @@ -39,7 +39,9 @@ struct GpuHnswSearchParams { int search_width = 4; int max_iterations = 0; int thread_block_size = 0; - int overflow_factor = 0; + // Keep in sync with SearchParametersGpuHNSW::overflow_factor: a zero here + // would disable the overflow candidate queue and silently drop recall. + int overflow_factor = 2; }; struct GpuHnswDeviceUpperLayer { From 452da1d177f74d6c7676aeca45354d8e71b7402c Mon Sep 17 00:00:00 2001 From: premal Date: Sat, 11 Jul 2026 20:19:01 +0000 Subject: [PATCH 03/24] fix: address Devin Review findings on GPU_HNSW index node - keep Count()/Dim() working after the CPU copy is freed by capturing n_rows/dim before reset and overriding Count()/Dim() (previously returned -1, making the loaded segment look empty) - make HasRawData()/GetVectorByIds() and StaticHasRawData() consistent: GPU_HNSW frees the CPU copy so it exposes no CPU-reconstructable raw data (return false / not_implemented once uploaded) - set gsp.overflow_factor=2 in the search path so the overflow candidate queue stays enabled (recall) - publish gpu_index_ through an atomic gpu_ready_ flag (release/acquire) to remove the unsynchronized double-checked read - harden the filtered-search guard to reject on bitset.data()!=nullptr instead of count() (which defaults to 0) - correct the class comment: GPU_HNSW is registered for F32 and INT8 only - add Count()/Dim() regression assertions to the GPU HNSW search test Signed-off-by: premal Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- src/index/hnsw/faiss_hnsw.cc | 86 ++++++++++++++++++++++++++++++++---- tests/ut/test_gpu_search.cc | 5 +++ 2 files changed, 83 insertions(+), 8 deletions(-) diff --git a/src/index/hnsw/faiss_hnsw.cc b/src/index/hnsw/faiss_hnsw.cc index 153ae6daf..f4d69a02f 100644 --- a/src/index/hnsw/faiss_hnsw.cc +++ b/src/index/hnsw/faiss_hnsw.cc @@ -21,12 +21,14 @@ #include #include +#include #include #include #include #include #include #include +#include #include #include #include @@ -3317,9 +3319,10 @@ GetGpuConstructionMutex() { return mtx; } -// Single GPU HNSW index node that handles all CPU storage formats (F32, SQ8, -// FP16, BF16) transparently. Uses faiss::gpu::GpuIndexHNSW for GPU search. -// Accepts CPU-serialized HNSW or HNSW_SQ binaries at load time. +// GPU HNSW index node. Registered for F32 and INT8 (see registrations at the +// bottom of this file); it loads CPU-serialized HNSW (F32) or HNSW_SQ (SQ8/INT8) +// binaries at load time and uploads them to the GPU via faiss::gpu::GpuIndexHNSW. +// The CPU copy is freed after upload, so this node exposes no CPU raw data. class GpuHnswIndexNode : public BaseFaissRegularIndexHNSWNode { public: GpuHnswIndexNode(const int32_t& version, const Object& object) @@ -3333,7 +3336,10 @@ class GpuHnswIndexNode : public BaseFaissRegularIndexHNSWNode { static bool StaticHasRawData(const knowhere::BaseConfig& config, const IndexVersion& version) { - return true; + // The CPU copy is freed after the GPU upload, so no CPU-reconstructable + // raw data is retained. Keep this consistent with the instance + // HasRawData() override below. + return false; } static expected @@ -3355,6 +3361,45 @@ class GpuHnswIndexNode : public BaseFaissRegularIndexHNSWNode { return IndexEnum::INDEX_GPU_HNSW; } + // After the eager GPU upload the CPU index (indexes[0]) is freed, so the + // inherited Count()/Dim()/HasRawData()/GetVectorByIds() would report an + // empty index. Serve vector count and dimension from the metadata captured + // before the CPU copy was released. + int64_t + Count() const override { + if (gpu_ready_.load(std::memory_order_acquire)) { + return gpu_ntotal_; + } + return BaseFaissRegularIndexHNSWNode::Count(); + } + + int64_t + Dim() const override { + if (gpu_ready_.load(std::memory_order_acquire)) { + return gpu_dim_; + } + return BaseFaissRegularIndexHNSWNode::Dim(); + } + + bool + HasRawData(const std::string& metric_type) const override { + // The CPU copy is freed after the GPU upload; GPU_HNSW does not expose + // CPU-reconstructable raw data (kept consistent with StaticHasRawData). + if (gpu_ready_.load(std::memory_order_acquire)) { + return false; + } + return BaseFaissRegularIndexHNSWNode::HasRawData(metric_type); + } + + expected + GetVectorByIds(const DataSetPtr dataset, milvus::OpContext* op_context) const override { + if (gpu_ready_.load(std::memory_order_acquire)) { + return expected::Err(Status::not_implemented, + "GPU_HNSW keeps vectors on GPU; raw data retrieval is not supported"); + } + return BaseFaissRegularIndexHNSWNode::GetVectorByIds(dataset, op_context); + } + protected: Status TrainInternal(const DataSetPtr /*dataset*/, const Config& /*cfg*/) override { @@ -3403,8 +3448,15 @@ class GpuHnswIndexNode : public BaseFaissRegularIndexHNSWNode { faiss_idx->metric_type); } gpu_index_->copyFromWithMetric(faiss_idx, use_ip, is_cosine); + // Capture count/dim before releasing the CPU copy so Count()/ + // Dim() keep working once indexes[0] is freed. + gpu_ntotal_ = faiss_idx->ntotal; + gpu_dim_ = faiss_idx->d; // Release CPU copy — vectors and graph are now on GPU. indexes[0].reset(); + // Publish gpu_index_ with release semantics; the lock-free + // reader in Search() pairs this with an acquire load. + gpu_ready_.store(true, std::memory_order_release); } catch (const std::exception& e) { fprintf(stderr, "[gpu_hnsw] eager GPU upload failed: %s\n", e.what()); gpu_index_.reset(); @@ -3416,14 +3468,19 @@ class GpuHnswIndexNode : public BaseFaissRegularIndexHNSWNode { expected Search(const DataSetPtr dataset, std::unique_ptr cfg, const BitsetView& bitset, milvus::OpContext* op_context) const override { - if (!bitset.empty() && bitset.count() > 0) { + // GPU_HNSW does not apply the bitset, so reject any search that carries + // one. Guard on data()!=nullptr rather than count(), because count() + // defaults to 0 when the caller omits num_filtered_out_bits and would + // let a real filter slip through. + if (!bitset.empty() && bitset.data() != nullptr) { return expected::Err(Status::invalid_args, "GPU_HNSW does not support filtered search"); } - // Fast path: gpu_index_ is set during Deserialize and never cleared. - if (!gpu_index_) { + // Fast path: gpu_ready_ is published (release) once gpu_index_ is fully + // constructed; the acquire load here avoids a data race on the pointer. + if (!gpu_ready_.load(std::memory_order_acquire)) { std::unique_lock lock(gpu_mutex_); - if (!gpu_index_) { + if (!gpu_ready_.load(std::memory_order_relaxed)) { const auto* faiss_idx = GetFaissHnswIndex(); if (!faiss_idx) { return expected::Err(Status::empty_index, "index not loaded"); @@ -3440,7 +3497,10 @@ class GpuHnswIndexNode : public BaseFaissRegularIndexHNSWNode { faiss_idx->metric_type); } gpu_index_->copyFromWithMetric(faiss_idx, use_ip, is_cosine); + gpu_ntotal_ = faiss_idx->ntotal; + gpu_dim_ = faiss_idx->d; const_cast&>(indexes[0]).reset(); + gpu_ready_.store(true, std::memory_order_release); } catch (const std::exception& e) { return expected::Err(Status::cuvs_inner_error, std::string("failed to build GPU HNSW index: ") + e.what()); @@ -3477,6 +3537,10 @@ class GpuHnswIndexNode : public BaseFaissRegularIndexHNSWNode { try { faiss::gpu::GpuHnswSearchParams gsp; gsp.ef = ef; + // Keep the overflow candidate queue enabled (matches + // SearchParametersGpuHNSW's default); leaving it 0 would disable + // the OCQ and reduce recall. + gsp.overflow_factor = 2; gpu_index_->searchHost(nq, h_queries, k, h_dist.get(), h_ids.get(), gsp); } catch (const std::exception& e) { LOG_KNOWHERE_ERROR_ << "GPU_HNSW search failed: " << e.what(); @@ -3508,6 +3572,12 @@ class GpuHnswIndexNode : public BaseFaissRegularIndexHNSWNode { mutable std::mutex gpu_mutex_; mutable std::shared_ptr gpu_resources_; mutable std::unique_ptr gpu_index_; + // Published (release) after gpu_index_ is fully built; read lock-free + // (acquire) on the search fast path and by Count()/Dim()/HasRawData(). + mutable std::atomic gpu_ready_{false}; + // Vector count/dim captured before the CPU copy is freed. + int64_t gpu_ntotal_ = 0; + int64_t gpu_dim_ = 0; }; // Register GPU_HNSW in the static config map at process startup. diff --git a/tests/ut/test_gpu_search.cc b/tests/ut/test_gpu_search.cc index 136131912..9a7a8d0a3 100644 --- a/tests/ut/test_gpu_search.cc +++ b/tests/ut/test_gpu_search.cc @@ -474,6 +474,11 @@ TEST_CASE("Test All GPU Index", "[search]") { auto deser_res = gpu_idx.Deserialize(bs); REQUIRE(deser_res == knowhere::Status::success); + // Count()/Dim() must survive the CPU-copy release after GPU upload; + // regression guard for returning -1 (segment treated as empty). + REQUIRE(gpu_idx.Count() == nb); + REQUIRE(gpu_idx.Dim() == dim); + // Search on GPU auto results = gpu_idx.Search(query_ds, hnsw_json, nullptr); REQUIRE(results.has_value()); From a224f003b15ca56ed8afe0676ac8878b478f7674 Mon Sep 17 00:00:00 2001 From: premal Date: Sat, 11 Jul 2026 20:49:50 +0000 Subject: [PATCH 04/24] fix: declare gpu_ntotal_/gpu_dim_ mutable (written from const Search()) Signed-off-by: premal Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- src/index/hnsw/faiss_hnsw.cc | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/src/index/hnsw/faiss_hnsw.cc b/src/index/hnsw/faiss_hnsw.cc index f4d69a02f..1f8c17417 100644 --- a/src/index/hnsw/faiss_hnsw.cc +++ b/src/index/hnsw/faiss_hnsw.cc @@ -3575,9 +3575,10 @@ class GpuHnswIndexNode : public BaseFaissRegularIndexHNSWNode { // Published (release) after gpu_index_ is fully built; read lock-free // (acquire) on the search fast path and by Count()/Dim()/HasRawData(). mutable std::atomic gpu_ready_{false}; - // Vector count/dim captured before the CPU copy is freed. - int64_t gpu_ntotal_ = 0; - int64_t gpu_dim_ = 0; + // Vector count/dim captured before the CPU copy is freed (written from the + // const lazy-upload path in Search(), hence mutable). + mutable int64_t gpu_ntotal_ = 0; + mutable int64_t gpu_dim_ = 0; }; // Register GPU_HNSW in the static config map at process startup. From 92e66b86c41e2855f4b2eb2bf014ac5b328de9c8 Mon Sep 17 00:00:00 2001 From: premal Date: Sat, 11 Jul 2026 21:07:01 +0000 Subject: [PATCH 05/24] refactor: remove non-functional OCQ overflow queue from GPU HNSW Mirror of 6si/faiss#1: the overflow candidate queue was never wired in (overflow_insert had no call site; d_overflow_count only zeroed+read), so it only wasted scratch VRAM + a memset per search. Remove the overflow machinery from the vendored faiss GPU kernel and drop the explicit gsp.overflow_factor=2 in faiss_hnsw.cc. Recall is tuned via ef. Signed-off-by: premal Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- src/index/hnsw/faiss_hnsw.cc | 4 -- thirdparty/faiss/faiss/gpu/GpuIndexHNSW.cu | 7 +- thirdparty/faiss/faiss/gpu/GpuIndexHNSW.h | 7 +- .../faiss/faiss/gpu/impl/GpuHnswSearch.cuh | 15 +---- .../faiss/gpu/impl/GpuHnswSearchKernel.cuh | 66 +------------------ .../faiss/faiss/gpu/impl/GpuHnswTypes.cu | 38 +---------- .../faiss/faiss/gpu/impl/GpuHnswTypes.h | 11 +--- 7 files changed, 9 insertions(+), 139 deletions(-) diff --git a/src/index/hnsw/faiss_hnsw.cc b/src/index/hnsw/faiss_hnsw.cc index 1f8c17417..e8713a651 100644 --- a/src/index/hnsw/faiss_hnsw.cc +++ b/src/index/hnsw/faiss_hnsw.cc @@ -3537,10 +3537,6 @@ class GpuHnswIndexNode : public BaseFaissRegularIndexHNSWNode { try { faiss::gpu::GpuHnswSearchParams gsp; gsp.ef = ef; - // Keep the overflow candidate queue enabled (matches - // SearchParametersGpuHNSW's default); leaving it 0 would disable - // the OCQ and reduce recall. - gsp.overflow_factor = 2; gpu_index_->searchHost(nq, h_queries, k, h_dist.get(), h_ids.get(), gsp); } catch (const std::exception& e) { LOG_KNOWHERE_ERROR_ << "GPU_HNSW search failed: " << e.what(); diff --git a/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.cu b/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.cu index c389c10b6..0df85cab5 100644 --- a/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.cu +++ b/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.cu @@ -162,7 +162,6 @@ void GpuIndexHNSW::searchImpl_( sp.search_width = params->search_width; sp.max_iterations = params->max_iterations; sp.thread_block_size = params->thread_block_size; - sp.overflow_factor = params->overflow_factor; got_params = true; } } @@ -174,8 +173,7 @@ void GpuIndexHNSW::searchImpl_( int nq = static_cast(n); int dim = static_cast(idx.dim); - int overflow_ef = sp.overflow_factor * sp.ef; - sc.ensure(nq, k, dim, static_cast(idx.n_rows), overflow_ef); + sc.ensure(nq, k, dim, static_cast(idx.n_rows)); // D2D: query vectors (GpuIndex::search passes device pointers) GPU_HNSW_CUDA_CHECK(cudaMemcpyAsync( @@ -242,8 +240,7 @@ void GpuIndexHNSW::searchHost( int nq = static_cast(n); int dim = static_cast(idx.dim); - int overflow_ef = sp.overflow_factor * sp.ef; - sc.ensure(nq, k, dim, static_cast(idx.n_rows), overflow_ef); + sc.ensure(nq, k, dim, static_cast(idx.n_rows)); GPU_HNSW_CUDA_CHECK(cudaMemcpyAsync( sc.d_queries, diff --git a/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.h b/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.h index 4437d8b1d..67445c844 100644 --- a/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.h +++ b/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.h @@ -54,17 +54,14 @@ struct SearchParametersGpuHNSW : SearchParameters { /// Thread block size (0 = auto, default 128). int thread_block_size = 0; - - /// Overflow queue factor: overflow_ef = overflow_factor * ef. - int overflow_factor = 2; }; /// GPU implementation of HNSW search. /// /// This index type does NOT build an HNSW graph on GPU — it takes a /// CPU-built faiss::IndexHNSW (Flat or SQ storage), converts the graph -/// to a GPU-friendly dense format, and runs the search on GPU using an -/// Overflow Candidate Queue (OCQ) beam search kernel. +/// to a GPU-friendly dense format, and runs the search on GPU using a +/// parallel beam search kernel. /// /// Supports L2, inner product, and cosine metrics. /// Supports float32 and int8 (QT_8bit_direct_signed) data. diff --git a/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearch.cuh b/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearch.cuh index 336e37a6f..04d1636d8 100644 --- a/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearch.cuh +++ b/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearch.cuh @@ -58,7 +58,6 @@ inline void gpu_hnsw_search( int ef = params.ef; int sw = params.search_width; - int overflow_ef = params.overflow_factor * ef; int max_iter = params.max_iterations > 0 ? params.max_iterations : 2 * ef / sw + 10; @@ -150,13 +149,6 @@ inline void gpu_hnsw_search( GPU_HNSW_CUDA_CHECK( cudaMemsetAsync(sc.d_visited_bitmaps, 0, bitmap_bytes, stream)); - if (overflow_ef > 0) { - GPU_HNSW_CUDA_CHECK(cudaMemsetAsync( - sc.d_overflow_count, - 0, - static_cast(num_queries) * sizeof(int), - stream)); - } hnsw_kernel::layer0_beam_search_kernel <<>>( @@ -176,12 +168,7 @@ inline void gpu_hnsw_search( ef, sw, max_iter, - idx.use_ip, - overflow_ef, - sc.d_overflow_ids, - sc.d_overflow_dists, - sc.d_overflow_expanded, - sc.d_overflow_count); + idx.use_ip); GPU_HNSW_CUDA_CHECK(cudaGetLastError()); }; diff --git a/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearchKernel.cuh b/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearchKernel.cuh index 84c5ba659..86342f92f 100644 --- a/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearchKernel.cuh +++ b/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearchKernel.cuh @@ -202,7 +202,7 @@ __global__ void upper_layer_search_kernel( } // ============================================================================ -// Phase 2: Layer-0 beam search kernel with Overflow Candidate Queue (OCQ) +// Phase 2: Layer-0 parallel beam search kernel // ============================================================================ __device__ __forceinline__ bool bitmap_visit( @@ -214,40 +214,6 @@ __device__ __forceinline__ bool bitmap_visit( return (old & bit) == 0; } -__device__ __forceinline__ void overflow_insert( - uint32_t* ovf_ids, - float* ovf_dists, - uint32_t* ovf_exp, - int& ovf_rc, - int overflow_ef, - uint32_t id, - float dist, - uint32_t expanded) { - if (ovf_rc >= overflow_ef && dist >= ovf_dists[ovf_rc - 1]) - return; - - int lo = 0, hi = ovf_rc; - while (lo < hi) { - int mid = (lo + hi) / 2; - if (ovf_dists[mid] < dist) - lo = mid + 1; - else - hi = mid; - } - - int insert_end = ovf_rc < overflow_ef ? ovf_rc : overflow_ef - 1; - for (int i = insert_end; i > lo; i--) { - ovf_ids[i] = ovf_ids[i - 1]; - ovf_dists[i] = ovf_dists[i - 1]; - ovf_exp[i] = ovf_exp[i - 1]; - } - ovf_ids[lo] = id; - ovf_dists[lo] = dist; - ovf_exp[lo] = expanded; - if (ovf_rc < overflow_ef) - ovf_rc++; -} - template __global__ void layer0_beam_search_kernel( const float* __restrict__ d_queries, @@ -266,12 +232,7 @@ __global__ void layer0_beam_search_kernel( int ef, int search_width, int max_iterations, - bool use_inner_product, - int overflow_ef, - uint32_t* __restrict__ d_overflow_ids, - float* __restrict__ d_overflow_dists, - uint32_t* __restrict__ d_overflow_expanded, - int* __restrict__ d_overflow_count) { + bool use_inner_product) { int query_idx = blockIdx.x; if (query_idx >= num_queries) return; @@ -295,17 +256,6 @@ __global__ void layer0_beam_search_kernel( uint32_t* visited_bmap = d_visited_bitmaps + static_cast(query_idx) * bitmap_words; - uint32_t* ovf_ids = overflow_ef > 0 - ? d_overflow_ids + static_cast(query_idx) * overflow_ef - : nullptr; - float* ovf_dists = overflow_ef > 0 - ? d_overflow_dists + static_cast(query_idx) * overflow_ef - : nullptr; - uint32_t* ovf_exp = overflow_ef > 0 - ? d_overflow_expanded + - static_cast(query_idx) * overflow_ef - : nullptr; - for (int i = threadIdx.x; i < ef; i += blockDim.x) { result_ids[i] = UINT32_MAX; result_dists[i] = FLT_MAX; @@ -315,8 +265,6 @@ __global__ void layer0_beam_search_kernel( meta[0] = 0; meta[1] = 0; meta[2] = 0; - if (overflow_ef > 0) - d_overflow_count[query_idx] = 0; } __syncthreads(); @@ -426,16 +374,6 @@ __global__ void layer0_beam_search_kernel( } } - if (num_parents == 0 && overflow_ef > 0) { - int ovf_rc = d_overflow_count[query_idx]; - for (int i = 0; i < ovf_rc && num_parents < search_width; - i++) { - if (!ovf_exp[i]) { - parent_ids[num_parents++] = ovf_ids[i]; - ovf_exp[i] = 1; - } - } - } meta[2] = num_parents; } __syncthreads(); diff --git a/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.cu b/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.cu index 1cd98bbdb..131a391c4 100644 --- a/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.cu +++ b/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.cu @@ -48,8 +48,7 @@ void GpuHnswSearchScratch::ensure( int nq, int k, int dim, - int N, - int overflow_ef) { + int N) { size_t need_q = static_cast(nq) * dim * sizeof(float); if (need_q > queries_bytes) { if (d_queries) @@ -88,33 +87,6 @@ void GpuHnswSearchScratch::ensure( SCRATCH_CUDA_CHECK(cudaMalloc(&d_visited_bitmaps, need_bm)); bitmap_bytes = need_bm; } - if (overflow_ef > 0) { - size_t ovf_entries = static_cast(nq) * overflow_ef; - size_t need_ovf = - ovf_entries * - (sizeof(uint32_t) + sizeof(float) + sizeof(uint32_t)) + - static_cast(nq) * sizeof(int); - if (need_ovf > overflow_bytes) { - if (d_overflow_ids) - cudaFree(d_overflow_ids); - if (d_overflow_dists) - cudaFree(d_overflow_dists); - if (d_overflow_expanded) - cudaFree(d_overflow_expanded); - if (d_overflow_count) - cudaFree(d_overflow_count); - SCRATCH_CUDA_CHECK(cudaMalloc( - &d_overflow_ids, ovf_entries * sizeof(uint32_t))); - SCRATCH_CUDA_CHECK(cudaMalloc( - &d_overflow_dists, ovf_entries * sizeof(float))); - SCRATCH_CUDA_CHECK(cudaMalloc( - &d_overflow_expanded, ovf_entries * sizeof(uint32_t))); - SCRATCH_CUDA_CHECK(cudaMalloc( - &d_overflow_count, - static_cast(nq) * sizeof(int))); - overflow_bytes = need_ovf; - } - } } GpuHnswSearchScratch::~GpuHnswSearchScratch() { @@ -128,14 +100,6 @@ GpuHnswSearchScratch::~GpuHnswSearchScratch() { cudaFree(d_entry_points); if (d_visited_bitmaps) cudaFree(d_visited_bitmaps); - if (d_overflow_ids) - cudaFree(d_overflow_ids); - if (d_overflow_dists) - cudaFree(d_overflow_dists); - if (d_overflow_expanded) - cudaFree(d_overflow_expanded); - if (d_overflow_count) - cudaFree(d_overflow_count); } GpuHnswScratchSlot::~GpuHnswScratchSlot() { diff --git a/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.h b/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.h index 9f2db4152..0093b3e11 100644 --- a/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.h +++ b/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.h @@ -39,9 +39,6 @@ struct GpuHnswSearchParams { int search_width = 4; int max_iterations = 0; int thread_block_size = 0; - // Keep in sync with SearchParametersGpuHNSW::overflow_factor: a zero here - // would disable the overflow candidate queue and silently drop recall. - int overflow_factor = 2; }; struct GpuHnswDeviceUpperLayer { @@ -58,19 +55,13 @@ struct GpuHnswSearchScratch { uint32_t* d_entry_points = nullptr; uint32_t* d_visited_bitmaps = nullptr; - uint32_t* d_overflow_ids = nullptr; - float* d_overflow_dists = nullptr; - uint32_t* d_overflow_expanded = nullptr; - int* d_overflow_count = nullptr; - size_t queries_bytes = 0; size_t neighbors_bytes = 0; size_t distances_bytes = 0; int entry_cap = 0; size_t bitmap_bytes = 0; - size_t overflow_bytes = 0; - void ensure(int nq, int k, int dim, int N, int overflow_ef); + void ensure(int nq, int k, int dim, int N); ~GpuHnswSearchScratch(); From 03a613f66b76bc7e8f3dc158781d870fbf447a9d Mon Sep 17 00:00:00 2001 From: premal Date: Sat, 11 Jul 2026 21:24:21 +0000 Subject: [PATCH 06/24] feat: register GPU_HNSW_SQ as an alias of GPU_HNSW GPU_HNSW_SQ was exposed by Milvus (constant, param spec, tests, GPU routing) but had no Knowhere backend, so creating it hit no impl. Add a thin GpuHnswSQIndexNode subclass of GpuHnswIndexNode that only overrides Type(); the GPU search path is identical (a CPU-built HNSW/HNSW_SQ index uploaded to GpuIndexHNSW). Register it for VECTOR_FLOAT and VECTOR_INT8, add the INDEX_GPU_HNSW_SQ constant + index-table entries, and accept the GPU_HNSW_SQ binset key in Deserialize's alias path. Co-Authored-By: Claude Opus 4.6 Signed-off-by: premal Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- include/knowhere/comp/index_param.h | 1 + include/knowhere/index/index_table.h | 2 ++ src/index/hnsw/faiss_hnsw.cc | 33 ++++++++++++++++++++++++++-- 3 files changed, 34 insertions(+), 2 deletions(-) diff --git a/include/knowhere/comp/index_param.h b/include/knowhere/comp/index_param.h index ab26efe52..93cb96d01 100644 --- a/include/knowhere/comp/index_param.h +++ b/include/knowhere/comp/index_param.h @@ -54,6 +54,7 @@ constexpr const char* INDEX_GPU_IVFFLAT = "GPU_IVF_FLAT"; constexpr const char* INDEX_GPU_IVFPQ = "GPU_IVF_PQ"; constexpr const char* INDEX_GPU_CAGRA = "GPU_CAGRA"; constexpr const char* INDEX_GPU_HNSW = "GPU_HNSW"; +constexpr const char* INDEX_GPU_HNSW_SQ = "GPU_HNSW_SQ"; constexpr const char* INDEX_HNSW = "HNSW"; constexpr const char* INDEX_HNSW_SQ = "HNSW_SQ"; diff --git a/include/knowhere/index/index_table.h b/include/knowhere/index/index_table.h index 88e12ce73..7d3fb628d 100644 --- a/include/knowhere/index/index_table.h +++ b/include/knowhere/index/index_table.h @@ -84,6 +84,8 @@ static std::set> legal_knowhere_index = { {IndexEnum::INDEX_GPU_CAGRA, VecType::VECTOR_BINARY}, {IndexEnum::INDEX_GPU_HNSW, VecType::VECTOR_FLOAT}, {IndexEnum::INDEX_GPU_HNSW, VecType::VECTOR_INT8}, + {IndexEnum::INDEX_GPU_HNSW_SQ, VecType::VECTOR_FLOAT}, + {IndexEnum::INDEX_GPU_HNSW_SQ, VecType::VECTOR_INT8}, // hnsw {IndexEnum::INDEX_HNSW, VecType::VECTOR_FLOAT}, diff --git a/src/index/hnsw/faiss_hnsw.cc b/src/index/hnsw/faiss_hnsw.cc index e8713a651..a215dcb87 100644 --- a/src/index/hnsw/faiss_hnsw.cc +++ b/src/index/hnsw/faiss_hnsw.cc @@ -3416,7 +3416,7 @@ class GpuHnswIndexNode : public BaseFaissRegularIndexHNSWNode { Status status; if (!binset.Contains(IndexEnum::INDEX_GPU_HNSW)) { BinarySet aliased = binset; - for (const char* key : {IndexEnum::INDEX_HNSW_SQ, IndexEnum::INDEX_HNSW}) { + for (const char* key : {IndexEnum::INDEX_GPU_HNSW_SQ, IndexEnum::INDEX_HNSW_SQ, IndexEnum::INDEX_HNSW}) { if (binset.Contains(key)) { aliased.Append(IndexEnum::INDEX_GPU_HNSW, binset.GetByName(key)); break; @@ -3577,11 +3577,28 @@ class GpuHnswIndexNode : public BaseFaissRegularIndexHNSWNode { mutable int64_t gpu_dim_ = 0; }; -// Register GPU_HNSW in the static config map at process startup. +// GPU_HNSW_SQ is an alias of GPU_HNSW: the GPU search path is identical (a +// CPU-built HNSW/HNSW_SQ index uploaded to GpuIndexHNSW), so it reuses +// GpuHnswIndexNode and only reports a distinct Type() so Milvus's configured +// index_type matches the loaded node. +class GpuHnswSQIndexNode : public GpuHnswIndexNode { + public: + GpuHnswSQIndexNode(const int32_t& version, const Object& object) : GpuHnswIndexNode(version, object) { + } + + std::string + Type() const override { + return IndexEnum::INDEX_GPU_HNSW_SQ; + } +}; + +// Register GPU_HNSW / GPU_HNSW_SQ in the static config map at process startup. __attribute__((constructor)) static void register_gpu_hnsw_static_config() { IndexStaticFaced::Instance().RegisterStaticFunc(IndexEnum::INDEX_GPU_HNSW); IndexStaticFaced::Instance().RegisterStaticFunc(IndexEnum::INDEX_GPU_HNSW); + IndexStaticFaced::Instance().RegisterStaticFunc(IndexEnum::INDEX_GPU_HNSW_SQ); + IndexStaticFaced::Instance().RegisterStaticFunc(IndexEnum::INDEX_GPU_HNSW_SQ); } KNOWHERE_REGISTER_GLOBAL( @@ -3595,6 +3612,18 @@ KNOWHERE_REGISTER_GLOBAL( return Index>::Create(std::make_unique(version, object)); }, int8, true, (feature::INT8 | feature::GPU)); + +KNOWHERE_REGISTER_GLOBAL( + GPU_HNSW_SQ, + [](const int32_t& version, const Object& object) { return Index::Create(version, object); }, + fp32, true, feature::GPU_ANN_FLOAT_INDEX); + +KNOWHERE_REGISTER_GLOBAL( + GPU_HNSW_SQ, + [](const int32_t& version, const Object& object) { + return Index>::Create(std::make_unique(version, object)); + }, + int8, true, (feature::INT8 | feature::GPU)); #endif // KNOWHERE_WITH_CUVS } // namespace knowhere From 18768e9b11683250b07ffb21e7763c39e18110cd Mon Sep 17 00:00:00 2001 From: premal Date: Sat, 11 Jul 2026 21:57:24 +0000 Subject: [PATCH 07/24] fix: mirror GPU HNSW re-review fixes into vendored faiss Sync of 6si/faiss@467593eae (byte-identical vendored copy): - searchImpl_ event-based stream sync before reading host queries (#15) - cudaSetDevice() before cudaFree in scratch/index destructors (#19/#20) - clear error when search_width exceeds shared-memory budget (#18) Co-Authored-By: Claude Opus 4.6 Signed-off-by: premal Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- thirdparty/faiss/faiss/gpu/GpuIndexHNSW.cu | 13 +++++++++++++ thirdparty/faiss/faiss/gpu/impl/GpuHnswBuild.cuh | 2 ++ thirdparty/faiss/faiss/gpu/impl/GpuHnswSearch.cuh | 12 ++++++++++++ thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.cu | 7 +++++++ thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.h | 8 ++++++++ 5 files changed, 42 insertions(+) diff --git a/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.cu b/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.cu index 0df85cab5..8c099e120 100644 --- a/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.cu +++ b/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.cu @@ -175,6 +175,19 @@ void GpuIndexHNSW::searchImpl_( int dim = static_cast(idx.dim); sc.ensure(nq, k, dim, static_cast(idx.n_rows)); + // The parent GpuIndex::search copied host queries into `x` (and will later + // copy outputs back) on the resources' default stream. Our scratch-pool + // slot has its own private stream, so make it wait for that H2D/D2D copy to + // complete before we read `x` — otherwise the search can race ahead and + // operate on partially-written query vectors. + cudaStream_t defaultStream = resources_->getDefaultStream(config_.device); + cudaEvent_t inputReady; + GPU_HNSW_CUDA_CHECK( + cudaEventCreateWithFlags(&inputReady, cudaEventDisableTiming)); + GPU_HNSW_CUDA_CHECK(cudaEventRecord(inputReady, defaultStream)); + GPU_HNSW_CUDA_CHECK(cudaStreamWaitEvent(stream, inputReady, 0)); + GPU_HNSW_CUDA_CHECK(cudaEventDestroy(inputReady)); + // D2D: query vectors (GpuIndex::search passes device pointers) GPU_HNSW_CUDA_CHECK(cudaMemcpyAsync( sc.d_queries, diff --git a/thirdparty/faiss/faiss/gpu/impl/GpuHnswBuild.cuh b/thirdparty/faiss/faiss/gpu/impl/GpuHnswBuild.cuh index 5f4ca5613..62ee675f0 100644 --- a/thirdparty/faiss/faiss/gpu/impl/GpuHnswBuild.cuh +++ b/thirdparty/faiss/faiss/gpu/impl/GpuHnswBuild.cuh @@ -290,6 +290,7 @@ inline std::unique_ptr from_faiss_hnsw_sq( idx->n_rows = n_rows; idx->dim = dim; idx->use_ip = use_ip; + idx->device = device; idx->scratch_pool = std::make_unique(4, device); bool is_direct_signed = (sq_storage->sq.qtype == @@ -332,6 +333,7 @@ inline std::unique_ptr from_faiss_hnsw_flat( idx->n_rows = n_rows; idx->dim = dim; idx->use_ip = use_ip; + idx->device = device; idx->scratch_pool = std::make_unique(4, device); upload_fp32_dataset(*idx, h_vectors, n_rows, is_cosine); diff --git a/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearch.cuh b/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearch.cuh index 04d1636d8..ed5d000b9 100644 --- a/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearch.cuh +++ b/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearch.cuh @@ -126,6 +126,18 @@ inline void gpu_hnsw_search( { int smem_overhead = sw * idx.max_degree0 * 8 + sw * 4 + 12; int max_ef = (smem_max - smem_overhead) / 12; + // A search_width so large that even ef=1 does not fit in shared + // memory leaves no valid beam. Fail clearly rather than clamping ef + // to a negative/zero value (which would launch with a bad geometry + // or read out of bounds). + if (max_ef < 1) { + throw std::runtime_error( + std::string("gpu_hnsw: search_width=") + + std::to_string(sw) + + " too large for device shared memory (" + + std::to_string(smem_max) + + " bytes); reduce search_width"); + } if (ef > max_ef) { ef = max_ef; } diff --git a/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.cu b/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.cu index 131a391c4..4895aa94a 100644 --- a/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.cu +++ b/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.cu @@ -90,6 +90,9 @@ void GpuHnswSearchScratch::ensure( } GpuHnswSearchScratch::~GpuHnswSearchScratch() { + // Set the owning device before freeing so cudaFree runs in the correct + // context on multi-GPU systems. + cudaSetDevice(device); if (d_queries) cudaFree(d_queries); if (d_neighbors) @@ -117,6 +120,7 @@ void GpuHnswScratchPool::init_once() { available_.reserve(pool_size_); for (int i = 0; i < pool_size_; i++) { auto slot = std::make_unique(); + slot->scratch.device = device_; SCRATCH_CUDA_CHECK(cudaSetDevice(device_)); SCRATCH_CUDA_CHECK( cudaStreamCreateWithFlags(&slot->stream, cudaStreamNonBlocking)); @@ -144,6 +148,9 @@ void GpuHnswScratchPool::release(GpuHnswScratchSlot* slot) { } GpuHnswDeviceIndex::~GpuHnswDeviceIndex() { + // Set the owning device before freeing so cudaFree runs in the correct + // context on multi-GPU systems. + cudaSetDevice(device); if (d_dataset) cudaFree(d_dataset); if (d_inv_norms) diff --git a/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.h b/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.h index 0093b3e11..6c69973e2 100644 --- a/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.h +++ b/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.h @@ -61,6 +61,10 @@ struct GpuHnswSearchScratch { int entry_cap = 0; size_t bitmap_bytes = 0; + // Device this scratch's allocations live on; used to set the CUDA device + // context before freeing in the destructor (multi-GPU correctness). + int device = 0; + void ensure(int nq, int k, int dim, int N); ~GpuHnswSearchScratch(); @@ -146,6 +150,10 @@ struct GpuHnswDeviceIndex { void* d_upper_layer_ptrs = nullptr; int num_upper_layers_built = 0; + // Device this index's allocations live on; used to set the CUDA device + // context before freeing in the destructor (multi-GPU correctness). + int device = 0; + mutable std::unique_ptr scratch_pool; ~GpuHnswDeviceIndex(); From 0165f9d88ed933178e963fc6b9c15ab019fc6b3f Mon Sep 17 00:00:00 2001 From: premal Date: Sat, 11 Jul 2026 22:40:53 +0000 Subject: [PATCH 08/24] feat: register native FP16/BF16 for GPU_HNSW and GPU_HNSW_SQ Mirror the faiss native fp16/bf16 device-storage change into the vendored faiss copy (byte-identical to 6si/faiss) and register fp16/bf16 for both GPU_HNSW and GPU_HNSW_SQ: - index_table.h: add VECTOR_FLOAT16 / VECTOR_BFLOAT16 legal-type rows for both GPU index types. - faiss_hnsw.cc: add fp16/bf16 KNOWHERE_REGISTER_GLOBAL (feature::FP16|GPU, feature::BF16|GPU) and RegisterStaticFunc entries for both nodes; update the class comment to describe native 2-byte fp16/bf16 device storage with per-element up-conversion to fp32 in the kernel. - test_gpu_search.cc: add GPU HNSW FP16/BF16 deserialize+search UTs (CPU HNSW_SQ QT_fp16/QT_bf16 build -> GPU load -> recall vs brute force). Compile-verified only; GPU-runtime validation pending capacity. Signed-off-by: premal Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- include/knowhere/index/index_table.h | 4 + src/index/hnsw/faiss_hnsw.cc | 43 +++++++++- tests/ut/test_gpu_search.cc | 80 +++++++++++++++++++ .../faiss/faiss/gpu/impl/GpuHnswBuild.cuh | 72 +++++++++++++++-- .../faiss/faiss/gpu/impl/GpuHnswSearch.cuh | 27 +++++-- .../faiss/gpu/impl/GpuHnswSearchKernel.cuh | 4 + .../faiss/faiss/gpu/impl/GpuHnswTypes.h | 13 ++- 7 files changed, 227 insertions(+), 16 deletions(-) diff --git a/include/knowhere/index/index_table.h b/include/knowhere/index/index_table.h index 7d3fb628d..0b300ef66 100644 --- a/include/knowhere/index/index_table.h +++ b/include/knowhere/index/index_table.h @@ -83,8 +83,12 @@ static std::set> legal_knowhere_index = { {IndexEnum::INDEX_GPU_CAGRA, VecType::VECTOR_INT8}, {IndexEnum::INDEX_GPU_CAGRA, VecType::VECTOR_BINARY}, {IndexEnum::INDEX_GPU_HNSW, VecType::VECTOR_FLOAT}, + {IndexEnum::INDEX_GPU_HNSW, VecType::VECTOR_FLOAT16}, + {IndexEnum::INDEX_GPU_HNSW, VecType::VECTOR_BFLOAT16}, {IndexEnum::INDEX_GPU_HNSW, VecType::VECTOR_INT8}, {IndexEnum::INDEX_GPU_HNSW_SQ, VecType::VECTOR_FLOAT}, + {IndexEnum::INDEX_GPU_HNSW_SQ, VecType::VECTOR_FLOAT16}, + {IndexEnum::INDEX_GPU_HNSW_SQ, VecType::VECTOR_BFLOAT16}, {IndexEnum::INDEX_GPU_HNSW_SQ, VecType::VECTOR_INT8}, // hnsw diff --git a/src/index/hnsw/faiss_hnsw.cc b/src/index/hnsw/faiss_hnsw.cc index a215dcb87..68a4bf2c5 100644 --- a/src/index/hnsw/faiss_hnsw.cc +++ b/src/index/hnsw/faiss_hnsw.cc @@ -3319,10 +3319,13 @@ GetGpuConstructionMutex() { return mtx; } -// GPU HNSW index node. Registered for F32 and INT8 (see registrations at the -// bottom of this file); it loads CPU-serialized HNSW (F32) or HNSW_SQ (SQ8/INT8) -// binaries at load time and uploads them to the GPU via faiss::gpu::GpuIndexHNSW. -// The CPU copy is freed after upload, so this node exposes no CPU raw data. +// GPU HNSW index node. Registered for F32, FP16, BF16 and INT8 (see +// registrations at the bottom of this file); it loads CPU-serialized HNSW (F32 +// Flat) or HNSW_SQ (SQ8/INT8, fp16, bf16) binaries at load time and uploads them +// to the GPU via faiss::gpu::GpuIndexHNSW. FP16/BF16/INT8 stay in their native +// low-precision layout on the device (2 bytes for fp16/bf16, 1 byte for int8) +// and are up-converted to fp32 per element inside the search kernel. The CPU +// copy is freed after upload, so this node exposes no CPU raw data. class GpuHnswIndexNode : public BaseFaissRegularIndexHNSWNode { public: GpuHnswIndexNode(const int32_t& version, const Object& object) @@ -3596,8 +3599,12 @@ class GpuHnswSQIndexNode : public GpuHnswIndexNode { __attribute__((constructor)) static void register_gpu_hnsw_static_config() { IndexStaticFaced::Instance().RegisterStaticFunc(IndexEnum::INDEX_GPU_HNSW); + IndexStaticFaced::Instance().RegisterStaticFunc(IndexEnum::INDEX_GPU_HNSW); + IndexStaticFaced::Instance().RegisterStaticFunc(IndexEnum::INDEX_GPU_HNSW); IndexStaticFaced::Instance().RegisterStaticFunc(IndexEnum::INDEX_GPU_HNSW); IndexStaticFaced::Instance().RegisterStaticFunc(IndexEnum::INDEX_GPU_HNSW_SQ); + IndexStaticFaced::Instance().RegisterStaticFunc(IndexEnum::INDEX_GPU_HNSW_SQ); + IndexStaticFaced::Instance().RegisterStaticFunc(IndexEnum::INDEX_GPU_HNSW_SQ); IndexStaticFaced::Instance().RegisterStaticFunc(IndexEnum::INDEX_GPU_HNSW_SQ); } @@ -3606,6 +3613,20 @@ KNOWHERE_REGISTER_GLOBAL( [](const int32_t& version, const Object& object) { return Index::Create(version, object); }, fp32, true, feature::GPU_ANN_FLOAT_INDEX); +KNOWHERE_REGISTER_GLOBAL( + GPU_HNSW, + [](const int32_t& version, const Object& object) { + return Index>::Create(std::make_unique(version, object)); + }, + fp16, true, (feature::FP16 | feature::GPU)); + +KNOWHERE_REGISTER_GLOBAL( + GPU_HNSW, + [](const int32_t& version, const Object& object) { + return Index>::Create(std::make_unique(version, object)); + }, + bf16, true, (feature::BF16 | feature::GPU)); + KNOWHERE_REGISTER_GLOBAL( GPU_HNSW, [](const int32_t& version, const Object& object) { @@ -3618,6 +3639,20 @@ KNOWHERE_REGISTER_GLOBAL( [](const int32_t& version, const Object& object) { return Index::Create(version, object); }, fp32, true, feature::GPU_ANN_FLOAT_INDEX); +KNOWHERE_REGISTER_GLOBAL( + GPU_HNSW_SQ, + [](const int32_t& version, const Object& object) { + return Index>::Create(std::make_unique(version, object)); + }, + fp16, true, (feature::FP16 | feature::GPU)); + +KNOWHERE_REGISTER_GLOBAL( + GPU_HNSW_SQ, + [](const int32_t& version, const Object& object) { + return Index>::Create(std::make_unique(version, object)); + }, + bf16, true, (feature::BF16 | feature::GPU)); + KNOWHERE_REGISTER_GLOBAL( GPU_HNSW_SQ, [](const int32_t& version, const Object& object) { diff --git a/tests/ut/test_gpu_search.cc b/tests/ut/test_gpu_search.cc index 9a7a8d0a3..9884c1839 100644 --- a/tests/ut/test_gpu_search.cc +++ b/tests/ut/test_gpu_search.cc @@ -643,6 +643,86 @@ TEST_CASE("Test All GPU Index", "[search]") { // SQ8 has lower recall than flat due to quantization REQUIRE(recall >= 0.7f); } + + SECTION("Test GPU HNSW FP16 Deserialization") { + // Build a CPU HNSW (fp16 -> HNSW_SQ QT_fp16 storage), then load and + // search on GPU. fp16 vectors stay in their native 2-byte layout on the + // device and are up-converted to fp32 per element in the search kernel. + knowhere::Json hnsw_json; + hnsw_json[knowhere::meta::DIM] = dim; + hnsw_json[knowhere::meta::METRIC_TYPE] = knowhere::metric::L2; + hnsw_json[knowhere::meta::TOPK] = 1; + hnsw_json[knowhere::indexparam::HNSW_M] = 16; + hnsw_json[knowhere::indexparam::EFCONSTRUCTION] = 200; + hnsw_json[knowhere::indexparam::EF] = 200; + + auto train_ds = knowhere::ConvertToDataTypeIfNeeded(GenDataSet(nb, dim, seed)); + auto query_ds = knowhere::ConvertToDataTypeIfNeeded(GenDataSet(nq, dim, seed + 2)); + + auto cpu_idx = + knowhere::IndexFactory::Instance().Create(knowhere::IndexEnum::INDEX_HNSW, version).value(); + auto res = cpu_idx.Build(train_ds, hnsw_json); + REQUIRE(res == knowhere::Status::success); + + knowhere::BinarySet bs; + cpu_idx.Serialize(bs); + + auto gpu_idx = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_GPU_HNSW, version) + .value(); + auto deser_res = gpu_idx.Deserialize(bs); + REQUIRE(deser_res == knowhere::Status::success); + + auto results = gpu_idx.Search(query_ds, hnsw_json, nullptr); + REQUIRE(results.has_value()); + + auto gt = knowhere::BruteForce::Search(train_ds, query_ds, hnsw_json, nullptr); + REQUIRE(gt.has_value()); + float recall = GetKNNRecall(*gt.value(), *results.value()); + // fp16 is near-lossless (10-bit mantissa); recall should stay close to + // the fp32 flat path. + REQUIRE(recall >= 0.9f); + } + + SECTION("Test GPU HNSW BF16 Deserialization") { + // Build a CPU HNSW (bf16 -> HNSW_SQ QT_bf16 storage), then load and + // search on GPU. bf16 vectors stay in their native 2-byte layout on the + // device and are up-converted to fp32 per element in the search kernel. + knowhere::Json hnsw_json; + hnsw_json[knowhere::meta::DIM] = dim; + hnsw_json[knowhere::meta::METRIC_TYPE] = knowhere::metric::L2; + hnsw_json[knowhere::meta::TOPK] = 1; + hnsw_json[knowhere::indexparam::HNSW_M] = 16; + hnsw_json[knowhere::indexparam::EFCONSTRUCTION] = 200; + hnsw_json[knowhere::indexparam::EF] = 200; + + auto train_ds = knowhere::ConvertToDataTypeIfNeeded(GenDataSet(nb, dim, seed)); + auto query_ds = knowhere::ConvertToDataTypeIfNeeded(GenDataSet(nq, dim, seed + 2)); + + auto cpu_idx = + knowhere::IndexFactory::Instance().Create(knowhere::IndexEnum::INDEX_HNSW, version).value(); + auto res = cpu_idx.Build(train_ds, hnsw_json); + REQUIRE(res == knowhere::Status::success); + + knowhere::BinarySet bs; + cpu_idx.Serialize(bs); + + auto gpu_idx = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_GPU_HNSW, version) + .value(); + auto deser_res = gpu_idx.Deserialize(bs); + REQUIRE(deser_res == knowhere::Status::success); + + auto results = gpu_idx.Search(query_ds, hnsw_json, nullptr); + REQUIRE(results.has_value()); + + auto gt = knowhere::BruteForce::Search(train_ds, query_ds, hnsw_json, nullptr); + REQUIRE(gt.has_value()); + float recall = GetKNNRecall(*gt.value(), *results.value()); + // bf16 has a 7-bit mantissa (coarser than fp16) but a wide exponent; + // recall stays high but with a slightly looser floor. + REQUIRE(recall >= 0.85f); + } } TEST_CASE("Test CPU vs GPU HNSW Comparison", "[gpu_hnsw_compare]") { diff --git a/thirdparty/faiss/faiss/gpu/impl/GpuHnswBuild.cuh b/thirdparty/faiss/faiss/gpu/impl/GpuHnswBuild.cuh index 62ee675f0..e86e2abd7 100644 --- a/thirdparty/faiss/faiss/gpu/impl/GpuHnswBuild.cuh +++ b/thirdparty/faiss/faiss/gpu/impl/GpuHnswBuild.cuh @@ -222,7 +222,48 @@ inline void upload_fp32_dataset( h_vectors.data(), dataset_bytes, cudaMemcpyHostToDevice)); - idx.dataset_int8 = false; + idx.dataset_type = GpuHnswDatasetType::FP32; +} + +// Upload fp16 (QT_fp16) or bf16 (QT_bf16) ScalarQuantizer codes to the GPU in +// their native 2-byte layout. faiss stores these codes row-major as raw IEEE +// half / bfloat16, which are bit-compatible with CUDA half / __nv_bfloat16, so +// the bytes are copied verbatim (no up-conversion to fp32). For cosine, inverse +// L2 norms are computed from the fp32-decoded values and applied at search time +// — mirroring the int8 path, since the stored codes are not normalized. +inline void upload_halfwidth_dataset( + GpuHnswDeviceIndex& idx, + const uint8_t* codes, + const float* decoded_for_norms, + int64_t n_rows, + bool is_cosine, + GpuHnswDatasetType dtype) { + int64_t dim = idx.dim; + size_t dataset_bytes = static_cast(n_rows) * dim * 2; + GPU_HNSW_BUILD_CUDA_CHECK(cudaMalloc(&idx.d_dataset, dataset_bytes)); + GPU_HNSW_BUILD_CUDA_CHECK(cudaMemcpy( + idx.d_dataset, codes, dataset_bytes, cudaMemcpyHostToDevice)); + idx.dataset_type = dtype; + + if (is_cosine && decoded_for_norms) { + std::vector h_inv_norms(n_rows); + for (int64_t i = 0; i < n_rows; i++) { + const float* row = decoded_for_norms + i * dim; + float sq_norm = 0.0f; + for (int64_t d = 0; d < dim; d++) { + sq_norm += row[d] * row[d]; + } + h_inv_norms[i] = + (sq_norm > 0.0f) ? (1.0f / std::sqrt(sq_norm)) : 0.0f; + } + size_t norms_bytes = static_cast(n_rows) * sizeof(float); + GPU_HNSW_BUILD_CUDA_CHECK(cudaMalloc(&idx.d_inv_norms, norms_bytes)); + GPU_HNSW_BUILD_CUDA_CHECK(cudaMemcpy( + idx.d_inv_norms, + h_inv_norms.data(), + norms_bytes, + cudaMemcpyHostToDevice)); + } } inline void upload_int8_dataset( @@ -245,7 +286,7 @@ inline void upload_int8_dataset( signed_codes.data(), dataset_bytes, cudaMemcpyHostToDevice)); - idx.dataset_int8 = true; + idx.dataset_type = GpuHnswDatasetType::INT8; if (is_cosine) { std::vector h_inv_norms(n_rows); @@ -293,11 +334,32 @@ inline std::unique_ptr from_faiss_hnsw_sq( idx->device = device; idx->scratch_pool = std::make_unique(4, device); - bool is_direct_signed = (sq_storage->sq.qtype == - faiss::ScalarQuantizer::QT_8bit_direct_signed); + auto qtype = sq_storage->sq.qtype; - if (is_direct_signed) { + if (qtype == faiss::ScalarQuantizer::QT_8bit_direct_signed) { upload_int8_dataset(*idx, sq_storage->codes.data(), n_rows, is_cosine); + } else if ( + qtype == faiss::ScalarQuantizer::QT_fp16 || + qtype == faiss::ScalarQuantizer::QT_bf16) { + // Keep fp16/bf16 in their native 2-byte layout on the GPU. Decode to + // fp32 only when cosine needs the row norms. + GpuHnswDatasetType dtype = + (qtype == faiss::ScalarQuantizer::QT_fp16) + ? GpuHnswDatasetType::FP16 + : GpuHnswDatasetType::BF16; + std::vector decoded; + if (is_cosine) { + decoded.resize(static_cast(n_rows) * dim); + sq_storage->sa_decode( + n_rows, sq_storage->codes.data(), decoded.data()); + } + upload_halfwidth_dataset( + *idx, + sq_storage->codes.data(), + is_cosine ? decoded.data() : nullptr, + n_rows, + is_cosine, + dtype); } else { std::vector h_vectors(n_rows * dim); sq_storage->sa_decode( diff --git a/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearch.cuh b/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearch.cuh index ed5d000b9..671c4d0a7 100644 --- a/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearch.cuh +++ b/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearch.cuh @@ -23,6 +23,8 @@ #pragma once +#include +#include #include #include @@ -184,12 +186,25 @@ inline void gpu_hnsw_search( GPU_HNSW_CUDA_CHECK(cudaGetLastError()); }; - if (idx.dataset_int8) { - launch_kernels( - static_cast(idx.d_dataset), idx.d_inv_norms); - } else { - launch_kernels( - static_cast(idx.d_dataset), idx.d_inv_norms); + switch (idx.dataset_type) { + case GpuHnswDatasetType::INT8: + launch_kernels( + static_cast(idx.d_dataset), idx.d_inv_norms); + break; + case GpuHnswDatasetType::FP16: + launch_kernels( + static_cast(idx.d_dataset), idx.d_inv_norms); + break; + case GpuHnswDatasetType::BF16: + launch_kernels( + static_cast(idx.d_dataset), + idx.d_inv_norms); + break; + case GpuHnswDatasetType::FP32: + default: + launch_kernels( + static_cast(idx.d_dataset), idx.d_inv_norms); + break; } } diff --git a/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearchKernel.cuh b/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearchKernel.cuh index 86342f92f..2133f4327 100644 --- a/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearchKernel.cuh +++ b/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearchKernel.cuh @@ -23,6 +23,7 @@ #pragma once +#include #include #include @@ -43,6 +44,9 @@ __device__ __forceinline__ float load_elem(const float* ptr, int idx) { __device__ __forceinline__ float load_elem(const half* ptr, int idx) { return __half2float(__ldg(&ptr[idx])); } +__device__ __forceinline__ float load_elem(const __nv_bfloat16* ptr, int idx) { + return __bfloat162float(__ldg(&ptr[idx])); +} __device__ __forceinline__ float load_elem(const int8_t* ptr, int idx) { return static_cast(__ldg(&ptr[idx])); } diff --git a/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.h b/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.h index 6c69973e2..95c6c5fd9 100644 --- a/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.h +++ b/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.h @@ -34,6 +34,17 @@ namespace faiss { namespace gpu { +// Element type of the device-resident dataset. The graph walk kernel is +// templated on this; each value selects a load_elem specialization so the +// vectors stay in their native precision on the GPU (no up-conversion to +// fp32 at upload time). +enum class GpuHnswDatasetType { + FP32 = 0, + INT8 = 1, + FP16 = 2, + BF16 = 3, +}; + struct GpuHnswSearchParams { int ef = 200; int search_width = 4; @@ -134,7 +145,7 @@ class ScratchPoolGuard { struct GpuHnswDeviceIndex { void* d_dataset = nullptr; - bool dataset_int8 = false; + GpuHnswDatasetType dataset_type = GpuHnswDatasetType::FP32; float* d_inv_norms = nullptr; uint32_t* d_layer0_graph = nullptr; std::vector upper_layers; From 5a9413c3bd0e8f7182674a0b78fd0feed292323a Mon Sep 17 00:00:00 2001 From: premal Date: Sat, 11 Jul 2026 22:46:23 +0000 Subject: [PATCH 09/24] fix: address GPU_HNSW review findings (gpu_ready_ reload, self-recall bound) - Deserialize (reload path) now clears gpu_ready_ under gpu_mutex_ before resetting gpu_index_, so a failed re-upload no longer leaves the node advertising gpu_ready_==true with a null gpu_index_ (the lock-free reader in Search() would otherwise dereference null). - Round-trip self-recall test searched train_ds (nb rows) but only checked the first nq of nb results; loop and denominator now use nb. - Document the single-visible-GPU-per-process device model of the shared StandardGpuResources (multi-GPU-per-process device selection is a follow-up). Signed-off-by: premal Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- src/index/hnsw/faiss_hnsw.cc | 14 ++++++++++++++ tests/ut/test_gpu_search.cc | 8 +++++--- 2 files changed, 19 insertions(+), 3 deletions(-) diff --git a/src/index/hnsw/faiss_hnsw.cc b/src/index/hnsw/faiss_hnsw.cc index 68a4bf2c5..42b407f07 100644 --- a/src/index/hnsw/faiss_hnsw.cc +++ b/src/index/hnsw/faiss_hnsw.cc @@ -3298,6 +3298,15 @@ namespace knowhere { // reopen namespace knowhere // Process-global StandardGpuResources shared by all GpuHnswIndexNode instances. // Avoids per-segment 256 MiB pinned memory, cuBLAS handle, and CUDA stream // allocations that accumulate to tens of GiB with many segments. +// +// Device model: this assumes a single visible GPU per process, which is how +// Milvus GPU querynodes are deployed (one device pinned via CUDA_VISIBLE_DEVICES +// / MIG). GpuIndexHNSW below is constructed with the default device (0), so all +// segments in a process land on that one GPU. StandardGpuResources can manage +// multiple devices, but true multi-GPU-per-process would additionally require +// explicit per-segment device selection in the GpuIndexHNSW config; that is a +// follow-up and is intentionally not wired here (cannot be validated without +// multi-GPU capacity). static std::shared_ptr& GetSharedGpuResources() { static std::once_flag flag; @@ -3413,6 +3422,11 @@ class GpuHnswIndexNode : public BaseFaissRegularIndexHNSWNode { Status Deserialize(const BinarySet& binset, std::shared_ptr cfg) override { std::unique_lock lock(gpu_mutex_); + // Unpublish before tearing down: on a reload a stale gpu_ready_==true + // combined with a reset (or failed re-upload of) gpu_index_ would let the + // lock-free reader in Search() dereference a null index. Clear it under + // the same lock so a failed re-upload leaves the node honestly "not ready". + gpu_ready_.store(false, std::memory_order_release); gpu_index_.reset(); // Accept CPU-built HNSW (F32) or HNSW_SQ (quantized) binaries. diff --git a/tests/ut/test_gpu_search.cc b/tests/ut/test_gpu_search.cc index 9884c1839..28f2a72b4 100644 --- a/tests/ut/test_gpu_search.cc +++ b/tests/ut/test_gpu_search.cc @@ -592,16 +592,18 @@ TEST_CASE("Test All GPU Index", "[search]") { auto deser_res = gpu_idx.Deserialize(bs); REQUIRE(deser_res == knowhere::Status::success); - // Self-search: each vector should find itself as nearest neighbor + // Self-search: each vector should find itself as nearest neighbor. + // The query set is train_ds (nb rows), so check all nb results — using + // nq only exercised the first nq of nb vectors. auto results = gpu_idx.Search(train_ds, hnsw_json, nullptr); REQUIRE(results.has_value()); auto ids = results.value()->GetIds(); int correct = 0; - for (int i = 0; i < nq; ++i) { + for (int i = 0; i < nb; ++i) { if (ids[i] == i) correct++; } - float self_recall = static_cast(correct) / nq; + float self_recall = static_cast(correct) / nb; REQUIRE(self_recall >= 0.95f); } From f4d7df318ddae07acdebf79285e4d6e7493e057a Mon Sep 17 00:00:00 2001 From: premal Date: Sat, 11 Jul 2026 22:49:34 +0000 Subject: [PATCH 10/24] test: add native int8 GPU_HNSW deserialize+search test The existing SQ8 test builds an unsigned QT_8bit HNSW_SQ, which routes through the decode-to-fp32 upload branch, so it does not exercise the native signed-int8 device path (upload_int8_dataset / QT_8bit_direct_signed). Add an int8 test that builds a CPU int8 HNSW and loads it as GPU_HNSW, covering all four native device representations (fp32/int8/fp16/bf16). Signed-off-by: premal Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/ut/test_gpu_search.cc | 42 +++++++++++++++++++++++++++++++++++++ 1 file changed, 42 insertions(+) diff --git a/tests/ut/test_gpu_search.cc b/tests/ut/test_gpu_search.cc index 28f2a72b4..db4713660 100644 --- a/tests/ut/test_gpu_search.cc +++ b/tests/ut/test_gpu_search.cc @@ -646,6 +646,48 @@ TEST_CASE("Test All GPU Index", "[search]") { REQUIRE(recall >= 0.7f); } + SECTION("Test GPU HNSW INT8 Deserialization") { + // Build a CPU int8 HNSW (HNSW_SQ QT_8bit_direct_signed storage), then + // load and search on GPU. int8 vectors stay in their native 1-byte + // layout on the device (upload_int8_dataset) and are up-converted to + // fp32 per element in the search kernel. This exercises the native + // signed-int8 device path, distinct from the unsigned SQ8 test above + // which routes through the decode-to-fp32 upload branch. + knowhere::Json hnsw_json; + hnsw_json[knowhere::meta::DIM] = dim; + hnsw_json[knowhere::meta::METRIC_TYPE] = knowhere::metric::L2; + hnsw_json[knowhere::meta::TOPK] = 1; + hnsw_json[knowhere::indexparam::HNSW_M] = 16; + hnsw_json[knowhere::indexparam::EFCONSTRUCTION] = 200; + hnsw_json[knowhere::indexparam::EF] = 200; + + auto train_ds = knowhere::ConvertToDataTypeIfNeeded(GenDataSet(nb, dim, seed)); + auto query_ds = knowhere::ConvertToDataTypeIfNeeded(GenDataSet(nq, dim, seed + 2)); + + auto cpu_idx = + knowhere::IndexFactory::Instance().Create(knowhere::IndexEnum::INDEX_HNSW, version).value(); + auto res = cpu_idx.Build(train_ds, hnsw_json); + REQUIRE(res == knowhere::Status::success); + + knowhere::BinarySet bs; + cpu_idx.Serialize(bs); + + auto gpu_idx = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_GPU_HNSW, version) + .value(); + auto deser_res = gpu_idx.Deserialize(bs); + REQUIRE(deser_res == knowhere::Status::success); + + auto results = gpu_idx.Search(query_ds, hnsw_json, nullptr); + REQUIRE(results.has_value()); + + auto gt = knowhere::BruteForce::Search(train_ds, query_ds, hnsw_json, nullptr); + REQUIRE(gt.has_value()); + float recall = GetKNNRecall(*gt.value(), *results.value()); + // int8 native path should match the CPU int8 HNSW recall closely. + REQUIRE(recall >= 0.9f); + } + SECTION("Test GPU HNSW FP16 Deserialization") { // Build a CPU HNSW (fp16 -> HNSW_SQ QT_fp16 storage), then load and // search on GPU. fp16 vectors stay in their native 2-byte layout on the From 71c4c785115ef38c774bd5cab336db404c96f18c Mon Sep 17 00:00:00 2001 From: premal Date: Sun, 12 Jul 2026 00:00:40 +0000 Subject: [PATCH 11/24] harden GPU HNSW: fail Deserialize on GPU upload error; mirror faiss review fixes - M1: GpuHnswIndexNode::Deserialize now returns Status::cuda_runtime_error when the eager GPU upload throws, instead of returning success with no GPU index (which only surfaced the failure on first search). Also log via LOG_KNOWHERE_ERROR_ rather than fprintf(stderr,...) so it lands in Milvus's structured logs. - Mirror the byte-identical vendored faiss changes (H1 scratch-release sync, L1/L2/C1/H5 documentation) to keep 6si/faiss and the vendored copy in lockstep. Signed-off-by: premal Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- src/index/hnsw/faiss_hnsw.cc | 7 ++++++- thirdparty/faiss/faiss/gpu/GpuIndexHNSW.h | 19 +++++++++++++++++-- .../faiss/faiss/gpu/impl/GpuHnswBuild.cuh | 13 +++++++++++-- .../faiss/faiss/gpu/impl/GpuHnswTypes.cu | 8 ++++++++ .../faiss/faiss/gpu/impl/GpuHnswTypes.h | 3 +++ 5 files changed, 45 insertions(+), 5 deletions(-) diff --git a/src/index/hnsw/faiss_hnsw.cc b/src/index/hnsw/faiss_hnsw.cc index 42b407f07..dcd46050a 100644 --- a/src/index/hnsw/faiss_hnsw.cc +++ b/src/index/hnsw/faiss_hnsw.cc @@ -3475,8 +3475,13 @@ class GpuHnswIndexNode : public BaseFaissRegularIndexHNSWNode { // reader in Search() pairs this with an acquire load. gpu_ready_.store(true, std::memory_order_release); } catch (const std::exception& e) { - fprintf(stderr, "[gpu_hnsw] eager GPU upload failed: %s\n", e.what()); + // Fail the load rather than reporting success with no GPU index: + // a silently "loaded" segment would only surface the error on the + // first search. Reset so gpu_ready_ stays false and return an error + // status so the caller treats the segment as failed to load. + LOG_KNOWHERE_ERROR_ << "eager GPU upload failed: " << e.what(); gpu_index_.reset(); + return Status::cuda_runtime_error; } } return Status::success; diff --git a/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.h b/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.h index 67445c844..0e4848117 100644 --- a/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.h +++ b/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.h @@ -64,7 +64,18 @@ struct SearchParametersGpuHNSW : SearchParameters { /// parallel beam search kernel. /// /// Supports L2, inner product, and cosine metrics. -/// Supports float32 and int8 (QT_8bit_direct_signed) data. +/// Supports float32, fp16, bf16, and int8 (QT_8bit_direct_signed) data. +/// Low-precision formats (int8/fp16/bf16) are kept in their native on-device +/// byte layout and up-converted to fp32 per element inside the search kernel. +/// +/// Two search entry points exist: +/// - searchHost(): the path Knowhere uses. Host in/out pointers, a single +/// device sync at the end. Preferred for production. +/// - searchImpl_(): the faiss-standard GpuIndex::search() override. It does a +/// D2H copy of labels, a CPU uint64->idx_t conversion, then an H2D copy +/// back, with a stream sync on each side. Correct but not latency-optimal; +/// a GPU-side label-conversion kernel to avoid the round-trip is a +/// documented follow-up. Knowhere does not use this path. struct GpuIndexHNSW : public GpuIndex { public: GpuIndexHNSW( @@ -78,7 +89,8 @@ struct GpuIndexHNSW : public GpuIndex { /// Load an HNSW index from CPU to GPU. /// The CPU index must have been built and trained already. /// Supports IndexHNSWFlat (float32) and IndexHNSWSQ - /// (QT_8bit_direct_signed for INT8, or dequantized for other SQ types). + /// (QT_8bit_direct_signed for INT8, QT_fp16/QT_bf16 kept native, or + /// dequantized for other SQ types). void copyFrom(const faiss::cppcontrib::knowhere::IndexHNSW* index); /// Load with explicit metric specification. @@ -100,6 +112,9 @@ struct GpuIndexHNSW : public GpuIndex { /// All input/output pointers must be host memory. /// This avoids the GpuIndex::search_ex temp allocation chain /// which can cause SIGSEGV from pointer lifetime issues. + /// This is the preferred entry point (single device sync); prefer it + /// over the searchImpl_/GpuIndex::search() path, which round-trips the + /// labels D2H then H2D. void searchHost( idx_t n, const float* x_host, diff --git a/thirdparty/faiss/faiss/gpu/impl/GpuHnswBuild.cuh b/thirdparty/faiss/faiss/gpu/impl/GpuHnswBuild.cuh index e86e2abd7..bce0ede9b 100644 --- a/thirdparty/faiss/faiss/gpu/impl/GpuHnswBuild.cuh +++ b/thirdparty/faiss/faiss/gpu/impl/GpuHnswBuild.cuh @@ -61,8 +61,17 @@ namespace gpu { /// Extract HNSW graph layers from a Knowhere HNSW struct. /// Template parameter HnswT can be faiss::cppcontrib::knowhere::HNSW -/// or faiss::HNSW — any type with neighbor_range(), nb_neighbors(), -/// neighbors, levels, entry_point, max_level. +/// or faiss::HNSW — any type exposing this interface: +/// - neighbor_range(node, layer, &begin, &end) +/// - nb_neighbors(layer) -> int +/// - neighbors : flat neighbor array indexed by neighbor_range +/// - levels : per-node array; levels[i] is the 1-based layer count, so +/// node i lives on layers 0..levels[i]-1 and "i is on layer L" +/// is the test levels[i] > L +/// - entry_point, max_level +/// The levels[] convention above matches faiss::HNSW (HNSW.cpp uses +/// pt_level = levels[i] - 1); a variant with different semantics would need a +/// different membership test at the levels[i] > layer check below. template inline void extract_hnsw_layers( const HnswT& hnsw, diff --git a/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.cu b/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.cu index 4895aa94a..5973b94ac 100644 --- a/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.cu +++ b/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.cu @@ -140,6 +140,14 @@ GpuHnswScratchSlot* GpuHnswScratchPool::acquire() { } void GpuHnswScratchPool::release(GpuHnswScratchSlot* slot) { + // Safety net: ensure all GPU work on this slot's stream has completed + // before returning it to the pool. Callers (searchHost / searchImpl_) + // already synchronize before release, so this is normally a no-op; it + // guards a future caller that skips its own sync, whose in-flight buffers + // could otherwise be freed+realloced by the next acquirer's ensure() + // (use-after-free). + cudaSetDevice(device_); + cudaStreamSynchronize(slot->stream); { std::lock_guard lock(mutex_); available_.push_back(slot); diff --git a/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.h b/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.h index 95c6c5fd9..25aa68f80 100644 --- a/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.h +++ b/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.h @@ -153,6 +153,9 @@ struct GpuHnswDeviceIndex { int64_t n_rows = 0; int64_t dim = 0; uint32_t entry_point = 0; + // Total layer count (max_level + 1), informational only; the search path + // uses num_upper_layers_built (below) for the number of uploaded upper + // layers. Kept for diagnostics/logging. int num_layers = 0; int M = 0; int max_degree0 = 0; From 08ced662099da5209f803a4b6a5f356383fd90e0 Mon Sep 17 00:00:00 2001 From: premal Date: Sun, 12 Jul 2026 00:58:14 +0000 Subject: [PATCH 12/24] fix(gpu_hnsw): quantized-cosine norms, reload lifetime safety, peak load accounting Codex P1 review fixes for the GPU HNSW path: #1 Quantized-cosine inverse-norm semantics. The CPU cosine index records inverse L2 norms from the ORIGINAL input vectors (HasInverseL2Norms:: get_inverse_l2_norms()). The GPU upload previously recomputed norms from decoded/quantized codes (and normalized decoded fp32 vectors in place), which diverges from the CPU scores/traversal for lossy SQ8/int8/fp16/bf16. Now upload the CPU index's stored inverse norms and apply them in the kernel; only flat-fp32 cosine (decoded==original) keeps normalizing in place. Vendored faiss copy kept byte-identical with 6si/faiss. #2 Reload/search use-after-free race. gpu_index_ is now an atomic>; Search() takes a shared-ownership snapshot for the whole search so a concurrent Deserialize() reload that resets/rebuilds the member cannot free an in-use index. Atomic readiness only guarantees visibility, not lifetime. #3 Phase-separated load-resource accounting. Resource gains maxMemoryCost (transient peak during load; defaults to 0 so other indexes keep their existing heuristic). GpuHnswIndexNode reports memoryCost=0 (CPU copy freed after upload) and maxMemoryCost covering the download buffer + deserialized CPU index + decode/graph staging. Tests (tests/ut/test_gpu_search.cc, [gpu_hnsw_p1]): native low-precision device paths (fp16/bf16/sq8 via non-mock fp32 node), SQ8/fp16/bf16 cosine CPU-vs-GPU parity asserting per-query top-1 id + per-result score on same-direction/ different-magnitude vectors, search-vs-concurrent-reload lifetime stress, >4 concurrent searches with varying nq/k/ef vs own ground truth, and load-resource estimate assertions. Co-Authored-By: Claude Opus 4.6 Signed-off-by: premal Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- include/knowhere/index/index_static.h | 10 +- src/index/hnsw/faiss_hnsw.cc | 71 +++-- tests/ut/test_gpu_search.cc | 297 ++++++++++++++++++ .../faiss/faiss/gpu/impl/GpuHnswBuild.cuh | 135 ++++---- 4 files changed, 429 insertions(+), 84 deletions(-) diff --git a/include/knowhere/index/index_static.h b/include/knowhere/index/index_static.h index f6279f613..001e7796f 100644 --- a/include/knowhere/index/index_static.h +++ b/include/knowhere/index/index_static.h @@ -19,8 +19,16 @@ namespace knowhere { struct Resource { - uint64_t memoryCost; // in bytes + uint64_t memoryCost; // retained host memory after load completes, in bytes uint64_t diskCost; // in bytes + // Peak transient host memory required *during* load, in bytes. Defaults to 0 + // so existing indexes are unaffected: consumers should fall back to their + // memoryCost-based heuristic when this is 0. Indexes whose peak load + // footprint differs from the retained footprint (e.g. GPU_HNSW frees its CPU + // copy after uploading to VRAM, so memoryCost≈0 but the peak covers the + // download buffer + deserialized CPU index + decode/graph staging) set this + // to the peak so the loader reserves enough host RAM and does not OOM. + uint64_t maxMemoryCost = 0; }; #define DEFINE_HAS_STATIC_FUNC(func_name) \ diff --git a/src/index/hnsw/faiss_hnsw.cc b/src/index/hnsw/faiss_hnsw.cc index dcd46050a..6cdb2cd8d 100644 --- a/src/index/hnsw/faiss_hnsw.cc +++ b/src/index/hnsw/faiss_hnsw.cc @@ -3357,10 +3357,21 @@ class GpuHnswIndexNode : public BaseFaissRegularIndexHNSWNode { static expected StaticEstimateLoadResource(const uint64_t file_size_in_bytes, const int64_t num_rows, const int64_t dim, const knowhere::BaseConfig& config, const IndexVersion& version) { - // GPU HNSW stores vectors and graph in VRAM; the CPU copy is freed - // after upload in Deserialize(). Report zero CPU memory cost so the - // Milvus segment loader does not over-commit host RAM reservations. - return Resource{.memoryCost = 0, .diskCost = 0}; + // GPU HNSW stores vectors and graph in VRAM; the CPU copy is freed after + // upload in Deserialize(). Phase-separated accounting: + // memoryCost (retained after load): ~0 — the CPU index is released + // once the graph/vectors are on the GPU. + // maxMemoryCost (transient peak during load): the loader must reserve + // enough host RAM to cover the simultaneously-live buffers before the + // GPU upload frees the CPU copy: + // - the serialized download buffer (~file_size), + // - the deserialized CPU HNSW index: vectors + graph (~file_size), + // - fp32 decode/graph staging for the upload (num_rows*dim*4). + const uint64_t rows = num_rows > 0 ? static_cast(num_rows) : 0; + const uint64_t d = dim > 0 ? static_cast(dim) : 0; + const uint64_t decode_staging = rows * d * sizeof(float); + const uint64_t peak = file_size_in_bytes + file_size_in_bytes + decode_staging; + return Resource{.memoryCost = 0, .diskCost = 0, .maxMemoryCost = peak}; } std::unique_ptr @@ -3427,7 +3438,7 @@ class GpuHnswIndexNode : public BaseFaissRegularIndexHNSWNode { // lock-free reader in Search() dereference a null index. Clear it under // the same lock so a failed re-upload leaves the node honestly "not ready". gpu_ready_.store(false, std::memory_order_release); - gpu_index_.reset(); + gpu_index_.store(nullptr, std::memory_order_release); // Accept CPU-built HNSW (F32) or HNSW_SQ (quantized) binaries. Status status; @@ -3458,29 +3469,31 @@ class GpuHnswIndexNode : public BaseFaissRegularIndexHNSWNode { dynamic_cast(faiss_idx) != nullptr; bool use_ip = is_cosine || (faiss_idx->metric_type == ::faiss::METRIC_INNER_PRODUCT); + std::shared_ptr local; { std::lock_guard gpu_ctor_lock(GetGpuConstructionMutex()); gpu_resources_ = GetSharedGpuResources(); - gpu_index_ = std::make_unique(gpu_resources_.get(), faiss_idx->d, - faiss_idx->metric_type); + local = std::make_shared(gpu_resources_.get(), faiss_idx->d, + faiss_idx->metric_type); } - gpu_index_->copyFromWithMetric(faiss_idx, use_ip, is_cosine); + local->copyFromWithMetric(faiss_idx, use_ip, is_cosine); // Capture count/dim before releasing the CPU copy so Count()/ // Dim() keep working once indexes[0] is freed. gpu_ntotal_ = faiss_idx->ntotal; gpu_dim_ = faiss_idx->d; // Release CPU copy — vectors and graph are now on GPU. indexes[0].reset(); - // Publish gpu_index_ with release semantics; the lock-free - // reader in Search() pairs this with an acquire load. + // Publish the fully-built index (release); Search() takes an + // acquire snapshot. gpu_ready_ is published last for Count()/Dim(). + gpu_index_.store(local, std::memory_order_release); gpu_ready_.store(true, std::memory_order_release); } catch (const std::exception& e) { // Fail the load rather than reporting success with no GPU index: // a silently "loaded" segment would only surface the error on the - // first search. Reset so gpu_ready_ stays false and return an error - // status so the caller treats the segment as failed to load. + // first search. Leave the member null so gpu_ready_ stays false + // and return an error so the caller treats the load as failed. LOG_KNOWHERE_ERROR_ << "eager GPU upload failed: " << e.what(); - gpu_index_.reset(); + gpu_index_.store(nullptr, std::memory_order_release); return Status::cuda_runtime_error; } } @@ -3498,11 +3511,15 @@ class GpuHnswIndexNode : public BaseFaissRegularIndexHNSWNode { return expected::Err(Status::invalid_args, "GPU_HNSW does not support filtered search"); } - // Fast path: gpu_ready_ is published (release) once gpu_index_ is fully - // constructed; the acquire load here avoids a data race on the pointer. - if (!gpu_ready_.load(std::memory_order_acquire)) { + // Take a shared-ownership snapshot of the GPU index up front. It keeps + // the index alive for the whole search even if a concurrent + // Deserialize() reload resets/rebuilds gpu_index_; a bare atomic + // readiness flag would guarantee visibility but not lifetime. + std::shared_ptr gpu_snapshot = gpu_index_.load(std::memory_order_acquire); + if (!gpu_snapshot) { std::unique_lock lock(gpu_mutex_); - if (!gpu_ready_.load(std::memory_order_relaxed)) { + gpu_snapshot = gpu_index_.load(std::memory_order_relaxed); + if (!gpu_snapshot) { const auto* faiss_idx = GetFaissHnswIndex(); if (!faiss_idx) { return expected::Err(Status::empty_index, "index not loaded"); @@ -3512,17 +3529,20 @@ class GpuHnswIndexNode : public BaseFaissRegularIndexHNSWNode { bool is_cosine = IsMetricType(hnsw_cfg.metric_type.value(), metric::COSINE); bool use_ip = IsMetricType(hnsw_cfg.metric_type.value(), metric::IP) || is_cosine; + std::shared_ptr local; { std::lock_guard gpu_ctor_lock(GetGpuConstructionMutex()); gpu_resources_ = GetSharedGpuResources(); - gpu_index_ = std::make_unique(gpu_resources_.get(), faiss_idx->d, - faiss_idx->metric_type); + local = std::make_shared(gpu_resources_.get(), faiss_idx->d, + faiss_idx->metric_type); } - gpu_index_->copyFromWithMetric(faiss_idx, use_ip, is_cosine); + local->copyFromWithMetric(faiss_idx, use_ip, is_cosine); gpu_ntotal_ = faiss_idx->ntotal; gpu_dim_ = faiss_idx->d; const_cast&>(indexes[0]).reset(); + gpu_index_.store(local, std::memory_order_release); gpu_ready_.store(true, std::memory_order_release); + gpu_snapshot = local; } catch (const std::exception& e) { return expected::Err(Status::cuvs_inner_error, std::string("failed to build GPU HNSW index: ") + e.what()); @@ -3559,7 +3579,7 @@ class GpuHnswIndexNode : public BaseFaissRegularIndexHNSWNode { try { faiss::gpu::GpuHnswSearchParams gsp; gsp.ef = ef; - gpu_index_->searchHost(nq, h_queries, k, h_dist.get(), h_ids.get(), gsp); + gpu_snapshot->searchHost(nq, h_queries, k, h_dist.get(), h_ids.get(), gsp); } catch (const std::exception& e) { LOG_KNOWHERE_ERROR_ << "GPU_HNSW search failed: " << e.what(); return expected::Err(Status::cuvs_inner_error, @@ -3589,9 +3609,14 @@ class GpuHnswIndexNode : public BaseFaissRegularIndexHNSWNode { mutable std::mutex gpu_mutex_; mutable std::shared_ptr gpu_resources_; - mutable std::unique_ptr gpu_index_; + // Held as an atomic shared_ptr so Search() can take a shared-ownership + // snapshot that keeps the GPU index alive for the entire search even if a + // concurrent Deserialize() reload resets/rebuilds the member. Atomic + // readiness (gpu_ready_) only guarantees visibility, not lifetime, so a + // bare pointer read on the search fast path would be a use-after-free. + mutable std::atomic> gpu_index_; // Published (release) after gpu_index_ is fully built; read lock-free - // (acquire) on the search fast path and by Count()/Dim()/HasRawData(). + // (acquire) by Count()/Dim()/HasRawData()/GetVectorByIds(). mutable std::atomic gpu_ready_{false}; // Vector count/dim captured before the CPU copy is freed (written from the // const lazy-upload path in Search(), hence mutable). diff --git a/tests/ut/test_gpu_search.cc b/tests/ut/test_gpu_search.cc index db4713660..c5be0b903 100644 --- a/tests/ut/test_gpu_search.cc +++ b/tests/ut/test_gpu_search.cc @@ -9,7 +9,11 @@ // is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express // or implied. See the License for the specific language governing permissions and limitations under the License. +#include +#include #include +#include +#include #include "catch2/catch_approx.hpp" #include "catch2/catch_test_macros.hpp" @@ -954,4 +958,297 @@ TEST_CASE("Test CPU vs GPU HNSW Comparison", "[gpu_hnsw_compare]") { fprintf(stderr, "===\n\n"); } } + +// Regression tests for the three Codex P1 findings on the gpu-hnsw-faiss branch: +// #1 quantized-cosine inverse-norm semantics (native low-precision kernels + +// SQ8/FP16/BF16 cosine parity), +// #2 search-vs-reload lifetime safety, +// #3 phase-separated load-resource accounting. +TEST_CASE("Test GPU HNSW Codex P1 Regressions", "[gpu_hnsw_p1]") { + using Catch::Approx; + + const int64_t nb = 4000, nq = 200; + const int64_t dim = 64; + const int64_t seed = 42; + const auto version = GenTestVersionList(); + + auto sq_json = [&](const std::string& metric, const std::string& sq_type, int64_t topk) { + knowhere::Json json; + json[knowhere::meta::DIM] = dim; + json[knowhere::meta::METRIC_TYPE] = metric; + json[knowhere::meta::TOPK] = topk; + json[knowhere::indexparam::HNSW_M] = 16; + json[knowhere::indexparam::EFCONSTRUCTION] = 200; + json[knowhere::indexparam::EF] = 200; + json[knowhere::indexparam::SQ_TYPE] = sq_type; + return json; + }; + + // Build an HNSW_SQ index on the fp32 (real, non-mock) node so the serialized + // storage carries the requested SQ qtype. Deserializing that as GPU_HNSW on + // the fp32 node selects the native device representation from the storage + // qtype (FP16 -> half kernel, BF16 -> __nv_bfloat16 kernel, SQ8 -> fp32 + // decode). This bypasses the int8/fp16/bf16 IndexNodeDataMockWrapper, which + // otherwise converts data to fp32 and routes the fp32 kernel. + + // #1a: native low-precision (FP16/BF16) and SQ8 device paths — L2. + SECTION("Native low-precision device path (L2): FP16 / BF16 / SQ8") { + const auto sq_type = GENERATE(std::string("fp16"), std::string("bf16"), std::string("sq8")); + CAPTURE(sq_type); + auto json = sq_json(knowhere::metric::L2, sq_type, 10); + + auto train_ds = GenDataSet(nb, dim, seed); + auto query_ds = GenDataSet(nq, dim, seed + 2); + + auto cpu_idx = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_HNSW_SQ, version) + .value(); + REQUIRE(cpu_idx.Build(train_ds, json) == knowhere::Status::success); + knowhere::BinarySet bs; + cpu_idx.Serialize(bs); + + auto gpu_idx = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_GPU_HNSW, version) + .value(); + REQUIRE(gpu_idx.Deserialize(bs) == knowhere::Status::success); + + auto results = gpu_idx.Search(query_ds, json, nullptr); + REQUIRE(results.has_value()); + auto gt = knowhere::BruteForce::Search(train_ds, query_ds, json, nullptr); + REQUIRE(gt.has_value()); + float recall = GetKNNRecall(*gt.value(), *results.value()); + // Lossy SQ paths; a conservative floor. fp16/bf16 are near-lossless, sq8 + // coarser. + REQUIRE(recall >= (sq_type == "sq8" ? 0.65f : 0.85f)); + } + + // #1b: quantized-cosine parity. The discriminating input is vectors that + // share directions but have widely different magnitudes: cosine must ignore + // magnitude. The CPU cosine index records inverse norms from the ORIGINAL + // input vectors; the GPU path must upload those same norms rather than + // recomputing them from lossily-decoded codes. With the pre-fix behavior + // (recompute-from-decoded / normalize-decoded) the SQ8 scores and top-1 IDs + // diverge from the CPU index. Assert per-query top-1 ID and per-result score, + // not just aggregate recall. + SECTION("Quantized cosine parity (CPU vs GPU): SQ8 / FP16 / BF16") { + const auto sq_type = GENERATE(std::string("sq8"), std::string("fp16"), std::string("bf16")); + CAPTURE(sq_type); + const int64_t topk = 10; + auto json = sq_json(knowhere::metric::COSINE, sq_type, topk); + + // Directions from a fixed RNG, then scale each row by a magnitude that + // spans three orders of magnitude so magnitude cannot be ignored by luck. + auto make_scaled = [&](int64_t rows, int64_t s) { + auto ds = GenDataSet(rows, dim, s); + auto* data = const_cast(static_cast(ds->GetTensor())); + std::mt19937 rng(static_cast(s + 7)); + std::uniform_real_distribution mag(0.5f, 500.0f); + for (int64_t r = 0; r < rows; ++r) { + float scale = mag(rng); + for (int64_t d = 0; d < dim; ++d) { + data[r * dim + d] *= scale; + } + } + return ds; + }; + auto train_ds = make_scaled(nb, seed); + auto query_ds = make_scaled(nq, seed + 2); + + auto cpu_idx = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_HNSW_SQ, version) + .value(); + REQUIRE(cpu_idx.Build(train_ds, json) == knowhere::Status::success); + auto cpu_results = cpu_idx.Search(query_ds, json, nullptr); + REQUIRE(cpu_results.has_value()); + + knowhere::BinarySet bs; + cpu_idx.Serialize(bs); + auto gpu_idx = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_GPU_HNSW, version) + .value(); + REQUIRE(gpu_idx.Deserialize(bs) == knowhere::Status::success); + auto gpu_results = gpu_idx.Search(query_ds, json, nullptr); + REQUIRE(gpu_results.has_value()); + + const auto* cpu_ids = cpu_results.value()->GetIds(); + const auto* gpu_ids = gpu_results.value()->GetIds(); + const auto* cpu_dist = cpu_results.value()->GetDistance(); + const auto* gpu_dist = gpu_results.value()->GetDistance(); + + // Top-1 must match the CPU index for the overwhelming majority of + // queries: identical cosine semantics (same norms) means identical + // nearest neighbor modulo rare quantization ties. + int top1_match = 0; + for (int64_t q = 0; q < nq; ++q) { + if (cpu_ids[q * topk] == gpu_ids[q * topk]) { + top1_match++; + } + } + float top1_ratio = static_cast(top1_match) / nq; + CAPTURE(top1_ratio); + REQUIRE(top1_ratio >= 0.95f); + + // Per-result cosine scores must agree closely (both decode identical SQ + // codes and apply the SAME original inverse norms). A loose floor would + // hide the pre-fix norm mismatch, so keep the tolerance tight. + int compared = 0, close = 0; + for (int64_t q = 0; q < nq; ++q) { + // Only compare positions where CPU and GPU agree on the id, so score + // comparison is like-for-like. + for (int64_t j = 0; j < topk; ++j) { + if (cpu_ids[q * topk + j] == gpu_ids[q * topk + j]) { + compared++; + if (GetRelativeLoss(cpu_dist[q * topk + j], gpu_dist[q * topk + j]) < 0.02f) { + close++; + } + } + } + } + REQUIRE(compared > 0); + float close_ratio = static_cast(close) / compared; + CAPTURE(close_ratio); + REQUIRE(close_ratio >= 0.98f); + } + + // #2: search must not use-after-free a GPU index that a concurrent reload + // resets/rebuilds. Search continuously from several threads while another + // thread repeatedly Deserialize()s (reloads) the same index. Run under + // TSAN/ASAN to catch the lifetime race the atomic snapshot fixes. + SECTION("Search vs concurrent reload (lifetime safety)") { + auto json = sq_json(knowhere::metric::L2, "fp16", 10); + json.erase(knowhere::indexparam::SQ_TYPE); // plain fp32 HNSW for the reload driver + auto train_ds = GenDataSet(nb, dim, seed); + auto query_ds = GenDataSet(nq, dim, seed + 2); + + auto cpu_idx = + knowhere::IndexFactory::Instance().Create(knowhere::IndexEnum::INDEX_HNSW, version).value(); + REQUIRE(cpu_idx.Build(train_ds, json) == knowhere::Status::success); + knowhere::BinarySet bs; + cpu_idx.Serialize(bs); + + auto gpu_idx = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_GPU_HNSW, version) + .value(); + REQUIRE(gpu_idx.Deserialize(bs) == knowhere::Status::success); + + std::atomic stop{false}; + std::atomic failed{false}; + std::vector searchers; + for (int t = 0; t < 6; ++t) { + searchers.emplace_back([&]() { + while (!stop.load(std::memory_order_relaxed)) { + auto r = gpu_idx.Search(query_ds, json, nullptr); + // A search may race a reload window; success or a benign + // "not ready" error are both acceptable, a crash is not. + if (!r.has_value() && r.error() != knowhere::Status::empty_index && + r.error() != knowhere::Status::cuda_runtime_error) { + failed.store(true, std::memory_order_relaxed); + } + } + }); + } + std::thread reloader([&]() { + for (int i = 0; i < 30 && !stop.load(std::memory_order_relaxed); ++i) { + if (gpu_idx.Deserialize(bs) != knowhere::Status::success) { + failed.store(true, std::memory_order_relaxed); + } + } + stop.store(true, std::memory_order_relaxed); + }); + reloader.join(); + stop.store(true, std::memory_order_relaxed); + for (auto& th : searchers) { + th.join(); + } + REQUIRE(!failed.load()); + // Index is still usable after the reload storm. + auto after = gpu_idx.Search(query_ds, json, nullptr); + REQUIRE(after.has_value()); + } + + // Codex coverage #3: more than four simultaneous searches (the device scratch + // pool has four slots), with varying nq/k/ef, each validated against its own + // brute-force ground truth. Exercises scratch-pool contention/serialization. + SECTION("More than four concurrent searches with varying nq/k/ef") { + auto build_json = sq_json(knowhere::metric::L2, "fp16", 10); + build_json.erase(knowhere::indexparam::SQ_TYPE); + auto train_ds = GenDataSet(nb, dim, seed); + + auto cpu_idx = + knowhere::IndexFactory::Instance().Create(knowhere::IndexEnum::INDEX_HNSW, version).value(); + REQUIRE(cpu_idx.Build(train_ds, build_json) == knowhere::Status::success); + knowhere::BinarySet bs; + cpu_idx.Serialize(bs); + auto gpu_idx = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_GPU_HNSW, version) + .value(); + REQUIRE(gpu_idx.Deserialize(bs) == knowhere::Status::success); + + struct Case { + int64_t nq; + int64_t k; + int ef; + }; + const std::vector cases = {{50, 5, 32}, {120, 10, 64}, {200, 20, 128}, {80, 50, 256}, + {30, 10, 512}, {150, 1, 48}, {64, 100, 300}, {90, 25, 96}}; + std::atomic ok{true}; + std::vector threads; + for (const auto& c : cases) { + threads.emplace_back([&, c]() { + knowhere::Json j; + j[knowhere::meta::DIM] = dim; + j[knowhere::meta::METRIC_TYPE] = knowhere::metric::L2; + j[knowhere::meta::TOPK] = c.k; + j[knowhere::indexparam::EF] = c.ef; + auto q = GenDataSet(c.nq, dim, seed + 100 + c.ef); + auto r = gpu_idx.Search(q, j, nullptr); + if (!r.has_value()) { + ok.store(false, std::memory_order_relaxed); + return; + } + auto gt = knowhere::BruteForce::Search(train_ds, q, j, nullptr); + if (!gt.has_value() || GetKNNRecall(*gt.value(), *r.value()) < 0.7f) { + ok.store(false, std::memory_order_relaxed); + } + }); + } + for (auto& th : threads) { + th.join(); + } + REQUIRE(ok.load()); + } + + // #3: phase-separated load-resource accounting. GPU_HNSW frees its CPU copy + // after uploading to VRAM, so retained memoryCost is 0 while the transient + // peak (maxMemoryCost) must be non-zero and cover the download buffer + + // deserialized CPU index + decode/graph staging. This estimate is a static + // computation and needs no live GPU. + SECTION("Load-resource estimate: memoryCost=0, maxMemoryCost covers peak") { + knowhere::Json params; + params[knowhere::meta::METRIC_TYPE] = knowhere::metric::L2; + params[knowhere::meta::DIM] = dim; + + const uint64_t file_size = static_cast(nb) * dim * sizeof(float); + auto gpu_res = knowhere::IndexStaticFaced::EstimateLoadResource( + knowhere::IndexEnum::INDEX_GPU_HNSW, version, file_size, nb, dim, params); + REQUIRE(gpu_res.has_value()); + // Retained host memory is ~0 after the CPU copy is freed. + REQUIRE(gpu_res.value().memoryCost == 0); + REQUIRE(gpu_res.value().diskCost == 0); + // Transient peak is non-zero and at least covers the fp32 decode staging + // plus the download buffer (2*file_size + rows*dim*4 in the estimator). + REQUIRE(gpu_res.value().maxMemoryCost > 0); + const uint64_t decode_staging = static_cast(nb) * dim * sizeof(float); + REQUIRE(gpu_res.value().maxMemoryCost >= file_size + decode_staging); + + // A plain CPU HNSW keeps its data resident: memoryCost is non-zero and it + // does not opt into the peak field (maxMemoryCost stays 0 -> Milvus falls + // back to its 2*memoryCost heuristic). + auto cpu_res = knowhere::IndexStaticFaced::EstimateLoadResource( + knowhere::IndexEnum::INDEX_HNSW, version, file_size, nb, dim, params); + REQUIRE(cpu_res.has_value()); + REQUIRE(cpu_res.value().memoryCost > 0); + REQUIRE(cpu_res.value().maxMemoryCost == 0); + } +} #endif diff --git a/thirdparty/faiss/faiss/gpu/impl/GpuHnswBuild.cuh b/thirdparty/faiss/faiss/gpu/impl/GpuHnswBuild.cuh index bce0ede9b..29e121e8f 100644 --- a/thirdparty/faiss/faiss/gpu/impl/GpuHnswBuild.cuh +++ b/thirdparty/faiss/faiss/gpu/impl/GpuHnswBuild.cuh @@ -32,6 +32,7 @@ #include #include #include +#include #include #include @@ -215,13 +216,36 @@ inline void upload_graph_to_gpu( } +// Upload precomputed per-row inverse L2 norms to the device. These must be the +// norms of the *original* input vectors as recorded by the CPU cosine index +// (HasInverseL2Norms::get_inverse_l2_norms()), NOT norms recomputed from +// lossily-decoded codes — otherwise cosine scores and graph traversal diverge +// from the CPU index for lossy SQ (SQ8 / fp16 / bf16). The search kernel +// computes score = IP(query, decoded_db) * inv_norm[row], mirroring the CPU +// WithCosineNormDistanceComputer. +inline void upload_inv_norms( + GpuHnswDeviceIndex& idx, + const float* inv_norms, + int64_t n_rows) { + size_t norms_bytes = static_cast(n_rows) * sizeof(float); + GPU_HNSW_BUILD_CUDA_CHECK(cudaMalloc(&idx.d_inv_norms, norms_bytes)); + GPU_HNSW_BUILD_CUDA_CHECK(cudaMemcpy( + idx.d_inv_norms, inv_norms, norms_bytes, cudaMemcpyHostToDevice)); +} + inline void upload_fp32_dataset( GpuHnswDeviceIndex& idx, std::vector& h_vectors, int64_t n_rows, - bool is_cosine) { + bool is_cosine, + const float* stored_inv_norms = nullptr) { int64_t dim = idx.dim; - if (is_cosine) + // Flat cosine (stored_inv_norms == nullptr): the stored vectors are the + // exact originals, so normalizing them in place yields cosine via plain + // inner product. Lossy-SQ cosine decoded to fp32 (stored_inv_norms != null): + // keep the decoded vectors un-normalized and apply the CPU index's original + // inverse norms at search time, matching CPU semantics for lossy codes. + if (is_cosine && stored_inv_norms == nullptr) normalize_vectors(h_vectors, n_rows, dim); size_t dataset_bytes = static_cast(n_rows) * dim * sizeof(float); @@ -232,21 +256,24 @@ inline void upload_fp32_dataset( dataset_bytes, cudaMemcpyHostToDevice)); idx.dataset_type = GpuHnswDatasetType::FP32; + + if (is_cosine && stored_inv_norms != nullptr) + upload_inv_norms(idx, stored_inv_norms, n_rows); } // Upload fp16 (QT_fp16) or bf16 (QT_bf16) ScalarQuantizer codes to the GPU in // their native 2-byte layout. faiss stores these codes row-major as raw IEEE // half / bfloat16, which are bit-compatible with CUDA half / __nv_bfloat16, so -// the bytes are copied verbatim (no up-conversion to fp32). For cosine, inverse -// L2 norms are computed from the fp32-decoded values and applied at search time -// — mirroring the int8 path, since the stored codes are not normalized. +// the bytes are copied verbatim (no up-conversion to fp32). For cosine, the CPU +// index's original inverse L2 norms are applied at search time — mirroring the +// int8 path, since the stored codes are not normalized. inline void upload_halfwidth_dataset( GpuHnswDeviceIndex& idx, const uint8_t* codes, - const float* decoded_for_norms, int64_t n_rows, bool is_cosine, - GpuHnswDatasetType dtype) { + GpuHnswDatasetType dtype, + const float* stored_inv_norms) { int64_t dim = idx.dim; size_t dataset_bytes = static_cast(n_rows) * dim * 2; GPU_HNSW_BUILD_CUDA_CHECK(cudaMalloc(&idx.d_dataset, dataset_bytes)); @@ -254,32 +281,19 @@ inline void upload_halfwidth_dataset( idx.d_dataset, codes, dataset_bytes, cudaMemcpyHostToDevice)); idx.dataset_type = dtype; - if (is_cosine && decoded_for_norms) { - std::vector h_inv_norms(n_rows); - for (int64_t i = 0; i < n_rows; i++) { - const float* row = decoded_for_norms + i * dim; - float sq_norm = 0.0f; - for (int64_t d = 0; d < dim; d++) { - sq_norm += row[d] * row[d]; - } - h_inv_norms[i] = - (sq_norm > 0.0f) ? (1.0f / std::sqrt(sq_norm)) : 0.0f; - } - size_t norms_bytes = static_cast(n_rows) * sizeof(float); - GPU_HNSW_BUILD_CUDA_CHECK(cudaMalloc(&idx.d_inv_norms, norms_bytes)); - GPU_HNSW_BUILD_CUDA_CHECK(cudaMemcpy( - idx.d_inv_norms, - h_inv_norms.data(), - norms_bytes, - cudaMemcpyHostToDevice)); - } + // The stored fp16/bf16 codes are not normalized, so cosine needs the CPU + // index's original inverse norms (not norms recomputed from the lossy + // decoded values). Mirrors the int8 path. + if (is_cosine) + upload_inv_norms(idx, stored_inv_norms, n_rows); } inline void upload_int8_dataset( GpuHnswDeviceIndex& idx, const uint8_t* codes, int64_t n_rows, - bool is_cosine) { + bool is_cosine, + const float* stored_inv_norms) { int64_t dim = idx.dim; size_t dataset_bytes = static_cast(n_rows) * dim; @@ -297,27 +311,12 @@ inline void upload_int8_dataset( cudaMemcpyHostToDevice)); idx.dataset_type = GpuHnswDatasetType::INT8; - if (is_cosine) { - std::vector h_inv_norms(n_rows); - for (int64_t i = 0; i < n_rows; i++) { - const int8_t* row = signed_codes.data() + i * dim; - float sq_norm = 0.0f; - for (int64_t d = 0; d < dim; d++) { - float v = static_cast(row[d]); - sq_norm += v * v; - } - h_inv_norms[i] = - (sq_norm > 0.0f) ? (1.0f / std::sqrt(sq_norm)) : 0.0f; - } - size_t norms_bytes = static_cast(n_rows) * sizeof(float); - GPU_HNSW_BUILD_CUDA_CHECK( - cudaMalloc(&idx.d_inv_norms, norms_bytes)); - GPU_HNSW_BUILD_CUDA_CHECK(cudaMemcpy( - idx.d_inv_norms, - h_inv_norms.data(), - norms_bytes, - cudaMemcpyHostToDevice)); - } + // Apply the CPU index's original inverse norms for cosine. For + // QT_8bit_direct_signed the codes are the original data (lossless), so these + // match a recompute; using the stored norms keeps all cosine paths uniform + // and bit-exact with the CPU index. + if (is_cosine) + upload_inv_norms(idx, stored_inv_norms, n_rows); } /// Build from Knowhere's HNSW index with SQ storage. @@ -345,35 +344,51 @@ inline std::unique_ptr from_faiss_hnsw_sq( auto qtype = sq_storage->sq.qtype; + // For cosine, use the CPU index's inverse L2 norms computed from the + // *original* input vectors (HasInverseL2Norms::get_inverse_l2_norms()). + // Recomputing norms from lossily-decoded/quantized codes would diverge from + // the CPU index's scores and graph traversal for lossy SQ. + const float* stored_inv_norms = nullptr; + if (is_cosine) { + // The cosine norms live on the SQ storage (IndexScalarQuantizerCosine), + // which implements HasInverseL2Norms — not on the outer IndexHNSW. + const auto* cos = dynamic_cast< + const faiss::cppcontrib::knowhere::HasInverseL2Norms*>( + sq_storage); + if (!cos || !cos->get_inverse_l2_norms()) + throw std::runtime_error( + "gpu_hnsw: cosine SQ index missing inverse L2 norms"); + stored_inv_norms = cos->get_inverse_l2_norms(); + } + if (qtype == faiss::ScalarQuantizer::QT_8bit_direct_signed) { - upload_int8_dataset(*idx, sq_storage->codes.data(), n_rows, is_cosine); + upload_int8_dataset( + *idx, sq_storage->codes.data(), n_rows, is_cosine, + stored_inv_norms); } else if ( qtype == faiss::ScalarQuantizer::QT_fp16 || qtype == faiss::ScalarQuantizer::QT_bf16) { - // Keep fp16/bf16 in their native 2-byte layout on the GPU. Decode to - // fp32 only when cosine needs the row norms. + // Keep fp16/bf16 in their native 2-byte layout on the GPU. GpuHnswDatasetType dtype = (qtype == faiss::ScalarQuantizer::QT_fp16) ? GpuHnswDatasetType::FP16 : GpuHnswDatasetType::BF16; - std::vector decoded; - if (is_cosine) { - decoded.resize(static_cast(n_rows) * dim); - sq_storage->sa_decode( - n_rows, sq_storage->codes.data(), decoded.data()); - } upload_halfwidth_dataset( *idx, sq_storage->codes.data(), - is_cosine ? decoded.data() : nullptr, n_rows, is_cosine, - dtype); + dtype, + stored_inv_norms); } else { + // Other SQ types (e.g. unsigned QT_8bit / SQ8) are decoded to fp32 and + // uploaded un-normalized; the stored original inverse norms are applied + // in the kernel, matching the CPU cosine computer for lossy codes. std::vector h_vectors(n_rows * dim); sq_storage->sa_decode( n_rows, sq_storage->codes.data(), h_vectors.data()); - upload_fp32_dataset(*idx, h_vectors, n_rows, is_cosine); + upload_fp32_dataset( + *idx, h_vectors, n_rows, is_cosine, stored_inv_norms); } upload_graph_to_gpu(*idx, hnsw_index.hnsw, n_rows); From 5d201532c816cbefb87bdea8bc577f17764e20d2 Mon Sep 17 00:00:00 2001 From: premal Date: Sun, 12 Jul 2026 01:36:38 +0000 Subject: [PATCH 13/24] fix(gpu_hnsw): guard gpu_index_ with mutex instead of atomic std::atomic> requires C++20 / libstdc++-12; the build toolchain is GCC 11 (libstdc++-11), where it fails the trivially-copyable static_assert. Keep the same lifetime guarantee with a plain shared_ptr guarded by the existing gpu_mutex_: Search() copies it into a local snapshot under a short lock (not held during the search), and Deserialize() already mutates it under the same lock. Reload still cannot free an in-use GPU index. Co-Authored-By: Claude Opus 4.6 Signed-off-by: premal Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- src/index/hnsw/faiss_hnsw.cc | 46 +++++++++++++++++++++--------------- 1 file changed, 27 insertions(+), 19 deletions(-) diff --git a/src/index/hnsw/faiss_hnsw.cc b/src/index/hnsw/faiss_hnsw.cc index 6cdb2cd8d..b3d8c4b06 100644 --- a/src/index/hnsw/faiss_hnsw.cc +++ b/src/index/hnsw/faiss_hnsw.cc @@ -3438,7 +3438,7 @@ class GpuHnswIndexNode : public BaseFaissRegularIndexHNSWNode { // lock-free reader in Search() dereference a null index. Clear it under // the same lock so a failed re-upload leaves the node honestly "not ready". gpu_ready_.store(false, std::memory_order_release); - gpu_index_.store(nullptr, std::memory_order_release); + gpu_index_ = nullptr; // Accept CPU-built HNSW (F32) or HNSW_SQ (quantized) binaries. Status status; @@ -3483,9 +3483,10 @@ class GpuHnswIndexNode : public BaseFaissRegularIndexHNSWNode { gpu_dim_ = faiss_idx->d; // Release CPU copy — vectors and graph are now on GPU. indexes[0].reset(); - // Publish the fully-built index (release); Search() takes an - // acquire snapshot. gpu_ready_ is published last for Count()/Dim(). - gpu_index_.store(local, std::memory_order_release); + // Publish the fully-built index under gpu_mutex_ (held for this + // whole method); Search() takes a locked snapshot. gpu_ready_ is + // published last for the lock-free Count()/Dim() readers. + gpu_index_ = local; gpu_ready_.store(true, std::memory_order_release); } catch (const std::exception& e) { // Fail the load rather than reporting success with no GPU index: @@ -3493,7 +3494,7 @@ class GpuHnswIndexNode : public BaseFaissRegularIndexHNSWNode { // first search. Leave the member null so gpu_ready_ stays false // and return an error so the caller treats the load as failed. LOG_KNOWHERE_ERROR_ << "eager GPU upload failed: " << e.what(); - gpu_index_.store(nullptr, std::memory_order_release); + gpu_index_ = nullptr; return Status::cuda_runtime_error; } } @@ -3511,14 +3512,20 @@ class GpuHnswIndexNode : public BaseFaissRegularIndexHNSWNode { return expected::Err(Status::invalid_args, "GPU_HNSW does not support filtered search"); } - // Take a shared-ownership snapshot of the GPU index up front. It keeps - // the index alive for the whole search even if a concurrent - // Deserialize() reload resets/rebuilds gpu_index_; a bare atomic - // readiness flag would guarantee visibility but not lifetime. - std::shared_ptr gpu_snapshot = gpu_index_.load(std::memory_order_acquire); + // Take a shared-ownership snapshot of the GPU index under gpu_mutex_. The + // snapshot keeps the index alive for the whole search even if a concurrent + // Deserialize() reload resets/rebuilds gpu_index_; the lock is held only + // for the pointer copy, not the search itself. (std::atomic + // is C++20/libstdc++-12 only, so guard the plain shared_ptr with the + // existing mutex instead.) + std::shared_ptr gpu_snapshot; + { + std::lock_guard lock(gpu_mutex_); + gpu_snapshot = gpu_index_; + } if (!gpu_snapshot) { - std::unique_lock lock(gpu_mutex_); - gpu_snapshot = gpu_index_.load(std::memory_order_relaxed); + std::unique_lock lock(gpu_mutex_); + gpu_snapshot = gpu_index_; if (!gpu_snapshot) { const auto* faiss_idx = GetFaissHnswIndex(); if (!faiss_idx) { @@ -3540,7 +3547,7 @@ class GpuHnswIndexNode : public BaseFaissRegularIndexHNSWNode { gpu_ntotal_ = faiss_idx->ntotal; gpu_dim_ = faiss_idx->d; const_cast&>(indexes[0]).reset(); - gpu_index_.store(local, std::memory_order_release); + gpu_index_ = local; gpu_ready_.store(true, std::memory_order_release); gpu_snapshot = local; } catch (const std::exception& e) { @@ -3609,12 +3616,13 @@ class GpuHnswIndexNode : public BaseFaissRegularIndexHNSWNode { mutable std::mutex gpu_mutex_; mutable std::shared_ptr gpu_resources_; - // Held as an atomic shared_ptr so Search() can take a shared-ownership - // snapshot that keeps the GPU index alive for the entire search even if a - // concurrent Deserialize() reload resets/rebuilds the member. Atomic - // readiness (gpu_ready_) only guarantees visibility, not lifetime, so a - // bare pointer read on the search fast path would be a use-after-free. - mutable std::atomic> gpu_index_; + // Guarded by gpu_mutex_. Search() copies this into a local shared_ptr under + // the lock so the GPU index stays alive for the entire (unlocked) search even + // if a concurrent Deserialize() reload resets/rebuilds the member. A bare + // pointer read on the search fast path would be a use-after-free, and + // std::atomic is C++20/libstdc++-12 only (the build + // toolchain is GCC 11), so the mutex-guarded snapshot is used instead. + mutable std::shared_ptr gpu_index_; // Published (release) after gpu_index_ is fully built; read lock-free // (acquire) by Count()/Dim()/HasRawData()/GetVectorByIds(). mutable std::atomic gpu_ready_{false}; From 22fd143b360a0b37bc89a2128c496b95abfb3914 Mon Sep 17 00:00:00 2001 From: premal Date: Sun, 12 Jul 2026 02:43:23 +0000 Subject: [PATCH 14/24] test(gpu_hnsw): raise GPU HNSW cosine recall floor 0.65 -> 0.80 L40S measures GPU HNSW cosine recall ~0.98 (comparable to L2), so the 0.65 floor was far too loose to catch a real regression. Align it with the L2/IP/TopK sections (0.85-ish) at 0.80. Co-Authored-By: Claude Opus 4.6 Signed-off-by: premal Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/ut/test_gpu_search.cc | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/ut/test_gpu_search.cc b/tests/ut/test_gpu_search.cc index c5be0b903..b0beeeb02 100644 --- a/tests/ut/test_gpu_search.cc +++ b/tests/ut/test_gpu_search.cc @@ -526,7 +526,9 @@ TEST_CASE("Test All GPU Index", "[search]") { auto gt = knowhere::BruteForce::Search(train_ds, query_ds, hnsw_json, nullptr); REQUIRE(gt.has_value()); float recall = GetKNNRecall(*gt.value(), *results.value()); - REQUIRE(recall >= 0.65f); + // GPU HNSW cosine recall tracks L2 (measured ~0.98 on L40S); keep the + // floor in line with the L2/IP sections so real regressions are caught. + REQUIRE(recall >= 0.80f); } SECTION("Test GPU HNSW Search TopK") { From 073344acce6ecfcb56df7d6375221bafdd5631f2 Mon Sep 17 00:00:00 2001 From: premal Date: Sun, 12 Jul 2026 02:54:38 +0000 Subject: [PATCH 15/24] test(gpu_hnsw): CUDA upload fault-injection during Deserialize Add a test-only GpuHnswUploadFaultInjection hook (host-safe, in GpuHnswTypes.h) that arms the Nth wrapped CUDA call in the device-upload path to simulate a failure. New [gpu_hnsw_p1] section forces the eager upload to fail both immediately and mid-upload, and asserts: load returns an error, no half-published index (a re-armed search errors cleanly, no crash), VRAM returns to baseline after teardown (no leak), and a fault-free retry loads + searches correctly. Vendored faiss copy mirrors the faiss change byte-identically. Co-Authored-By: Claude Opus 4.6 Signed-off-by: premal Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/ut/test_gpu_search.cc | 72 +++++++++++++++++++ .../faiss/faiss/gpu/impl/GpuHnswBuild.cuh | 26 ++++--- .../faiss/faiss/gpu/impl/GpuHnswTypes.h | 31 ++++++++ 3 files changed, 120 insertions(+), 9 deletions(-) diff --git a/tests/ut/test_gpu_search.cc b/tests/ut/test_gpu_search.cc index b0beeeb02..5e4191058 100644 --- a/tests/ut/test_gpu_search.cc +++ b/tests/ut/test_gpu_search.cc @@ -27,6 +27,10 @@ #ifdef KNOWHERE_WITH_CUVS +// Host-safe header (no device kernels) exposing GpuHnswUploadFaultInjection, the +// test-only hook used by the CUDA-upload fault-injection section below. +#include + template void check_search(const int64_t nb, const int64_t nq, const int64_t dim, const int64_t seed, std::string name, int version, @@ -1252,5 +1256,73 @@ TEST_CASE("Test GPU HNSW Codex P1 Regressions", "[gpu_hnsw_p1]") { REQUIRE(cpu_res.value().memoryCost > 0); REQUIRE(cpu_res.value().maxMemoryCost == 0); } + + // Codex coverage #5: CUDA-upload fault injection during Deserialize(). Using + // the test-only GpuHnswUploadFaultInjection hook, force the eager GPU upload + // to fail both at the very first wrapped CUDA call (immediate) and mid-upload + // (after several allocations have already succeeded, exercising the device + // index destructor's partial-allocation cleanup). Assert: (1) the load + // reports failure, (2) nothing is half-published (a search while re-armed + // errors cleanly rather than using a partial index or crashing), (3) VRAM + // returns to baseline after the failed loads + node teardown (no leak), and + // (4) a subsequent fault-free load succeeds and searches correctly. + SECTION("CUDA upload fault injection: load fails, nothing half-published, retry succeeds") { + auto json = sq_json(knowhere::metric::L2, "fp16", 10); + json.erase(knowhere::indexparam::SQ_TYPE); // plain fp32 HNSW + auto train_ds = GenDataSet(nb, dim, seed); + auto query_ds = GenDataSet(nq, dim, seed + 3); + + auto cpu_idx = + knowhere::IndexFactory::Instance().Create(knowhere::IndexEnum::INDEX_HNSW, version).value(); + REQUIRE(cpu_idx.Build(train_ds, json) == knowhere::Status::success); + knowhere::BinarySet bs; + cpu_idx.Serialize(bs); + + size_t free_before = 0, total = 0; + REQUIRE(cudaMemGetInfo(&free_before, &total) == cudaSuccess); + + // fail_at = 1 fails the first cudaMalloc (immediate); fail_at = 4 fails + // after several buffers are already on the device (mid-upload cleanup). + for (int fail_at : {1, 4}) { + CAPTURE(fail_at); + auto gpu_idx = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_GPU_HNSW, version) + .value(); + + faiss::gpu::GpuHnswUploadFaultInjection::arm(fail_at); + auto st = gpu_idx.Deserialize(bs); + faiss::gpu::GpuHnswUploadFaultInjection::disarm(); + // (1) load fails rather than reporting a silently-broken segment. + REQUIRE(st != knowhere::Status::success); + + // (2) no half-published index: re-arm so the lazy-upload retry inside + // Search() also faults, proving the failed load left no usable index + // behind and the search path errors cleanly (no crash / no results). + faiss::gpu::GpuHnswUploadFaultInjection::arm(fail_at); + auto r = gpu_idx.Search(query_ds, json, nullptr); + faiss::gpu::GpuHnswUploadFaultInjection::disarm(); + REQUIRE(!r.has_value()); + } + + // (3) no leak: after the failed loads and node teardown, free VRAM must + // return to ~baseline (allow slack for allocator/context fragmentation). + REQUIRE(cudaDeviceSynchronize() == cudaSuccess); + size_t free_after = 0; + REQUIRE(cudaMemGetInfo(&free_after, &total) == cudaSuccess); + const size_t slack = static_cast(64) * 1024 * 1024; // 64 MiB + REQUIRE(free_after + slack >= free_before); + + // (4) retry with the fault disarmed succeeds and returns valid results. + faiss::gpu::GpuHnswUploadFaultInjection::disarm(); + auto gpu_ok = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_GPU_HNSW, version) + .value(); + REQUIRE(gpu_ok.Deserialize(bs) == knowhere::Status::success); + auto r_ok = gpu_ok.Search(query_ds, json, nullptr); + REQUIRE(r_ok.has_value()); + auto gt = knowhere::BruteForce::Search(train_ds, query_ds, json, nullptr); + REQUIRE(gt.has_value()); + REQUIRE(GetKNNRecall(*gt.value(), *r_ok.value()) >= 0.8f); + } } #endif diff --git a/thirdparty/faiss/faiss/gpu/impl/GpuHnswBuild.cuh b/thirdparty/faiss/faiss/gpu/impl/GpuHnswBuild.cuh index 29e121e8f..598daab74 100644 --- a/thirdparty/faiss/faiss/gpu/impl/GpuHnswBuild.cuh +++ b/thirdparty/faiss/faiss/gpu/impl/GpuHnswBuild.cuh @@ -46,15 +46,23 @@ #include #include -#define GPU_HNSW_BUILD_CUDA_CHECK(expr) \ - do { \ - cudaError_t _e = (expr); \ - if (_e != cudaSuccess) { \ - throw std::runtime_error( \ - std::string("CUDA error: ") + \ - cudaGetErrorString(_e) + " at " + __FILE__ + ":" + \ - std::to_string(__LINE__)); \ - } \ +// GpuHnswUploadFaultInjection (used below) lives in GpuHnswTypes.h so it is +// reachable from host-compiled unit tests without pulling in device kernels. +#define GPU_HNSW_BUILD_CUDA_CHECK(expr) \ + do { \ + if (faiss::gpu::GpuHnswUploadFaultInjection::should_fail()) { \ + throw std::runtime_error( \ + std::string("CUDA error (injected): simulated ") + \ + "upload failure at " + __FILE__ + ":" + \ + std::to_string(__LINE__)); \ + } \ + cudaError_t _e = (expr); \ + if (_e != cudaSuccess) { \ + throw std::runtime_error( \ + std::string("CUDA error: ") + \ + cudaGetErrorString(_e) + " at " + __FILE__ + ":" + \ + std::to_string(__LINE__)); \ + } \ } while (0) namespace faiss { diff --git a/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.h b/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.h index 25aa68f80..98a001c03 100644 --- a/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.h +++ b/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.h @@ -25,6 +25,7 @@ #include +#include #include #include #include @@ -34,6 +35,36 @@ namespace faiss { namespace gpu { +// Test-only fault injection for the device-upload path (consulted by +// GPU_HNSW_BUILD_CUDA_CHECK in GpuHnswBuild.cuh). Production code never arms it +// (the countdown stays 0), so the check is a single relaxed atomic load per +// wrapped CUDA call with no runtime effect. Unit tests call arm(n) so the n-th +// subsequent wrapped CUDA call reports a simulated failure, exercising +// Deserialize()'s upload error / partial-allocation cleanup / retry path +// without a real OOM. Defined here (host-safe header) rather than in the .cuh so +// host-compiled tests can arm it without pulling in device kernels. The static +// lives in an inline function, so all translation units share one instance. +struct GpuHnswUploadFaultInjection { + static std::atomic& countdown() { + static std::atomic c{0}; + return c; + } + // Arm so the n-th following wrapped CUDA call fails (n >= 1). 0 disarms. + static void arm(int n) { + countdown().store(n, std::memory_order_relaxed); + } + static void disarm() { + countdown().store(0, std::memory_order_relaxed); + } + // Returns true exactly once, on the n-th call after arm(n); self-disarms. + static bool should_fail() { + int cur = countdown().load(std::memory_order_relaxed); + if (cur <= 0) + return false; + return countdown().fetch_sub(1, std::memory_order_relaxed) == 1; + } +}; + // Element type of the device-resident dataset. The graph walk kernel is // templated on this; each value selects a load_elem specialization so the // vectors stay in their native precision on the GPU (no up-conversion to From bac5b7046eaa851530d0c320392e3cd388113192 Mon Sep 17 00:00:00 2001 From: premal Date: Sun, 12 Jul 2026 03:27:24 +0000 Subject: [PATCH 16/24] refactor(gpu_hnsw): drop redundant h_node_ids copy (vendored faiss sync) Mirrors faiss byte-identically: node_ids is read-only after construction, so cudaMemcpy reads it directly. Behavior-neutral. Co-Authored-By: Claude Opus 4.6 Signed-off-by: premal Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- thirdparty/faiss/faiss/gpu/impl/GpuHnswBuild.cuh | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/thirdparty/faiss/faiss/gpu/impl/GpuHnswBuild.cuh b/thirdparty/faiss/faiss/gpu/impl/GpuHnswBuild.cuh index 598daab74..f051bcd73 100644 --- a/thirdparty/faiss/faiss/gpu/impl/GpuHnswBuild.cuh +++ b/thirdparty/faiss/faiss/gpu/impl/GpuHnswBuild.cuh @@ -127,7 +127,6 @@ inline void extract_hnsw_layers( ul.num_nodes = static_cast(node_ids.size()); std::vector h_neighbors(ul.num_nodes * maxM, UINT32_MAX); - std::vector h_node_ids = node_ids; for (uint32_t idx = 0; idx < ul.num_nodes; idx++) { int64_t i = node_ids[idx]; @@ -145,7 +144,7 @@ inline void extract_hnsw_layers( cudaMalloc(&ul.d_node_ids, ul.num_nodes * sizeof(uint32_t))); GPU_HNSW_BUILD_CUDA_CHECK(cudaMemcpy( ul.d_node_ids, - h_node_ids.data(), + node_ids.data(), ul.num_nodes * sizeof(uint32_t), cudaMemcpyHostToDevice)); From 6906c993fa08811edd244ad34b6a08bbe2c2ede1 Mon Sep 17 00:00:00 2001 From: premal Date: Sun, 12 Jul 2026 08:11:44 +0000 Subject: [PATCH 17/24] fix(gpu_hnsw): run GPU upload on the DeserializeFromFile (mmap) load path Milvus loads sealed vector-index segments through VectorMemIndex::LoadFromFile -> index_.DeserializeFromFile(), not Deserialize(BinarySet). GpuHnswIndexNode only overrode Deserialize(BinarySet), so the file/mmap load path fell through to the base CPU loader and never ran the eager GPU upload: the CPU-built HNSW stayed resident in host RAM (fp32-expanded for int8/fp16/bf16 inputs, ~5.9 GiB/segment for the mpd_v2 collection) while VRAM sat near-idle, OOMKilling querynodes during the load storm before any search could trigger the lazy-upload fallback. Override DeserializeFromFile to load the CPU index via the base and then run the same eager upload + host-copy release. The shared upload/unpublish logic is factored into UploadCpuIndexToGpuLocked()/UnpublishGpuIndexLocked() so both entrypoints behave identically; GpuHnswSQIndexNode inherits the fix. Adds a [gpu_hnsw_p1] section that loads via DeserializeFromFile (enable_mmap true/false) and asserts Size()==0 before any search (i.e. indexes[0] freed), which fails pre-fix and cannot be masked by the first-search lazy upload. Signed-off-by: premal Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- src/index/hnsw/faiss_hnsw.cc | 127 ++++++++++++++++++++++------------- tests/ut/test_gpu_search.cc | 67 ++++++++++++++++++ 2 files changed, 148 insertions(+), 46 deletions(-) diff --git a/src/index/hnsw/faiss_hnsw.cc b/src/index/hnsw/faiss_hnsw.cc index b3d8c4b06..30167929e 100644 --- a/src/index/hnsw/faiss_hnsw.cc +++ b/src/index/hnsw/faiss_hnsw.cc @@ -3433,12 +3433,7 @@ class GpuHnswIndexNode : public BaseFaissRegularIndexHNSWNode { Status Deserialize(const BinarySet& binset, std::shared_ptr cfg) override { std::unique_lock lock(gpu_mutex_); - // Unpublish before tearing down: on a reload a stale gpu_ready_==true - // combined with a reset (or failed re-upload of) gpu_index_ would let the - // lock-free reader in Search() dereference a null index. Clear it under - // the same lock so a failed re-upload leaves the node honestly "not ready". - gpu_ready_.store(false, std::memory_order_release); - gpu_index_ = nullptr; + UnpublishGpuIndexLocked(); // Accept CPU-built HNSW (F32) or HNSW_SQ (quantized) binaries. Status status; @@ -3458,47 +3453,26 @@ class GpuHnswIndexNode : public BaseFaissRegularIndexHNSWNode { return status; } - // Eager GPU upload via faiss::gpu::GpuIndexHNSW. - const auto* faiss_idx = GetFaissHnswIndex(); - if (faiss_idx) { - try { - // Detect metric from the FAISS index type rather than config, - // because Deserialize may be called without metric_type in the config - // (e.g. empty json defaults metric_type to L2). - bool is_cosine = - dynamic_cast(faiss_idx) != nullptr; - bool use_ip = is_cosine || (faiss_idx->metric_type == ::faiss::METRIC_INNER_PRODUCT); - - std::shared_ptr local; - { - std::lock_guard gpu_ctor_lock(GetGpuConstructionMutex()); - gpu_resources_ = GetSharedGpuResources(); - local = std::make_shared(gpu_resources_.get(), faiss_idx->d, - faiss_idx->metric_type); - } - local->copyFromWithMetric(faiss_idx, use_ip, is_cosine); - // Capture count/dim before releasing the CPU copy so Count()/ - // Dim() keep working once indexes[0] is freed. - gpu_ntotal_ = faiss_idx->ntotal; - gpu_dim_ = faiss_idx->d; - // Release CPU copy — vectors and graph are now on GPU. - indexes[0].reset(); - // Publish the fully-built index under gpu_mutex_ (held for this - // whole method); Search() takes a locked snapshot. gpu_ready_ is - // published last for the lock-free Count()/Dim() readers. - gpu_index_ = local; - gpu_ready_.store(true, std::memory_order_release); - } catch (const std::exception& e) { - // Fail the load rather than reporting success with no GPU index: - // a silently "loaded" segment would only surface the error on the - // first search. Leave the member null so gpu_ready_ stays false - // and return an error so the caller treats the load as failed. - LOG_KNOWHERE_ERROR_ << "eager GPU upload failed: " << e.what(); - gpu_index_ = nullptr; - return Status::cuda_runtime_error; - } + return UploadCpuIndexToGpuLocked(); + } + + // Milvus loads sealed segments through the mmap/file path + // (VectorMemIndex::LoadFromFile -> DeserializeFromFile), NOT Deserialize(). + // The base only materializes the CPU index, so without this override the + // eager GPU upload would never run on that path: the CPU-built HNSW would + // stay resident in host RAM (fp32-expanded for quantized inputs) and the GPU + // would sit unused. Load the CPU index via the base, then upload + free it. + Status + DeserializeFromFile(const std::string& filename, std::shared_ptr config) override { + std::unique_lock lock(gpu_mutex_); + UnpublishGpuIndexLocked(); + + auto status = BaseFaissRegularIndexHNSWNode::DeserializeFromFile(filename, config); + if (status != Status::success) { + return status; } - return Status::success; + + return UploadCpuIndexToGpuLocked(); } expected @@ -3607,6 +3581,67 @@ class GpuHnswIndexNode : public BaseFaissRegularIndexHNSWNode { ~GpuHnswIndexNode() override = default; private: + // Requires gpu_mutex_ held. Unpublish before tearing down: on a reload a + // stale gpu_ready_==true combined with a reset (or failed re-upload of) + // gpu_index_ would let the lock-free reader in Search() dereference a null + // index. Clearing it here leaves the node honestly "not ready" if the + // subsequent (re-)upload fails. + void + UnpublishGpuIndexLocked() { + gpu_ready_.store(false, std::memory_order_release); + gpu_index_ = nullptr; + } + + // Requires gpu_mutex_ held. Uploads the just-deserialized CPU HNSW to the GPU + // via faiss::gpu::GpuIndexHNSW, releases the host copy, and publishes the + // device index. Shared by Deserialize() (BinarySet) and DeserializeFromFile() + // (Milvus mmap/file load path) so both entrypoints upload identically. A + // no-op success when there is no CPU index to upload (empty segment). + Status + UploadCpuIndexToGpuLocked() { + const auto* faiss_idx = GetFaissHnswIndex(); + if (faiss_idx == nullptr) { + return Status::success; + } + try { + // Detect metric from the FAISS index type rather than config, because + // deserialize may be called without metric_type in the config (e.g. + // an empty json defaults metric_type to L2). + bool is_cosine = + dynamic_cast(faiss_idx) != nullptr; + bool use_ip = is_cosine || (faiss_idx->metric_type == ::faiss::METRIC_INNER_PRODUCT); + + std::shared_ptr local; + { + std::lock_guard gpu_ctor_lock(GetGpuConstructionMutex()); + gpu_resources_ = GetSharedGpuResources(); + local = std::make_shared(gpu_resources_.get(), faiss_idx->d, + faiss_idx->metric_type); + } + local->copyFromWithMetric(faiss_idx, use_ip, is_cosine); + // Capture count/dim before releasing the CPU copy so Count()/Dim() + // keep working once indexes[0] is freed. + gpu_ntotal_ = faiss_idx->ntotal; + gpu_dim_ = faiss_idx->d; + // Release CPU copy — vectors and graph are now on GPU. This is what + // drops the transient host-RAM peak back to ~0. + indexes[0].reset(); + // Publish last: Search() takes a locked snapshot of gpu_index_, and + // gpu_ready_ is released for the lock-free Count()/Dim() readers. + gpu_index_ = local; + gpu_ready_.store(true, std::memory_order_release); + } catch (const std::exception& e) { + // Fail the load rather than reporting success with no GPU index: a + // silently "loaded" segment would only surface the error on the first + // search. Leave the member null so gpu_ready_ stays false and return + // an error so the caller treats the load as failed. + LOG_KNOWHERE_ERROR_ << "eager GPU upload failed: " << e.what(); + gpu_index_ = nullptr; + return Status::cuda_runtime_error; + } + return Status::success; + } + const ::faiss::cppcontrib::knowhere::IndexHNSW* GetFaissHnswIndex() const { if (indexes.empty() || !indexes[0]) diff --git a/tests/ut/test_gpu_search.cc b/tests/ut/test_gpu_search.cc index 5e4191058..1565fceba 100644 --- a/tests/ut/test_gpu_search.cc +++ b/tests/ut/test_gpu_search.cc @@ -11,6 +11,8 @@ #include #include +#include +#include #include #include #include @@ -1028,6 +1030,71 @@ TEST_CASE("Test GPU HNSW Codex P1 Regressions", "[gpu_hnsw_p1]") { REQUIRE(recall >= (sq_type == "sq8" ? 0.65f : 0.85f)); } + // #4: Milvus loads sealed segments through the mmap/file path + // (VectorMemIndex::LoadFromFile -> DeserializeFromFile), NOT + // Deserialize(BinarySet). The eager GPU upload must fire on that path too; + // otherwise the CPU-built HNSW stays resident in host RAM and the GPU is + // never used. This was the root cause of the mpd_v2 querynode OOM: each + // segment's (fp32-expanded) CPU index — ~5.9 GiB — was retained in host + // anon memory instead of being uploaded to VRAM and freed, so ~14 + // segments/pod overran the memory limit while VRAM stayed near-idle. + // + // The discriminating assertion is Size()==0 *before any search*: the eager + // upload releases indexes[0], so the base Size() (computed from indexes[0]) + // is 0. Pre-fix, DeserializeFromFile skipped the upload and left the CPU + // copy resident, so Size() reported the full retained host footprint. A + // recall-only check would NOT catch the regression, because the first + // Search() lazily uploads and frees the copy — masking the load-time leak + // that OOMs before any search runs. + SECTION("DeserializeFromFile uploads to GPU and frees the host copy") { + const bool enable_mmap = GENERATE(false, true); + CAPTURE(enable_mmap); + auto json = sq_json(knowhere::metric::COSINE, "sq8", 10); + + auto train_ds = GenDataSet(nb, dim, seed); + auto query_ds = GenDataSet(nq, dim, seed + 2); + + auto cpu_idx = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_HNSW_SQ, version) + .value(); + REQUIRE(cpu_idx.Build(train_ds, json) == knowhere::Status::success); + knowhere::BinarySet bs; + REQUIRE(cpu_idx.Serialize(bs) == knowhere::Status::success); + auto binary = bs.GetByName(cpu_idx.Type()); + REQUIRE(binary != nullptr); + + const auto path = std::filesystem::temp_directory_path() / + ("knowhere_gpu_hnsw_deserialize_from_file_" + std::to_string(enable_mmap) + ".index"); + std::filesystem::remove(path); + { + std::ofstream out(path, std::ios::binary); + REQUIRE(out.is_open()); + out.write(reinterpret_cast(binary->data.get()), binary->size); + } + + auto gpu_idx = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_GPU_HNSW, version) + .value(); + auto load_json = json; + load_json["enable_mmap"] = enable_mmap; + REQUIRE(gpu_idx.DeserializeFromFile(path.string(), load_json) == knowhere::Status::success); + + // Eager upload ran during load (no search yet): host copy released. + REQUIRE(gpu_idx.Size() == 0); + REQUIRE(gpu_idx.Count() == nb); + REQUIRE(gpu_idx.HasRawData(knowhere::metric::COSINE) == false); + + // And the GPU index actually serves searches. + auto results = gpu_idx.Search(query_ds, json, nullptr); + REQUIRE(results.has_value()); + auto gt = knowhere::BruteForce::Search(train_ds, query_ds, json, nullptr); + REQUIRE(gt.has_value()); + float recall = GetKNNRecall(*gt.value(), *results.value()); + REQUIRE(recall >= 0.65f); + + std::filesystem::remove(path); + } + // #1b: quantized-cosine parity. The discriminating input is vectors that // share directions but have widely different magnitudes: cosine must ignore // magnitude. The CPU cosine index records inverse norms from the ORIGINAL From 4d3d97d8f801541ebd3afe6829c47531de13f295 Mon Sep 17 00:00:00 2001 From: premal Date: Sun, 12 Jul 2026 16:15:36 +0000 Subject: [PATCH 18/24] fix(gpu_hnsw): size load estimate to fp32 resident graph; restore mmap eligibility Two coupled fixes for the mpd_v2 host-RAM OOM regression that appeared once the HNSW->GPU_HNSW load override was introduced. 1. StaticEstimateLoadResource previously proxied the resident deserialized index as ~file_size. For a compressed (int8) on-disk index this grossly under-counts the fp32-expanded hnswlib graph materialized in host RAM, so the loader admitted far too many concurrent uploads and OOMed the host cgroup before any GPU upload completed (VRAM stayed at 0). The estimate now sizes the transient peak as file + resident fp32 vectors (rows*dim*4) + graph neighbor lists (rows*M*2*4) + fp32 decode staging (rows*dim*4), reads M from the load config when present, and falls back to a conservative multiple of file_size when row/dim metadata is missing. 2. GPU_HNSW / GPU_HNSW_SQ were not advertised as mmap-capable, so the override HNSW->GPU_HNSW silently flipped enableMmap off and the whole CPU index landed in anonymous host RAM instead of being file-backed as plain HNSW was before the override. Re-add feature::MMAP to the registrations and list the types in legal_support_mmap_knowhere_index; the CPU index is deserialized via the mmap-capable file path and freed after the VRAM upload. Strengthen the [gpu_hnsw_p1] load-resource test to assert the estimate tracks the fp32 resident footprint (not file_size) for a compressed index and does not collapse when num_rows arrives as 0. Co-Authored-By: Claude Opus 4.6 Signed-off-by: premal Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- include/knowhere/index/index_table.h | 6 +++ src/index/hnsw/faiss_hnsw.cc | 69 ++++++++++++++++++++++------ tests/ut/test_gpu_search.cc | 28 ++++++++++- 3 files changed, 89 insertions(+), 14 deletions(-) diff --git a/include/knowhere/index/index_table.h b/include/knowhere/index/index_table.h index 0b300ef66..5c8838511 100644 --- a/include/knowhere/index/index_table.h +++ b/include/knowhere/index/index_table.h @@ -167,6 +167,12 @@ static std::set legal_support_mmap_knowhere_index = { IndexEnum::INDEX_HNSW_PQ, IndexEnum::INDEX_HNSW_PRQ, + // gpu hnsw: the CPU index is deserialized via the mmap-capable file path + // during load (the vectors/graph then move to VRAM and the CPU copy is + // freed), so the transient CPU index can be file-backed rather than anon. + IndexEnum::INDEX_GPU_HNSW, + IndexEnum::INDEX_GPU_HNSW_SQ, + // sparse index IndexEnum::INDEX_SPARSE_INVERTED_INDEX, IndexEnum::INDEX_SPARSE_WAND, diff --git a/src/index/hnsw/faiss_hnsw.cc b/src/index/hnsw/faiss_hnsw.cc index 30167929e..0b7e50365 100644 --- a/src/index/hnsw/faiss_hnsw.cc +++ b/src/index/hnsw/faiss_hnsw.cc @@ -3358,19 +3358,55 @@ class GpuHnswIndexNode : public BaseFaissRegularIndexHNSWNode { StaticEstimateLoadResource(const uint64_t file_size_in_bytes, const int64_t num_rows, const int64_t dim, const knowhere::BaseConfig& config, const IndexVersion& version) { // GPU HNSW stores vectors and graph in VRAM; the CPU copy is freed after - // upload in Deserialize(). Phase-separated accounting: + // upload in Deserialize()/DeserializeFromFile(). Phase-separated accounting: // memoryCost (retained after load): ~0 — the CPU index is released // once the graph/vectors are on the GPU. // maxMemoryCost (transient peak during load): the loader must reserve // enough host RAM to cover the simultaneously-live buffers before the // GPU upload frees the CPU copy: - // - the serialized download buffer (~file_size), - // - the deserialized CPU HNSW index: vectors + graph (~file_size), - // - fp32 decode/graph staging for the upload (num_rows*dim*4). + // - the serialized download buffer (~file_size), + // - the resident fp32 vectors (num_rows*dim*4), + // - the resident hnswlib graph lists (num_rows*M*2*4), + // - fp32 decode staging for the upload (num_rows*dim*4). + // The on-disk file_size is NOT a proxy for the resident index: hnswlib + // materializes vectors as fp32 inline (dim*4 per node) even when the + // stored form is compressed (e.g. int8), so the deserialized index is + // several times the file. copyFromWithMetric() then reconstructs a + // separate fp32 staging buffer before the host->device copy. Under- + // counting this peak lets the loader admit too many concurrent uploads + // and OOMs the host cgroup before any upload completes. const uint64_t rows = num_rows > 0 ? static_cast(num_rows) : 0; const uint64_t d = dim > 0 ? static_cast(dim) : 0; + + // hnsw connectivity: read the train-time M when the load config carries + // it; otherwise fall back to a conservative value (over-reserving the + // graph term is safe for admission control). + uint64_t m = 32; + if (const auto* hnsw_cfg = dynamic_cast(&config)) { + if (hnsw_cfg->M.has_value()) { + m = static_cast(hnsw_cfg->M.value()); + } + } + + const uint64_t resident_vectors = rows * d * sizeof(float); + const uint64_t resident_graph = rows * m * 2 * sizeof(int32_t); const uint64_t decode_staging = rows * d * sizeof(float); - const uint64_t peak = file_size_in_bytes + file_size_in_bytes + decode_staging; + uint64_t peak = file_size_in_bytes + resident_vectors + resident_graph + decode_staging; + + if (rows == 0 || d == 0) { + // Row/dim metadata is unavailable at estimate time (num_rows can + // arrive as 0 from the load info). rows*dim*4 collapses to 0 and the + // estimate degenerates to ~file_size, grossly under-reserving the + // fp32-expanded peak. Guard with a conservative multiple of the file + // size so admission still throttles concurrent GPU uploads. + LOG_KNOWHERE_WARNING_ << "GPU_HNSW load estimate missing row/dim metadata (num_rows=" << num_rows + << ", dim=" << dim << "); falling back to file-size-based peak"; + const uint64_t fallback = file_size_in_bytes * 9; + if (peak < fallback) { + peak = fallback; + } + } + return Resource{.memoryCost = 0, .diskCost = 0, .maxMemoryCost = peak}; } @@ -3695,57 +3731,64 @@ register_gpu_hnsw_static_config() { IndexStaticFaced::Instance().RegisterStaticFunc(IndexEnum::INDEX_GPU_HNSW_SQ); } +// NOTE: feature::MMAP is intentionally set. Although GPU_HNSW ultimately holds +// its vectors/graph in VRAM (the CPU copy is freed after upload), the CPU index +// is deserialized via the mmap-capable file path during load. Advertising MMAP +// lets Milvus keep enableMmap=true so the transient CPU HNSW is file-backed +// rather than anonymous host RAM before the GPU upload frees it -- without this, +// the override HNSW->GPU_HNSW silently disables mmap and the full fp32-expanded +// index lands in anon RAM (the post-override host OOM regression). KNOWHERE_REGISTER_GLOBAL( GPU_HNSW, [](const int32_t& version, const Object& object) { return Index::Create(version, object); }, fp32, - true, feature::GPU_ANN_FLOAT_INDEX); + true, (feature::GPU_ANN_FLOAT_INDEX | feature::MMAP)); KNOWHERE_REGISTER_GLOBAL( GPU_HNSW, [](const int32_t& version, const Object& object) { return Index>::Create(std::make_unique(version, object)); }, - fp16, true, (feature::FP16 | feature::GPU)); + fp16, true, (feature::FP16 | feature::GPU | feature::MMAP)); KNOWHERE_REGISTER_GLOBAL( GPU_HNSW, [](const int32_t& version, const Object& object) { return Index>::Create(std::make_unique(version, object)); }, - bf16, true, (feature::BF16 | feature::GPU)); + bf16, true, (feature::BF16 | feature::GPU | feature::MMAP)); KNOWHERE_REGISTER_GLOBAL( GPU_HNSW, [](const int32_t& version, const Object& object) { return Index>::Create(std::make_unique(version, object)); }, - int8, true, (feature::INT8 | feature::GPU)); + int8, true, (feature::INT8 | feature::GPU | feature::MMAP)); KNOWHERE_REGISTER_GLOBAL( GPU_HNSW_SQ, [](const int32_t& version, const Object& object) { return Index::Create(version, object); }, - fp32, true, feature::GPU_ANN_FLOAT_INDEX); + fp32, true, (feature::GPU_ANN_FLOAT_INDEX | feature::MMAP)); KNOWHERE_REGISTER_GLOBAL( GPU_HNSW_SQ, [](const int32_t& version, const Object& object) { return Index>::Create(std::make_unique(version, object)); }, - fp16, true, (feature::FP16 | feature::GPU)); + fp16, true, (feature::FP16 | feature::GPU | feature::MMAP)); KNOWHERE_REGISTER_GLOBAL( GPU_HNSW_SQ, [](const int32_t& version, const Object& object) { return Index>::Create(std::make_unique(version, object)); }, - bf16, true, (feature::BF16 | feature::GPU)); + bf16, true, (feature::BF16 | feature::GPU | feature::MMAP)); KNOWHERE_REGISTER_GLOBAL( GPU_HNSW_SQ, [](const int32_t& version, const Object& object) { return Index>::Create(std::make_unique(version, object)); }, - int8, true, (feature::INT8 | feature::GPU)); + int8, true, (feature::INT8 | feature::GPU | feature::MMAP)); #endif // KNOWHERE_WITH_CUVS } // namespace knowhere diff --git a/tests/ut/test_gpu_search.cc b/tests/ut/test_gpu_search.cc index 1565fceba..2a334f727 100644 --- a/tests/ut/test_gpu_search.cc +++ b/tests/ut/test_gpu_search.cc @@ -1309,11 +1309,37 @@ TEST_CASE("Test GPU HNSW Codex P1 Regressions", "[gpu_hnsw_p1]") { REQUIRE(gpu_res.value().memoryCost == 0); REQUIRE(gpu_res.value().diskCost == 0); // Transient peak is non-zero and at least covers the fp32 decode staging - // plus the download buffer (2*file_size + rows*dim*4 in the estimator). + // plus the download buffer. REQUIRE(gpu_res.value().maxMemoryCost > 0); const uint64_t decode_staging = static_cast(nb) * dim * sizeof(float); REQUIRE(gpu_res.value().maxMemoryCost >= file_size + decode_staging); + // The estimate must be driven by the fp32-expanded resident index, NOT + // by the on-disk file size. A compressed (e.g. int8) index stores far + // fewer bytes on disk than the fp32 graph hnswlib materializes in host + // RAM, so the peak has to cover num_rows*dim*4 (resident vectors) + + // num_rows*dim*4 (decode staging) independent of file_size. Simulate an + // int8-sized file (1 byte/element) and assert the peak still reflects + // the fp32 in-memory footprint (this is the production OOM regression: + // an ~file_size-tracking estimate under-reserves and OOMs the host). + const uint64_t int8_file = static_cast(nb) * dim * sizeof(int8_t); + auto compressed_res = knowhere::IndexStaticFaced::EstimateLoadResource( + knowhere::IndexEnum::INDEX_GPU_HNSW, version, int8_file, nb, dim, params); + REQUIRE(compressed_res.has_value()); + const uint64_t fp32_resident_plus_staging = 2 * static_cast(nb) * dim * sizeof(float); + REQUIRE(compressed_res.value().maxMemoryCost >= fp32_resident_plus_staging); + // And it must dwarf the naive 2*file_size an ~file-tracking estimate + // would have produced. + REQUIRE(compressed_res.value().maxMemoryCost > 2 * int8_file); + + // When row/dim metadata is missing (num_rows arrives as 0 from the load + // info), the rows*dim terms vanish; the estimate must not collapse to + // ~file_size but fall back to a conservative multiple of it. + auto norows_res = knowhere::IndexStaticFaced::EstimateLoadResource( + knowhere::IndexEnum::INDEX_GPU_HNSW, version, int8_file, 0, dim, params); + REQUIRE(norows_res.has_value()); + REQUIRE(norows_res.value().maxMemoryCost >= int8_file * 9); + // A plain CPU HNSW keeps its data resident: memoryCost is non-zero and it // does not opt into the peak field (maxMemoryCost stays 0 -> Milvus falls // back to its 2*memoryCost heuristic). From 75474cdf543775bc673aac648487299847b63b5d Mon Sep 17 00:00:00 2001 From: premal Date: Sun, 12 Jul 2026 17:42:14 +0000 Subject: [PATCH 19/24] revert(gpu_hnsw): drop MMAP eligibility; native int8 CPU copy is compact The HNSW->GPU_HNSW load deserializes the on-disk faiss index (IndexHNSWSQCosine with QT_8bit_direct_signed for an int8 collection) into a compact host copy, uploads it to VRAM, then frees it. That transient int8 copy is ~1 byte/dim, matching how the CPU HNSW int8 node loads -- so GPU_HNSW does not need to be memory-mapped. Remove feature::MMAP from the GPU_HNSW/GPU_HNSW_SQ registrations and the mmap support table; do not rely on enable_mmap. Signed-off-by: premal Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- include/knowhere/index/index_table.h | 6 ------ src/index/hnsw/faiss_hnsw.cc | 27 ++++++++++++--------------- 2 files changed, 12 insertions(+), 21 deletions(-) diff --git a/include/knowhere/index/index_table.h b/include/knowhere/index/index_table.h index 5c8838511..0b300ef66 100644 --- a/include/knowhere/index/index_table.h +++ b/include/knowhere/index/index_table.h @@ -167,12 +167,6 @@ static std::set legal_support_mmap_knowhere_index = { IndexEnum::INDEX_HNSW_PQ, IndexEnum::INDEX_HNSW_PRQ, - // gpu hnsw: the CPU index is deserialized via the mmap-capable file path - // during load (the vectors/graph then move to VRAM and the CPU copy is - // freed), so the transient CPU index can be file-backed rather than anon. - IndexEnum::INDEX_GPU_HNSW, - IndexEnum::INDEX_GPU_HNSW_SQ, - // sparse index IndexEnum::INDEX_SPARSE_INVERTED_INDEX, IndexEnum::INDEX_SPARSE_WAND, diff --git a/src/index/hnsw/faiss_hnsw.cc b/src/index/hnsw/faiss_hnsw.cc index 0b7e50365..865a5d7fe 100644 --- a/src/index/hnsw/faiss_hnsw.cc +++ b/src/index/hnsw/faiss_hnsw.cc @@ -3731,64 +3731,61 @@ register_gpu_hnsw_static_config() { IndexStaticFaced::Instance().RegisterStaticFunc(IndexEnum::INDEX_GPU_HNSW_SQ); } -// NOTE: feature::MMAP is intentionally set. Although GPU_HNSW ultimately holds -// its vectors/graph in VRAM (the CPU copy is freed after upload), the CPU index -// is deserialized via the mmap-capable file path during load. Advertising MMAP -// lets Milvus keep enableMmap=true so the transient CPU HNSW is file-backed -// rather than anonymous host RAM before the GPU upload frees it -- without this, -// the override HNSW->GPU_HNSW silently disables mmap and the full fp32-expanded -// index lands in anon RAM (the post-override host OOM regression). +// GPU_HNSW is not registered as MMAP-capable: the CPU index is deserialized +// into host RAM, uploaded to VRAM, then the CPU copy is freed. For a native +// int8 index the transient CPU copy is compact (~1 byte/dim), matching how the +// CPU HNSW int8 node loads, so no memory-mapped-file path is required. KNOWHERE_REGISTER_GLOBAL( GPU_HNSW, [](const int32_t& version, const Object& object) { return Index::Create(version, object); }, fp32, - true, (feature::GPU_ANN_FLOAT_INDEX | feature::MMAP)); + true, (feature::GPU_ANN_FLOAT_INDEX)); KNOWHERE_REGISTER_GLOBAL( GPU_HNSW, [](const int32_t& version, const Object& object) { return Index>::Create(std::make_unique(version, object)); }, - fp16, true, (feature::FP16 | feature::GPU | feature::MMAP)); + fp16, true, (feature::FP16 | feature::GPU)); KNOWHERE_REGISTER_GLOBAL( GPU_HNSW, [](const int32_t& version, const Object& object) { return Index>::Create(std::make_unique(version, object)); }, - bf16, true, (feature::BF16 | feature::GPU | feature::MMAP)); + bf16, true, (feature::BF16 | feature::GPU)); KNOWHERE_REGISTER_GLOBAL( GPU_HNSW, [](const int32_t& version, const Object& object) { return Index>::Create(std::make_unique(version, object)); }, - int8, true, (feature::INT8 | feature::GPU | feature::MMAP)); + int8, true, (feature::INT8 | feature::GPU)); KNOWHERE_REGISTER_GLOBAL( GPU_HNSW_SQ, [](const int32_t& version, const Object& object) { return Index::Create(version, object); }, - fp32, true, (feature::GPU_ANN_FLOAT_INDEX | feature::MMAP)); + fp32, true, (feature::GPU_ANN_FLOAT_INDEX)); KNOWHERE_REGISTER_GLOBAL( GPU_HNSW_SQ, [](const int32_t& version, const Object& object) { return Index>::Create(std::make_unique(version, object)); }, - fp16, true, (feature::FP16 | feature::GPU | feature::MMAP)); + fp16, true, (feature::FP16 | feature::GPU)); KNOWHERE_REGISTER_GLOBAL( GPU_HNSW_SQ, [](const int32_t& version, const Object& object) { return Index>::Create(std::make_unique(version, object)); }, - bf16, true, (feature::BF16 | feature::GPU | feature::MMAP)); + bf16, true, (feature::BF16 | feature::GPU)); KNOWHERE_REGISTER_GLOBAL( GPU_HNSW_SQ, [](const int32_t& version, const Object& object) { return Index>::Create(std::make_unique(version, object)); }, - int8, true, (feature::INT8 | feature::GPU | feature::MMAP)); + int8, true, (feature::INT8 | feature::GPU)); #endif // KNOWHERE_WITH_CUVS } // namespace knowhere From 7d8a448adb68ad1634e3e6eba6a0727df090817f Mon Sep 17 00:00:00 2001 From: premal Date: Sun, 12 Jul 2026 18:51:19 +0000 Subject: [PATCH 20/24] fix: GPU_HNSW load estimate to compact int8 transient (~2x file), not fp32 The prior StaticEstimateLoadResource sized the transient host peak as an fp32/hnswlib index (rows*dim*4 resident + rows*dim*4 decode staging + graph), yielding ~20 GiB/seg. mpd_v2's on-disk index is the compact FAISS int8-SQ form (IndexHNSWSQCosine): it deserializes to ~file_size in host RAM and the signed int8 codes are uploaded to the device directly, so there is no fp32 reconstruct/decode staging on this path. The CPU HNSW cluster loads the same index resident at ~file_size (~1.1 GiB/seg measured). Because the estimate is charged into the querynode's committed-loading memory (indexLoadingMemoryContribution), the fp32 overestimate made the soft-OOM admission guard trip after ~3 segments (3x20 > threshold) even though real RSS was only ~16 GiB and freed back to baseline after every upload. Sizing the transient at ~2x file_size (serialized read buffer + deserialized compact index) makes admission reflect the true compact footprint. Retained cost stays 0 (data lives in VRAM after the CPU copy is freed). Co-Authored-By: Claude Opus 4.6 Signed-off-by: premal Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- src/index/hnsw/faiss_hnsw.cc | 75 +++++++++++++----------------------- 1 file changed, 27 insertions(+), 48 deletions(-) diff --git a/src/index/hnsw/faiss_hnsw.cc b/src/index/hnsw/faiss_hnsw.cc index 865a5d7fe..79d3e6b2c 100644 --- a/src/index/hnsw/faiss_hnsw.cc +++ b/src/index/hnsw/faiss_hnsw.cc @@ -3357,55 +3357,34 @@ class GpuHnswIndexNode : public BaseFaissRegularIndexHNSWNode { static expected StaticEstimateLoadResource(const uint64_t file_size_in_bytes, const int64_t num_rows, const int64_t dim, const knowhere::BaseConfig& config, const IndexVersion& version) { - // GPU HNSW stores vectors and graph in VRAM; the CPU copy is freed after - // upload in Deserialize()/DeserializeFromFile(). Phase-separated accounting: + // GPU HNSW stores the vectors and graph in VRAM; the CPU copy is freed + // after upload in Deserialize()/DeserializeFromFile(). Phase-separated + // accounting: // memoryCost (retained after load): ~0 — the CPU index is released - // once the graph/vectors are on the GPU. - // maxMemoryCost (transient peak during load): the loader must reserve - // enough host RAM to cover the simultaneously-live buffers before the - // GPU upload frees the CPU copy: - // - the serialized download buffer (~file_size), - // - the resident fp32 vectors (num_rows*dim*4), - // - the resident hnswlib graph lists (num_rows*M*2*4), - // - fp32 decode staging for the upload (num_rows*dim*4). - // The on-disk file_size is NOT a proxy for the resident index: hnswlib - // materializes vectors as fp32 inline (dim*4 per node) even when the - // stored form is compressed (e.g. int8), so the deserialized index is - // several times the file. copyFromWithMetric() then reconstructs a - // separate fp32 staging buffer before the host->device copy. Under- - // counting this peak lets the loader admit too many concurrent uploads - // and OOMs the host cgroup before any upload completes. - const uint64_t rows = num_rows > 0 ? static_cast(num_rows) : 0; - const uint64_t d = dim > 0 ? static_cast(dim) : 0; - - // hnsw connectivity: read the train-time M when the load config carries - // it; otherwise fall back to a conservative value (over-reserving the - // graph term is safe for admission control). - uint64_t m = 32; - if (const auto* hnsw_cfg = dynamic_cast(&config)) { - if (hnsw_cfg->M.has_value()) { - m = static_cast(hnsw_cfg->M.value()); - } - } - - const uint64_t resident_vectors = rows * d * sizeof(float); - const uint64_t resident_graph = rows * m * 2 * sizeof(int32_t); - const uint64_t decode_staging = rows * d * sizeof(float); - uint64_t peak = file_size_in_bytes + resident_vectors + resident_graph + decode_staging; - - if (rows == 0 || d == 0) { - // Row/dim metadata is unavailable at estimate time (num_rows can - // arrive as 0 from the load info). rows*dim*4 collapses to 0 and the - // estimate degenerates to ~file_size, grossly under-reserving the - // fp32-expanded peak. Guard with a conservative multiple of the file - // size so admission still throttles concurrent GPU uploads. - LOG_KNOWHERE_WARNING_ << "GPU_HNSW load estimate missing row/dim metadata (num_rows=" << num_rows - << ", dim=" << dim << "); falling back to file-size-based peak"; - const uint64_t fallback = file_size_in_bytes * 9; - if (peak < fallback) { - peak = fallback; - } - } + // once the graph/vectors are on the GPU, so a loaded segment holds + // no steady-state host RAM (its data lives in VRAM). + // maxMemoryCost (transient peak during load): the host RAM that is + // briefly live before the GPU upload frees the CPU copy. + // + // The on-disk index for this collection is the compact FAISS int8-SQ + // form (IndexHNSWSQCosine: signed-int8 codes + graph), the SAME form the + // CPU HNSW cluster loads resident at ~file_size. Deserialize reads it + // into host RAM at ~file_size (no fp32 expansion — the int8 codes are + // uploaded to the device directly, so there is NO fp32 reconstruct/decode + // staging on this path). The transient peak is therefore dominated by the + // serialized read buffer plus the deserialized compact index, ~2x the + // on-disk file size — not the fp32 rows*dim*4 blowup that a Flat/hnswlib + // path would incur. Basing the estimate on file_size also avoids relying + // on num_rows/dim, which can arrive as 0 at estimate time. + (void)num_rows; + (void)dim; + (void)config; + (void)version; + + const uint64_t peak = file_size_in_bytes * 2; + + LOG_KNOWHERE_INFO_ << "GPU_HNSW load estimate (compact int8): file_size=" << file_size_in_bytes + << " transient_peak=" << peak << " retained=0"; return Resource{.memoryCost = 0, .diskCost = 0, .maxMemoryCost = peak}; } From db139ada8354c8fcc5eade3ba92e54aa9a5b4b0b Mon Sep 17 00:00:00 2001 From: Premal Shah Date: Mon, 13 Jul 2026 06:23:40 +0000 Subject: [PATCH 21/24] fix(gpu_hnsw): update load-resource test to match file_size*2 estimate The test was written for the intermediate 4d3d97d commit which used an fp32-expansion-based estimate. After 7d8a448 changed the implementation to file_size*2, the test assertions were not updated and would fail on any CUDA-enabled build. Update the int8-file and no-rows test cases to assert >= 2*file_size instead of the old fp32-based bounds. Also clarify the StaticEstimateLoadResource comment to note that the no-fp32-staging claim applies to native int8/fp16/bf16 paths only, not unsigned SQ8 (QT_8bit) which decodes via sa_decode(). Co-Authored-By: Claude Opus 4.8 Signed-off-by: Premal Shah --- src/index/hnsw/faiss_hnsw.cc | 17 ++++++++++------- tests/ut/test_gpu_search.cc | 36 +++++++++++++----------------------- 2 files changed, 23 insertions(+), 30 deletions(-) diff --git a/src/index/hnsw/faiss_hnsw.cc b/src/index/hnsw/faiss_hnsw.cc index 79d3e6b2c..668e18c54 100644 --- a/src/index/hnsw/faiss_hnsw.cc +++ b/src/index/hnsw/faiss_hnsw.cc @@ -3369,13 +3369,16 @@ class GpuHnswIndexNode : public BaseFaissRegularIndexHNSWNode { // The on-disk index for this collection is the compact FAISS int8-SQ // form (IndexHNSWSQCosine: signed-int8 codes + graph), the SAME form the // CPU HNSW cluster loads resident at ~file_size. Deserialize reads it - // into host RAM at ~file_size (no fp32 expansion — the int8 codes are - // uploaded to the device directly, so there is NO fp32 reconstruct/decode - // staging on this path). The transient peak is therefore dominated by the - // serialized read buffer plus the deserialized compact index, ~2x the - // on-disk file size — not the fp32 rows*dim*4 blowup that a Flat/hnswlib - // path would incur. Basing the estimate on file_size also avoids relying - // on num_rows/dim, which can arrive as 0 at estimate time. + // into host RAM at ~file_size. For native int8/fp16/bf16 upload paths + // the codes are uploaded to the device directly with no fp32 + // reconstruct/decode staging. (Note: unsigned SQ8 via QT_8bit does call + // sa_decode() into a float buffer, so that path would incur fp32 + // staging — but it is not used in production GPU HNSW loads here.) + // The transient peak is therefore dominated by the serialized read + // buffer plus the deserialized compact index, ~2x the on-disk file + // size — not the fp32 rows*dim*4 blowup that a Flat/hnswlib path would + // incur. Basing the estimate on file_size also avoids relying on + // num_rows/dim, which can arrive as 0 at estimate time. (void)num_rows; (void)dim; (void)config; diff --git a/tests/ut/test_gpu_search.cc b/tests/ut/test_gpu_search.cc index 2a334f727..4d3a7274f 100644 --- a/tests/ut/test_gpu_search.cc +++ b/tests/ut/test_gpu_search.cc @@ -1293,8 +1293,8 @@ TEST_CASE("Test GPU HNSW Codex P1 Regressions", "[gpu_hnsw_p1]") { // #3: phase-separated load-resource accounting. GPU_HNSW frees its CPU copy // after uploading to VRAM, so retained memoryCost is 0 while the transient - // peak (maxMemoryCost) must be non-zero and cover the download buffer + - // deserialized CPU index + decode/graph staging. This estimate is a static + // peak (maxMemoryCost) must be non-zero and cover the serialized read buffer + // + deserialized compact index (~2x file_size). This estimate is a static // computation and needs no live GPU. SECTION("Load-resource estimate: memoryCost=0, maxMemoryCost covers peak") { knowhere::Json params; @@ -1308,37 +1308,27 @@ TEST_CASE("Test GPU HNSW Codex P1 Regressions", "[gpu_hnsw_p1]") { // Retained host memory is ~0 after the CPU copy is freed. REQUIRE(gpu_res.value().memoryCost == 0); REQUIRE(gpu_res.value().diskCost == 0); - // Transient peak is non-zero and at least covers the fp32 decode staging - // plus the download buffer. + // Transient peak is non-zero and at least 2x the on-disk file size + // (serialized read buffer + deserialized compact index). REQUIRE(gpu_res.value().maxMemoryCost > 0); - const uint64_t decode_staging = static_cast(nb) * dim * sizeof(float); - REQUIRE(gpu_res.value().maxMemoryCost >= file_size + decode_staging); - - // The estimate must be driven by the fp32-expanded resident index, NOT - // by the on-disk file size. A compressed (e.g. int8) index stores far - // fewer bytes on disk than the fp32 graph hnswlib materializes in host - // RAM, so the peak has to cover num_rows*dim*4 (resident vectors) + - // num_rows*dim*4 (decode staging) independent of file_size. Simulate an - // int8-sized file (1 byte/element) and assert the peak still reflects - // the fp32 in-memory footprint (this is the production OOM regression: - // an ~file_size-tracking estimate under-reserves and OOMs the host). + REQUIRE(gpu_res.value().maxMemoryCost >= 2 * file_size); + + // The estimate is driven by file_size * 2 (compact int8 upload path: + // no fp32 expansion staging). For a compressed (e.g. int8) file the + // peak is 2 * int8_file — NOT the fp32-expanded footprint that the + // old hnswlib-based estimate used. const uint64_t int8_file = static_cast(nb) * dim * sizeof(int8_t); auto compressed_res = knowhere::IndexStaticFaced::EstimateLoadResource( knowhere::IndexEnum::INDEX_GPU_HNSW, version, int8_file, nb, dim, params); REQUIRE(compressed_res.has_value()); - const uint64_t fp32_resident_plus_staging = 2 * static_cast(nb) * dim * sizeof(float); - REQUIRE(compressed_res.value().maxMemoryCost >= fp32_resident_plus_staging); - // And it must dwarf the naive 2*file_size an ~file-tracking estimate - // would have produced. - REQUIRE(compressed_res.value().maxMemoryCost > 2 * int8_file); + REQUIRE(compressed_res.value().maxMemoryCost >= 2 * int8_file); // When row/dim metadata is missing (num_rows arrives as 0 from the load - // info), the rows*dim terms vanish; the estimate must not collapse to - // ~file_size but fall back to a conservative multiple of it. + // info), the estimate still works because it depends only on file_size. auto norows_res = knowhere::IndexStaticFaced::EstimateLoadResource( knowhere::IndexEnum::INDEX_GPU_HNSW, version, int8_file, 0, dim, params); REQUIRE(norows_res.has_value()); - REQUIRE(norows_res.value().maxMemoryCost >= int8_file * 9); + REQUIRE(norows_res.value().maxMemoryCost >= 2 * int8_file); // A plain CPU HNSW keeps its data resident: memoryCost is non-zero and it // does not opt into the peak field (maxMemoryCost stays 0 -> Milvus falls From 2f5a6486fcebcf2455148722f90774bdab531681 Mon Sep 17 00:00:00 2001 From: Premal Shah Date: Mon, 13 Jul 2026 06:58:36 +0000 Subject: [PATCH 22/24] feat(gpu_hnsw): native int8 search via DP4A, bypass MockWrapper upcast - Add searchHostInt8() to GpuIndexHNSW: applies +128 query shift to match upload_int8_dataset encoding, uploads int8 queries to d_queries_i8 - Add layer0_beam_search_kernel_int8: uses __dp4a for 4x int8 MADs/cycle - Dispatch in GpuHnswIndexNode::Search when data_format==int8, bypassing IndexNodeDataMockWrapper fp32 upcast - dim=384: max accumulator 6.2M << INT32_MAX, no overflow risk --- src/index/hnsw/faiss_hnsw.cc | 67 ++++- thirdparty/faiss/faiss/gpu/GpuIndexHNSW.cu | 89 ++++++ thirdparty/faiss/faiss/gpu/GpuIndexHNSW.h | 12 + .../faiss/faiss/gpu/impl/GpuHnswSearch.cuh | 106 +++++++ .../faiss/gpu/impl/GpuHnswSearchKernel.cuh | 272 ++++++++++++++++++ .../faiss/faiss/gpu/impl/GpuHnswTypes.cu | 14 +- .../faiss/faiss/gpu/impl/GpuHnswTypes.h | 4 +- 7 files changed, 557 insertions(+), 7 deletions(-) diff --git a/src/index/hnsw/faiss_hnsw.cc b/src/index/hnsw/faiss_hnsw.cc index 668e18c54..64ef7d5f4 100644 --- a/src/index/hnsw/faiss_hnsw.cc +++ b/src/index/hnsw/faiss_hnsw.cc @@ -3341,6 +3341,11 @@ class GpuHnswIndexNode : public BaseFaissRegularIndexHNSWNode { : BaseFaissRegularIndexHNSWNode(version, object, DataFormatEnum::fp32) { } + // Constructor for native data-type variants (e.g. int8 DP4A path). + GpuHnswIndexNode(const int32_t& version, const Object& object, DataFormatEnum query_format) + : BaseFaissRegularIndexHNSWNode(version, object, query_format) { + } + static std::unique_ptr StaticCreateConfig() { return std::make_unique(); @@ -3554,6 +3559,53 @@ class GpuHnswIndexNode : public BaseFaissRegularIndexHNSWNode { auto nq = dataset->GetRows(); auto dim = dataset->GetDim(); auto ef = hnsw_cfg.ef.value_or(200); + + auto h_ids = std::make_unique(nq * k); + auto h_dist = std::make_unique(nq * k); + + // Native int8 DP4A path: queries arrive as raw int8 (no MockWrapper upcast). + if (data_format == DataFormatEnum::int8) { + const auto* h_queries_i8 = + reinterpret_cast(dataset->GetTensor()); + try { + faiss::gpu::GpuHnswSearchParams gsp; + gsp.ef = ef; + gpu_snapshot->searchHostInt8(nq, h_queries_i8, k, h_dist.get(), h_ids.get(), gsp); + } catch (const std::exception& e) { + LOG_KNOWHERE_ERROR_ << "GPU_HNSW int8 search failed: " << e.what(); + return expected::Err(Status::cuvs_inner_error, + std::string("GPU HNSW int8 search failed: ") + e.what()); + } + + // For COSINE, post-scale distances by 1/||q|| so scores are in [-1, 1]. + // The kernel computes -dot(shifted_q, shifted_db) which equals the + // unscaled negative cosine numerator; divide by query norm to finalize. + if (IsMetricType(hnsw_cfg.metric_type.value(), metric::COSINE)) { + for (int64_t i = 0; i < static_cast(nq); i++) { + const int8_t* q = h_queries_i8 + i * dim; + int32_t sq_norm = 0; + for (int64_t d = 0; d < static_cast(dim); d++) { + sq_norm += static_cast(q[d]) * static_cast(q[d]); + } + float inv_qnorm = (sq_norm > 0) ? (1.0f / std::sqrt(static_cast(sq_norm))) : 1.0f; + for (int64_t j = 0; j < static_cast(k); j++) { + h_dist[i * k + j] *= inv_qnorm; + } + } + } + + // Negate back to positive for IP and COSINE. + if (IsMetricType(hnsw_cfg.metric_type.value(), metric::IP) || + IsMetricType(hnsw_cfg.metric_type.value(), metric::COSINE)) { + for (int64_t i = 0; i < static_cast(nq * k); i++) { + h_dist[i] = -h_dist[i]; + } + } + + return GenResultDataSet(nq, k, h_ids.release(), h_dist.release()); + } + + // fp32 path (fp16/bf16 have already been upcast by MockWrapper). const auto* h_queries_raw = reinterpret_cast(dataset->GetTensor()); // For COSINE metric, normalize queries to unit length. @@ -3572,9 +3624,6 @@ class GpuHnswIndexNode : public BaseFaissRegularIndexHNSWNode { h_queries = normalized_queries.get(); } - auto h_ids = std::make_unique(nq * k); - auto h_dist = std::make_unique(nq * k); - try { faiss::gpu::GpuHnswSearchParams gsp; gsp.ef = ef; @@ -3694,6 +3743,10 @@ class GpuHnswSQIndexNode : public GpuHnswIndexNode { GpuHnswSQIndexNode(const int32_t& version, const Object& object) : GpuHnswIndexNode(version, object) { } + GpuHnswSQIndexNode(const int32_t& version, const Object& object, DataFormatEnum query_format) + : GpuHnswIndexNode(version, object, query_format) { + } + std::string Type() const override { return IndexEnum::INDEX_GPU_HNSW_SQ; @@ -3739,7 +3792,9 @@ KNOWHERE_REGISTER_GLOBAL( KNOWHERE_REGISTER_GLOBAL( GPU_HNSW, [](const int32_t& version, const Object& object) { - return Index>::Create(std::make_unique(version, object)); + // Native int8 path: bypass MockWrapper upcast, pass int8 queries directly + // to searchHostInt8() which applies the +128 shift and uses DP4A kernel. + return Index::Create(version, object, DataFormatEnum::int8); }, int8, true, (feature::INT8 | feature::GPU)); @@ -3765,7 +3820,9 @@ KNOWHERE_REGISTER_GLOBAL( KNOWHERE_REGISTER_GLOBAL( GPU_HNSW_SQ, [](const int32_t& version, const Object& object) { - return Index>::Create(std::make_unique(version, object)); + // Native int8 path: bypass MockWrapper upcast, pass int8 queries directly + // to searchHostInt8() which applies the +128 shift and uses DP4A kernel. + return Index::Create(version, object, DataFormatEnum::int8); }, int8, true, (feature::INT8 | feature::GPU)); #endif // KNOWHERE_WITH_CUVS diff --git a/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.cu b/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.cu index 8c099e120..ce5978e77 100644 --- a/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.cu +++ b/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.cu @@ -287,5 +287,94 @@ void GpuIndexHNSW::searchHost( } } +void GpuIndexHNSW::searchHostInt8( + idx_t n, + const int8_t* x_host, + int k, + float* distances_host, + idx_t* labels_host, + const GpuHnswSearchParams& sp) const { + FAISS_THROW_IF_NOT_MSG( + this->is_trained && deviceIndex_, + "Index not loaded. Call copyFrom() first."); + FAISS_THROW_IF_NOT_MSG(n > 0, "n must be > 0"); + + auto& idx = *deviceIndex_; + FAISS_THROW_IF_NOT_FMT( + idx.dim % 4 == 0, + "searchHostInt8: dim=%ld is not divisible by 4 (required for DP4A)", + idx.dim); + + GPU_HNSW_CUDA_CHECK(cudaSetDevice(config_.device)); + DeviceScope scope(config_.device); + + ScratchPoolGuard guard(*idx.scratch_pool); + auto* slot = guard.get(); + auto& sc = slot->scratch; + cudaStream_t stream = slot->stream; + + int nq = static_cast(n); + int dim = static_cast(idx.dim); + int64_t nelem = static_cast(nq) * dim; + + sc.ensure(nq, k, dim, static_cast(idx.n_rows), /*use_i8_queries=*/true); + + // Apply +128 shift to queries to match upload_int8_dataset's -128 encoding. + // upload_int8_dataset stores: gpu_val = user_val - 128 + // So query shift: shifted_q = user_q + 128 → same encoding as stored vectors. + auto shifted = std::make_unique(nelem); + for (int64_t i = 0; i < nelem; i++) { + shifted[i] = static_cast(static_cast(x_host[i]) + 128); + } + + // Upload shifted int8 queries to d_queries_i8 (for DP4A layer-0 kernel). + GPU_HNSW_CUDA_CHECK(cudaMemcpyAsync( + sc.d_queries_i8, + shifted.get(), + nelem * sizeof(int8_t), + cudaMemcpyHostToDevice, + stream)); + + // Also upload fp32 queries to d_queries for upper-layer greedy search. + auto fp32_q = std::make_unique(nelem); + for (int64_t i = 0; i < nelem; i++) { + fp32_q[i] = static_cast(x_host[i]); + } + GPU_HNSW_CUDA_CHECK(cudaMemcpyAsync( + sc.d_queries, + fp32_q.get(), + nelem * sizeof(float), + cudaMemcpyHostToDevice, + stream)); + + // Synchronize to ensure host buffers (shifted, fp32_q) are not freed + // while the async copies are in flight. + GPU_HNSW_CUDA_CHECK(cudaStreamSynchronize(stream)); + + gpu_hnsw_search(stream, sp, idx, sc, nq, k); + + GPU_HNSW_CUDA_CHECK(cudaMemcpyAsync( + distances_host, + sc.d_distances, + static_cast(nq) * k * sizeof(float), + cudaMemcpyDeviceToHost, + stream)); + + auto tmp = std::make_unique(nq * k); + GPU_HNSW_CUDA_CHECK(cudaMemcpyAsync( + tmp.get(), + sc.d_neighbors, + static_cast(nq) * k * sizeof(uint64_t), + cudaMemcpyDeviceToHost, + stream)); + + GPU_HNSW_CUDA_CHECK(cudaStreamSynchronize(stream)); + + for (int i = 0; i < nq * k; i++) { + labels_host[i] = + (tmp[i] == UINT64_MAX) ? -1 : static_cast(tmp[i]); + } +} + } // namespace gpu } // namespace faiss diff --git a/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.h b/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.h index 0e4848117..05be2b9c4 100644 --- a/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.h +++ b/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.h @@ -123,6 +123,18 @@ struct GpuIndexHNSW : public GpuIndex { idx_t* labels_host, const GpuHnswSearchParams& params) const; + /// Search with int8 host query vectors using the native DP4A path. + /// Applies +128 shift to queries to match upload_int8_dataset encoding + /// (which stores as value - 128). dim must be divisible by 4. + /// All input/output pointers must be host memory. + void searchHostInt8( + idx_t n, + const int8_t* x_host, + int k, + float* distances_host, + idx_t* labels_host, + const GpuHnswSearchParams& params) const; + protected: bool addImplRequiresIDs_() const override; diff --git a/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearch.cuh b/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearch.cuh index 671c4d0a7..be9d990f8 100644 --- a/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearch.cuh +++ b/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearch.cuh @@ -188,6 +188,112 @@ inline void gpu_hnsw_search( switch (idx.dataset_type) { case GpuHnswDatasetType::INT8: + // Use the DP4A native int8 kernel when int8 queries are available. + if (sc.d_queries_i8 != nullptr) { + // Upper-layer greedy search still uses fp32 queries (d_queries). + if (num_upper_layers > 0) { + auto* d_layer_ptrs = static_cast( + idx.d_upper_layer_ptrs); + int warps_per_block = 4; + int threads_per_block = warps_per_block * 32; + int num_blocks = + (num_queries + warps_per_block - 1) / warps_per_block; + hnsw_kernel::upper_layer_search_kernel + <<>>( + sc.d_queries, + static_cast(idx.d_dataset), + idx.d_inv_norms, + d_layer_ptrs, + sc.d_entry_points, + idx.entry_point, + num_queries, + dim, + num_upper_layers, + idx.use_ip); + GPU_HNSW_CUDA_CHECK(cudaGetLastError()); + } else { + std::vector h_eps(num_queries, idx.entry_point); + GPU_HNSW_CUDA_CHECK(cudaMemcpyAsync( + sc.d_entry_points, + h_eps.data(), + num_queries * sizeof(uint32_t), + cudaMemcpyHostToDevice, + stream)); + GPU_HNSW_CUDA_CHECK(cudaStreamSynchronize(stream)); + } + + int block_size = + params.thread_block_size > 0 ? params.thread_block_size : 128; + + int smem_max = 49152; + { + int device = 0; + int optin = 0; + if (cudaGetDevice(&device) == cudaSuccess && + cudaDeviceGetAttribute( + &optin, + cudaDevAttrMaxSharedMemoryPerBlockOptin, + device) == cudaSuccess && + optin > smem_max) { + smem_max = optin; + } + } + + { + int smem_overhead = sw * idx.max_degree0 * 8 + sw * 4 + 12; + int max_ef = (smem_max - smem_overhead) / 12; + if (max_ef < 1) { + throw std::runtime_error( + std::string("gpu_hnsw: search_width=") + + std::to_string(sw) + + " too large for device shared memory (" + + std::to_string(smem_max) + + " bytes); reduce search_width"); + } + if (ef > max_ef) { + ef = max_ef; + } + } + + size_t smem_size = hnsw_kernel::calc_layer0_smem_size( + ef, sw, idx.max_degree0); + + if (smem_size > 49152) { + GPU_HNSW_CUDA_CHECK(cudaFuncSetAttribute( + hnsw_kernel::layer0_beam_search_kernel_int8, + cudaFuncAttributeMaxDynamicSharedMemorySize, + static_cast(smem_size))); + } + + int N_int = static_cast(idx.n_rows); + size_t bitmap_bytes = hnsw_kernel::calc_visited_bitmap_size( + num_queries, N_int); + GPU_HNSW_CUDA_CHECK(cudaMemsetAsync( + sc.d_visited_bitmaps, 0, bitmap_bytes, stream)); + + hnsw_kernel::layer0_beam_search_kernel_int8 + <<>>( + reinterpret_cast(sc.d_queries_i8), + static_cast(idx.d_dataset), + idx.d_inv_norms, + idx.d_layer0_graph, + sc.d_entry_points, + sc.d_visited_bitmaps, + sc.d_neighbors, + sc.d_distances, + num_queries, + N_int, + dim, + idx.max_degree0, + k, + ef, + sw, + max_iter, + idx.use_ip); + GPU_HNSW_CUDA_CHECK(cudaGetLastError()); + break; + } + // Fall through to generic int8 fp32-query path if no i8 queries. launch_kernels( static_cast(idx.d_dataset), idx.d_inv_norms); break; diff --git a/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearchKernel.cuh b/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearchKernel.cuh index 2133f4327..b0083ddb1 100644 --- a/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearchKernel.cuh +++ b/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearchKernel.cuh @@ -78,6 +78,23 @@ __device__ __forceinline__ float thread_ip_distance( return -sum; } +// DP4A inner-product distance for int8 queries vs int8 dataset. +// Both query and vec must already be in the same encoding (both shifted +128 +// relative to the user's signed int8 values, matching upload_int8_dataset). +// Computes -dot(query, vec) as a float (max-IP / negative dot convention). +// dim4 = dim / 4 (caller must ensure dim % 4 == 0). +__device__ __forceinline__ float thread_ip_distance_dp4a( + const int32_t* __restrict__ query_packed, + const int32_t* __restrict__ vec_packed, + int dim4) { + int sum = 0; +#pragma unroll 8 + for (int d = 0; d < dim4; d++) { + sum = __dp4a(query_packed[d], vec_packed[d], sum); + } + return static_cast(-sum); // negate: max-IP convention +} + // ============================================================================ // Phase 1: Upper-layer greedy search // ============================================================================ @@ -476,6 +493,261 @@ __global__ void layer0_beam_search_kernel( } } +// ============================================================================ +// Phase 2 (int8/DP4A variant): Layer-0 parallel beam search kernel +// Uses __dp4a for 4x int8 MADs per cycle instead of fp32 FMA. +// d_queries_packed: int8 queries reinterpreted as int32 (shifted +128 to +// match upload_int8_dataset encoding); dim must be divisible by 4. +// d_dataset: raw int8 dataset on device (already shifted -128 at upload time, +// now consistent with query shift so DP4A computes the correct dot product). +// ============================================================================ +__global__ void layer0_beam_search_kernel_int8( + const int32_t* __restrict__ d_queries_packed, + const int8_t* __restrict__ d_dataset, + const float* __restrict__ d_inv_norms, + const uint32_t* __restrict__ d_layer0_graph, + const uint32_t* __restrict__ d_entry_points, + uint32_t* __restrict__ d_visited_bitmaps, + uint64_t* __restrict__ d_neighbors, + float* __restrict__ d_distances, + int num_queries, + int N, + int dim, + int max_degree0, + int k, + int ef, + int search_width, + int max_iterations, + bool use_inner_product) { + int query_idx = blockIdx.x; + if (query_idx >= num_queries) + return; + + int dim4 = dim / 4; + const int32_t* query = d_queries_packed + static_cast(query_idx) * dim4; + + extern __shared__ char smem[]; + + int max_staging = search_width * max_degree0; + + uint32_t* result_ids = reinterpret_cast(smem); + float* result_dists = reinterpret_cast(result_ids + ef); + uint32_t* is_expanded = reinterpret_cast(result_dists + ef); + uint32_t* staging_ids = is_expanded + ef; + float* staging_dists = reinterpret_cast(staging_ids + max_staging); + uint32_t* parent_ids = + reinterpret_cast(staging_dists + max_staging); + int* meta = reinterpret_cast(parent_ids + search_width); + + int bitmap_words = (N + 31) / 32; + uint32_t* visited_bmap = + d_visited_bitmaps + static_cast(query_idx) * bitmap_words; + + for (int i = threadIdx.x; i < ef; i += blockDim.x) { + result_ids[i] = UINT32_MAX; + result_dists[i] = FLT_MAX; + is_expanded[i] = 0; + } + if (threadIdx.x == 0) { + meta[0] = 0; + meta[1] = 0; + meta[2] = 0; + } + __syncthreads(); + + // --- Seed with entry point --- + uint32_t ep = d_entry_points[query_idx]; + if (threadIdx.x == 0) { + const int32_t* ep_packed = + reinterpret_cast(d_dataset + static_cast(ep) * dim); + float ep_dist = thread_ip_distance_dp4a(query, ep_packed, dim4); + if (!use_inner_product) { + // L2 fallback: should not happen for int8 DP4A path, but guard + ep_dist = ep_dist; // keep as-is (ip convention, distance is valid) + } + if (d_inv_norms) + ep_dist *= __ldg(&d_inv_norms[ep]); + result_ids[0] = ep; + result_dists[0] = ep_dist; + is_expanded[0] = 0; + meta[0] = 1; + bitmap_visit(visited_bmap, ep); + } + __syncthreads(); + + // --- Seed with entry point's neighbors --- + if (threadIdx.x == 0) + meta[1] = 0; + __syncthreads(); + + for (int j = threadIdx.x; j < max_degree0; j += blockDim.x) { + uint32_t nbr = + d_layer0_graph[static_cast(ep) * max_degree0 + j]; + if (nbr == UINT32_MAX || nbr >= static_cast(N)) + continue; + if (!bitmap_visit(visited_bmap, nbr)) + continue; + + const int32_t* nbr_packed = + reinterpret_cast(d_dataset + static_cast(nbr) * dim); + float dist = thread_ip_distance_dp4a(query, nbr_packed, dim4); + if (d_inv_norms) + dist *= d_inv_norms[nbr]; + + int slot = atomicAdd(&meta[1], 1); + if (slot < max_staging) { + staging_ids[slot] = nbr; + staging_dists[slot] = dist; + } + } + __syncthreads(); + + if (threadIdx.x == 0) { + int staging_count = min(meta[1], max_staging); + int rc = meta[0]; + + for (int s = 0; s < staging_count; s++) { + uint32_t sid = staging_ids[s]; + float sdist = staging_dists[s]; + if (rc >= ef && sdist >= result_dists[rc - 1]) + continue; + + int lo = 0, hi = rc; + while (lo < hi) { + int mid = (lo + hi) / 2; + if (result_dists[mid] < sdist) + lo = mid + 1; + else + hi = mid; + } + int insert_end = rc < ef ? rc : ef - 1; + for (int i = insert_end; i > lo; i--) { + result_ids[i] = result_ids[i - 1]; + result_dists[i] = result_dists[i - 1]; + is_expanded[i] = is_expanded[i - 1]; + } + result_ids[lo] = sid; + result_dists[lo] = sdist; + is_expanded[lo] = 0; + if (rc < ef) + rc++; + } + + for (int i = 0; i < rc; i++) { + if (result_ids[i] == ep) { + is_expanded[i] = 1; + break; + } + } + meta[0] = rc; + } + __syncthreads(); + + // --- Unified main loop --- + for (int iter = 0; iter < max_iterations; iter++) { + if (threadIdx.x == 0) { + int num_parents = 0; + int rc = meta[0]; + + for (int i = 0; i < rc && num_parents < search_width; i++) { + if (!is_expanded[i]) { + parent_ids[num_parents++] = result_ids[i]; + is_expanded[i] = 1; + } + } + + meta[2] = num_parents; + } + __syncthreads(); + + int num_parents = meta[2]; + if (num_parents == 0) + break; + + if (threadIdx.x == 0) + meta[1] = 0; + __syncthreads(); + + int total_work = num_parents * max_degree0; + for (int wi = threadIdx.x; wi < total_work; wi += blockDim.x) { + int parent_idx = wi / max_degree0; + int nbr_slot = wi % max_degree0; + + uint32_t parent = parent_ids[parent_idx]; + uint32_t nbr = + d_layer0_graph + [static_cast(parent) * max_degree0 + + nbr_slot]; + if (nbr == UINT32_MAX || nbr >= static_cast(N)) + continue; + if (!bitmap_visit(visited_bmap, nbr)) + continue; + + const int32_t* nbr_packed = + reinterpret_cast(d_dataset + static_cast(nbr) * dim); + float dist = thread_ip_distance_dp4a(query, nbr_packed, dim4); + if (d_inv_norms) + dist *= d_inv_norms[nbr]; + + int slot = atomicAdd(&meta[1], 1); + if (slot < max_staging) { + staging_ids[slot] = nbr; + staging_dists[slot] = dist; + } + } + __syncthreads(); + + if (threadIdx.x == 0) { + int staging_count = min(meta[1], max_staging); + int rc = meta[0]; + + for (int s = 0; s < staging_count; s++) { + uint32_t sid = staging_ids[s]; + float sdist = staging_dists[s]; + if (rc >= ef && sdist >= result_dists[rc - 1]) + continue; + + int lo = 0, hi = rc; + while (lo < hi) { + int mid = (lo + hi) / 2; + if (result_dists[mid] < sdist) + lo = mid + 1; + else + hi = mid; + } + int insert_end = rc < ef ? rc : ef - 1; + for (int i = insert_end; i > lo; i--) { + result_ids[i] = result_ids[i - 1]; + result_dists[i] = result_dists[i - 1]; + is_expanded[i] = is_expanded[i - 1]; + } + result_ids[lo] = sid; + result_dists[lo] = sdist; + is_expanded[lo] = 0; + if (rc < ef) + rc++; + } + + meta[0] = rc; + } + __syncthreads(); + } + + // --- Copy top-k results to global memory --- + int rc = meta[0]; + for (int i = threadIdx.x; i < k; i += blockDim.x) { + if (i < rc) { + d_neighbors[static_cast(query_idx) * k + i] = + static_cast(result_ids[i]); + d_distances[static_cast(query_idx) * k + i] = + result_dists[i]; + } else { + d_neighbors[static_cast(query_idx) * k + i] = UINT64_MAX; + d_distances[static_cast(query_idx) * k + i] = FLT_MAX; + } + } +} + inline size_t calc_layer0_smem_size(int ef, int search_width, int max_degree0) { int max_staging = search_width * max_degree0; diff --git a/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.cu b/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.cu index 5973b94ac..67d00365a 100644 --- a/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.cu +++ b/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.cu @@ -48,7 +48,8 @@ void GpuHnswSearchScratch::ensure( int nq, int k, int dim, - int N) { + int N, + bool use_i8_queries) { size_t need_q = static_cast(nq) * dim * sizeof(float); if (need_q > queries_bytes) { if (d_queries) @@ -87,6 +88,15 @@ void GpuHnswSearchScratch::ensure( SCRATCH_CUDA_CHECK(cudaMalloc(&d_visited_bitmaps, need_bm)); bitmap_bytes = need_bm; } + if (use_i8_queries) { + size_t need_i8 = static_cast(nq) * dim * sizeof(int8_t); + if (need_i8 > queries_i8_bytes) { + if (d_queries_i8) + cudaFree(d_queries_i8); + SCRATCH_CUDA_CHECK(cudaMalloc(&d_queries_i8, need_i8)); + queries_i8_bytes = need_i8; + } + } } GpuHnswSearchScratch::~GpuHnswSearchScratch() { @@ -103,6 +113,8 @@ GpuHnswSearchScratch::~GpuHnswSearchScratch() { cudaFree(d_entry_points); if (d_visited_bitmaps) cudaFree(d_visited_bitmaps); + if (d_queries_i8) + cudaFree(d_queries_i8); } GpuHnswScratchSlot::~GpuHnswScratchSlot() { diff --git a/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.h b/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.h index 98a001c03..1f15a3cca 100644 --- a/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.h +++ b/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.h @@ -96,18 +96,20 @@ struct GpuHnswSearchScratch { float* d_distances = nullptr; uint32_t* d_entry_points = nullptr; uint32_t* d_visited_bitmaps = nullptr; + int8_t* d_queries_i8 = nullptr; // int8 queries for DP4A path (shifted +128) size_t queries_bytes = 0; size_t neighbors_bytes = 0; size_t distances_bytes = 0; int entry_cap = 0; size_t bitmap_bytes = 0; + size_t queries_i8_bytes = 0; // Device this scratch's allocations live on; used to set the CUDA device // context before freeing in the destructor (multi-GPU correctness). int device = 0; - void ensure(int nq, int k, int dim, int N); + void ensure(int nq, int k, int dim, int N, bool use_i8_queries = false); ~GpuHnswSearchScratch(); From b5b1726c0d7b63f6a438cc22590e46cd55547724 Mon Sep 17 00:00:00 2001 From: Premal Shah Date: Mon, 13 Jul 2026 07:01:46 +0000 Subject: [PATCH 23/24] feat(gpu_hnsw): add DP4A dispatch in gpu_hnsw_search for INT8 dataset type Add INT8 path to gpu_hnsw_search() that dispatches to layer0_beam_search_kernel_int8 when sc.d_queries_i8 is populated. Upper-layer search still uses fp32 queries (d_queries) for the greedy entry-point descent through sparse upper layers. --- .../faiss/faiss/gpu/impl/GpuHnswSearch.cuh | 20 +++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearch.cuh b/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearch.cuh index be9d990f8..09f45b3f2 100644 --- a/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearch.cuh +++ b/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearch.cuh @@ -225,7 +225,7 @@ inline void gpu_hnsw_search( int block_size = params.thread_block_size > 0 ? params.thread_block_size : 128; - int smem_max = 49152; + int smem_max_i8 = 49152; { int device = 0; int optin = 0; @@ -234,20 +234,20 @@ inline void gpu_hnsw_search( &optin, cudaDevAttrMaxSharedMemoryPerBlockOptin, device) == cudaSuccess && - optin > smem_max) { - smem_max = optin; + optin > smem_max_i8) { + smem_max_i8 = optin; } } { int smem_overhead = sw * idx.max_degree0 * 8 + sw * 4 + 12; - int max_ef = (smem_max - smem_overhead) / 12; + int max_ef = (smem_max_i8 - smem_overhead) / 12; if (max_ef < 1) { throw std::runtime_error( - std::string("gpu_hnsw: search_width=") + + std::string("gpu_hnsw_int8: search_width=") + std::to_string(sw) + " too large for device shared memory (" + - std::to_string(smem_max) + + std::to_string(smem_max_i8) + " bytes); reduce search_width"); } if (ef > max_ef) { @@ -255,14 +255,14 @@ inline void gpu_hnsw_search( } } - size_t smem_size = hnsw_kernel::calc_layer0_smem_size( + size_t smem_size_i8 = hnsw_kernel::calc_layer0_smem_size( ef, sw, idx.max_degree0); - if (smem_size > 49152) { + if (smem_size_i8 > 49152) { GPU_HNSW_CUDA_CHECK(cudaFuncSetAttribute( hnsw_kernel::layer0_beam_search_kernel_int8, cudaFuncAttributeMaxDynamicSharedMemorySize, - static_cast(smem_size))); + static_cast(smem_size_i8))); } int N_int = static_cast(idx.n_rows); @@ -272,7 +272,7 @@ inline void gpu_hnsw_search( sc.d_visited_bitmaps, 0, bitmap_bytes, stream)); hnsw_kernel::layer0_beam_search_kernel_int8 - <<>>( + <<>>( reinterpret_cast(sc.d_queries_i8), static_cast(idx.d_dataset), idx.d_inv_norms, From 5a5c50285e437424e1a042d55d56134e072b03d7 Mon Sep 17 00:00:00 2001 From: Premal Shah Date: Mon, 13 Jul 2026 07:49:04 +0000 Subject: [PATCH 24/24] fix(gpu_hnsw): fix DP4A encoding bug and L2 guard MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit V5-C1: Remove +128 query shift in searchHostInt8. upload_int8_dataset already reverses FAISS's +128 bias (codes[i]-128), so GPU stores original signed user values. Queries arrive as the same signed int8 values — no shift needed. The +128 shift was causing (q+128)*d instead of q*d. V5-C2: Gate DP4A dispatch on idx.use_ip in GpuHnswSearch.cuh. The DP4A kernel only computes dot products; L2 now falls through to the existing fp32 generic kernel. V5-H1: Replace throw on dim%4!=0 with fp32 fallback via searchHost(). V5-L1: Fix misleading comments describing the encoding. --- thirdparty/faiss/faiss/gpu/GpuIndexHNSW.cu | 31 ++++++++++--------- .../faiss/faiss/gpu/impl/GpuHnswSearch.cuh | 6 ++-- .../faiss/gpu/impl/GpuHnswSearchKernel.cuh | 16 +++++----- 3 files changed, 28 insertions(+), 25 deletions(-) diff --git a/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.cu b/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.cu index ce5978e77..63eb18e91 100644 --- a/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.cu +++ b/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.cu @@ -300,10 +300,16 @@ void GpuIndexHNSW::searchHostInt8( FAISS_THROW_IF_NOT_MSG(n > 0, "n must be > 0"); auto& idx = *deviceIndex_; - FAISS_THROW_IF_NOT_FMT( - idx.dim % 4 == 0, - "searchHostInt8: dim=%ld is not divisible by 4 (required for DP4A)", - idx.dim); + // If dim is not divisible by 4, DP4A cannot be used; fall back to fp32 path. + if (idx.dim % 4 != 0) { + auto fp32_fallback = std::make_unique( + static_cast(n) * idx.dim); + for (int64_t i = 0; i < static_cast(n) * idx.dim; i++) { + fp32_fallback[i] = static_cast(x_host[i]); + } + searchHost(n, fp32_fallback.get(), k, distances_host, labels_host, sp); + return; + } GPU_HNSW_CUDA_CHECK(cudaSetDevice(config_.device)); DeviceScope scope(config_.device); @@ -319,18 +325,13 @@ void GpuIndexHNSW::searchHostInt8( sc.ensure(nq, k, dim, static_cast(idx.n_rows), /*use_i8_queries=*/true); - // Apply +128 shift to queries to match upload_int8_dataset's -128 encoding. - // upload_int8_dataset stores: gpu_val = user_val - 128 - // So query shift: shifted_q = user_q + 128 → same encoding as stored vectors. - auto shifted = std::make_unique(nelem); - for (int64_t i = 0; i < nelem; i++) { - shifted[i] = static_cast(static_cast(x_host[i]) + 128); - } - - // Upload shifted int8 queries to d_queries_i8 (for DP4A layer-0 kernel). + // Upload int8 queries directly — dataset on GPU is already in signed int8 + // (upload_int8_dataset applies codes[i]-128, reversing FAISS's +128 bias, + // yielding the original signed user values). Queries arrive as the same + // signed int8 user values; no shift is needed. GPU_HNSW_CUDA_CHECK(cudaMemcpyAsync( sc.d_queries_i8, - shifted.get(), + x_host, nelem * sizeof(int8_t), cudaMemcpyHostToDevice, stream)); @@ -347,7 +348,7 @@ void GpuIndexHNSW::searchHostInt8( cudaMemcpyHostToDevice, stream)); - // Synchronize to ensure host buffers (shifted, fp32_q) are not freed + // Synchronize to ensure host buffers (fp32_q) are not freed // while the async copies are in flight. GPU_HNSW_CUDA_CHECK(cudaStreamSynchronize(stream)); diff --git a/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearch.cuh b/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearch.cuh index 09f45b3f2..4422c20f0 100644 --- a/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearch.cuh +++ b/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearch.cuh @@ -188,8 +188,10 @@ inline void gpu_hnsw_search( switch (idx.dataset_type) { case GpuHnswDatasetType::INT8: - // Use the DP4A native int8 kernel when int8 queries are available. - if (sc.d_queries_i8 != nullptr) { + // Use the DP4A native int8 kernel when int8 queries are available + // and the metric is IP/COSINE. L2 falls through to the fp32 path + // because the DP4A kernel only computes dot products, not L2 distances. + if (sc.d_queries_i8 != nullptr && idx.use_ip) { // Upper-layer greedy search still uses fp32 queries (d_queries). if (num_upper_layers > 0) { auto* d_layer_ptrs = static_cast( diff --git a/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearchKernel.cuh b/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearchKernel.cuh index b0083ddb1..c0a3ae471 100644 --- a/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearchKernel.cuh +++ b/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearchKernel.cuh @@ -79,8 +79,9 @@ __device__ __forceinline__ float thread_ip_distance( } // DP4A inner-product distance for int8 queries vs int8 dataset. -// Both query and vec must already be in the same encoding (both shifted +128 -// relative to the user's signed int8 values, matching upload_int8_dataset). +// Both query and dataset must be in the same encoding: signed int8 user values. +// Dataset is stored as-is (upload_int8_dataset reverses FAISS's +128 bias). +// Queries are passed without any shift (user signed int8 values directly). // Computes -dot(query, vec) as a float (max-IP / negative dot convention). // dim4 = dim / 4 (caller must ensure dim % 4 == 0). __device__ __forceinline__ float thread_ip_distance_dp4a( @@ -496,8 +497,9 @@ __global__ void layer0_beam_search_kernel( // ============================================================================ // Phase 2 (int8/DP4A variant): Layer-0 parallel beam search kernel // Uses __dp4a for 4x int8 MADs per cycle instead of fp32 FMA. -// d_queries_packed: int8 queries reinterpreted as int32 (shifted +128 to -// match upload_int8_dataset encoding); dim must be divisible by 4. +// d_queries_packed: int8 queries reinterpreted as int32 (signed user values, +// no shift — matches dataset encoding after upload_int8_dataset -128 reversal); +// dim must be divisible by 4. // d_dataset: raw int8 dataset on device (already shifted -128 at upload time, // now consistent with query shift so DP4A computes the correct dot product). // ============================================================================ @@ -560,11 +562,9 @@ __global__ void layer0_beam_search_kernel_int8( if (threadIdx.x == 0) { const int32_t* ep_packed = reinterpret_cast(d_dataset + static_cast(ep) * dim); + // DP4A kernel is only dispatched for IP/COSINE (use_ip=true); L2 + // falls back to the fp32 kernel before this point. float ep_dist = thread_ip_distance_dp4a(query, ep_packed, dim4); - if (!use_inner_product) { - // L2 fallback: should not happen for int8 DP4A path, but guard - ep_dist = ep_dist; // keep as-is (ip convention, distance is valid) - } if (d_inv_norms) ep_dist *= __ldg(&d_inv_norms[ep]); result_ids[0] = ep;