From 8e24da10ba2c4e1aacbbb9c7fb2699a0bc7295b9 Mon Sep 17 00:00:00 2001 From: Devin AI Date: Thu, 9 Jul 2026 06:30:45 +0000 Subject: [PATCH 01/37] 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/37] 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/37] 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/37] 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/37] 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/37] 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/37] 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/37] 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/37] 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/37] 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/37] 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/37] 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/37] 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/37] 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/37] 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/37] 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/37] 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/37] 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/37] 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/37] 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/37] 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/37] 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/37] 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/37] 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; From a82eedbb3b726581104acd165b61f280784e3ca7 Mon Sep 17 00:00:00 2001 From: premal Date: Fri, 17 Jul 2026 01:39:04 +0000 Subject: [PATCH 25/37] feat(gpu_hnsw): re-vendor faiss GPU tree (unified int8/DP4A + parallel merge) Re-vendor faiss/gpu from 6si/faiss (devin/gpu-hnsw-int8-dp4a-unify) so the vendored copy stops diverging from upstream faiss. This pulls in: - the single layer-0 kernel (replaces the separate serial int8 kernel + serial fp32 kernel with one parallel bitonic-merge kernel that also handles native int8 DP4A); - staging padded to next_pow2(search_width*2*M) so non-power-of-two 2*M no longer hard-fails GPU search; - ef-clamp warning and corrected int8/negated-distance header docs. Also: - document the fp32-Flat/unsigned-SQ8 transient-peak under-count in StaticEstimateLoadResource (2x is correct for the compact int8 path production uses; a per-DataType estimator is the proper fix); - add a non-power-of-two 2*M UT covering both the fp32 and int8 DP4A parallel-merge paths (cosine, dim divisible by 4). Signed-off-by: premal Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- src/index/hnsw/faiss_hnsw.cc | 13 + tests/ut/test_gpu_search.cc | 78 +++ thirdparty/faiss/faiss/gpu/GpuIndexHNSW.h | 21 +- .../faiss/faiss/gpu/impl/GpuHnswSearch.cuh | 188 ++---- .../faiss/gpu/impl/GpuHnswSearchKernel.cuh | 574 +++++++----------- .../faiss/faiss/gpu/impl/GpuHnswTypes.h | 2 +- 6 files changed, 388 insertions(+), 488 deletions(-) diff --git a/src/index/hnsw/faiss_hnsw.cc b/src/index/hnsw/faiss_hnsw.cc index 64ef7d5f4..9215ba6a1 100644 --- a/src/index/hnsw/faiss_hnsw.cc +++ b/src/index/hnsw/faiss_hnsw.cc @@ -3384,6 +3384,19 @@ class GpuHnswIndexNode : public BaseFaissRegularIndexHNSWNode { // 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. + // + // KNOWN LIMITATION: this 2x is correct for the native int8/fp16/bf16 + // and int8-SQ upload paths (no fp32 staging). It UNDER-counts the two + // paths that materialize a full fp32 buffer during load — VECTOR_FLOAT + // IndexHNSWFlat (from_faiss_hnsw_flat does a reconstruct_n into a + // host float[]) and unsigned QT_8bit SQ8 (sa_decode into float) — whose + // true transient peak is closer to file_size + rows*dim*4. It is not + // corrected here because this static entry point receives no data-type + // discriminator (the same non-templated StaticEstimateLoadResource is + // registered for fp32/fp16/bf16/int8), so adding a blanket rows*dim*4 + // term would badly over-reserve host RAM on the compact int8 path that + // production actually uses. Fixing it properly needs a per-DataType + // estimator; tracked as a follow-up. (void)num_rows; (void)dim; (void)config; diff --git a/tests/ut/test_gpu_search.cc b/tests/ut/test_gpu_search.cc index 4d3a7274f..62fc13921 100644 --- a/tests/ut/test_gpu_search.cc +++ b/tests/ut/test_gpu_search.cc @@ -779,6 +779,84 @@ TEST_CASE("Test All GPU Index", "[search]") { // recall stays high but with a slightly looser floor. REQUIRE(recall >= 0.85f); } + + SECTION("Test GPU HNSW non-power-of-two 2*M staging (fp32 + int8 DP4A)") { + // M=24 => max_degree0 = 2*M = 48, which is NOT a power of two. The + // parallel bitonic-merge kernel pads the layer-0 staging capacity up to + // the next power of two (64) and grows the block accordingly. Before the + // padding fix this configuration hard-failed every GPU search. Cover + // both the generic fp32-query kernel and the native int8 DP4A kernel + // (dim=128 is divisible by 4, so int8 COSINE takes the DP4A path). + const int64_t m = 24; + + SECTION("fp32 cosine, non-pow2 2*M") { + knowhere::Json hnsw_json; + hnsw_json[knowhere::meta::DIM] = dim; + hnsw_json[knowhere::meta::METRIC_TYPE] = knowhere::metric::COSINE; + hnsw_json[knowhere::meta::TOPK] = 10; + hnsw_json[knowhere::indexparam::HNSW_M] = m; + 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(); + REQUIRE(cpu_idx.Build(train_ds, hnsw_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, 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.80f); + } + + SECTION("int8 cosine DP4A, non-pow2 2*M") { + knowhere::Json hnsw_json; + hnsw_json[knowhere::meta::DIM] = dim; + hnsw_json[knowhere::meta::METRIC_TYPE] = knowhere::metric::COSINE; + hnsw_json[knowhere::meta::TOPK] = 10; + hnsw_json[knowhere::indexparam::HNSW_M] = m; + 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(); + REQUIRE(cpu_idx.Build(train_ds, hnsw_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, 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.80f); + } + } } TEST_CASE("Test CPU vs GPU HNSW Comparison", "[gpu_hnsw_compare]") { diff --git a/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.h b/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.h index 05be2b9c4..b52018a32 100644 --- a/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.h +++ b/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.h @@ -30,11 +30,13 @@ #include namespace faiss { + namespace cppcontrib { namespace knowhere { struct IndexHNSW; } // namespace knowhere } // namespace cppcontrib + namespace gpu { struct GpuHnswDeviceIndex; @@ -93,7 +95,10 @@ struct GpuIndexHNSW : public GpuIndex { /// dequantized for other SQ types). void copyFrom(const faiss::cppcontrib::knowhere::IndexHNSW* index); - /// Load with explicit metric specification. + /// Like copyFrom(), but with the metric interpretation supplied by the + /// caller instead of being detected from the index type. + /// \param use_ip treat the metric as inner product + /// \param is_cosine the storage carries cosine (inverse L2) norms void copyFromWithMetric( const faiss::cppcontrib::knowhere::IndexHNSW* index, bool use_ip, @@ -115,6 +120,11 @@ struct GpuIndexHNSW : public GpuIndex { /// 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. + /// + /// Distance convention: for inner-product and cosine metrics the returned + /// distances are the *negated* dot product (smaller == more similar), so + /// callers that want raw similarity scores must negate them (Knowhere does + /// this on the way out). L2 distances are returned as-is. void searchHost( idx_t n, const float* x_host, @@ -124,9 +134,12 @@ struct GpuIndexHNSW : public GpuIndex { 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. + /// The queries are the caller's signed int8 values and are uploaded + /// verbatim — no bias/shift is applied, matching the dataset upload which + /// already reverses FAISS's +128 SQ bias. When dim % 4 != 0 (DP4A requires + /// groups of four int8 lanes) this transparently falls back to the fp32 + /// searchHost() path. Same negated IP/cosine distance convention as + /// searchHost(). All input/output pointers must be host memory. void searchHostInt8( idx_t n, const int8_t* x_host, diff --git a/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearch.cuh b/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearch.cuh index 4422c20f0..f97d93ae4 100644 --- a/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearch.cuh +++ b/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearch.cuh @@ -66,9 +66,13 @@ inline void gpu_hnsw_search( int dim = static_cast(idx.dim); int num_upper_layers = idx.num_upper_layers_built; - auto launch_kernels = [&]( + // Layer-0 is templated on the dataset type (DataT), the layer-0 query type + // (QueryT: float generic, int8_t for the native DP4A path) and USE_DP4A. + // The upper-layer greedy descent always uses the fp32 queries (sc.d_queries). + auto launch_kernels = [&]( const DataT* d_data, - const float* d_inv_norms) { + const float* d_inv_norms, + const QueryT* d_layer0_queries) { if (num_upper_layers > 0) { auto* d_layer_ptrs = static_cast( idx.d_upper_layer_ptrs); @@ -126,12 +130,28 @@ 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). + // The bitonic sort in the parallel merge needs a power-of-two + // staging capacity, one thread per slot. Pad up to the next power + // of two (handles non-power-of-two 2*M) and grow the block so every + // staging slot is owned by a thread. + int max_staging = hnsw_kernel::padded_staging_capacity( + sw, idx.max_degree0); + if (block_size < max_staging) { + block_size = max_staging; + } + if (block_size > 1024) { + throw std::runtime_error( + std::string("gpu_hnsw: padded staging capacity ") + + std::to_string(max_staging) + + " exceeds the max CUDA block size (1024); reduce " + "search_width or M (max_degree0=" + + std::to_string(idx.max_degree0) + ")"); + } + // Fixed overhead: staging (max_staging*8) + parent_ids (sw*4) + // + meta (12) + int smem_overhead = max_staging * 8 + sw * 4 + 12; + // Per-ef cost: 3 result arrays + 3 merge arrays = 6 × 4 = 24 bytes/slot + int max_ef = (smem_max - smem_overhead) / 24; if (max_ef < 1) { throw std::runtime_error( std::string("gpu_hnsw: search_width=") + @@ -141,6 +161,12 @@ inline void gpu_hnsw_search( " bytes); reduce search_width"); } if (ef > max_ef) { + fprintf(stderr, + "[gpu_hnsw] warning: ef=%d exceeds the per-block " + "shared-memory budget (%d bytes); clamping ef to %d. " + "Recall may be reduced; raise thread_block_size or " + "lower search_width to restore the requested ef.\n", + ef, smem_max, max_ef); ef = max_ef; } } @@ -152,7 +178,8 @@ inline void gpu_hnsw_search( // 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, + hnsw_kernel::layer0_beam_search_kernel< + DataT, QueryT, USE_DP4A>, cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast(smem_size))); } @@ -164,9 +191,9 @@ inline void gpu_hnsw_search( GPU_HNSW_CUDA_CHECK( cudaMemsetAsync(sc.d_visited_bitmaps, 0, bitmap_bytes, stream)); - hnsw_kernel::layer0_beam_search_kernel + hnsw_kernel::layer0_beam_search_kernel <<>>( - sc.d_queries, + d_layer0_queries, d_data, d_inv_norms, idx.d_layer0_graph, @@ -188,130 +215,39 @@ inline void gpu_hnsw_search( switch (idx.dataset_type) { case GpuHnswDatasetType::INT8: - // 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( - 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_i8 = 49152; - { - int device = 0; - int optin = 0; - if (cudaGetDevice(&device) == cudaSuccess && - cudaDeviceGetAttribute( - &optin, - cudaDevAttrMaxSharedMemoryPerBlockOptin, - device) == cudaSuccess && - optin > smem_max_i8) { - smem_max_i8 = optin; - } - } - - { - int smem_overhead = sw * idx.max_degree0 * 8 + sw * 4 + 12; - int max_ef = (smem_max_i8 - smem_overhead) / 12; - if (max_ef < 1) { - throw std::runtime_error( - std::string("gpu_hnsw_int8: search_width=") + - std::to_string(sw) + - " too large for device shared memory (" + - std::to_string(smem_max_i8) + - " bytes); reduce search_width"); - } - if (ef > max_ef) { - ef = max_ef; - } - } - - size_t smem_size_i8 = hnsw_kernel::calc_layer0_smem_size( - ef, sw, idx.max_degree0); - - if (smem_size_i8 > 49152) { - GPU_HNSW_CUDA_CHECK(cudaFuncSetAttribute( - hnsw_kernel::layer0_beam_search_kernel_int8, - cudaFuncAttributeMaxDynamicSharedMemorySize, - static_cast(smem_size_i8))); - } - - 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; + // Native DP4A path: requires int8 queries staged on device, an + // inner-product/cosine metric (DP4A only computes dot products) and + // dim % 4 == 0. Otherwise fall back to the generic fp32-query path. + if (sc.d_queries_i8 != nullptr && idx.use_ip && (dim % 4 == 0)) { + launch_kernels.template operator()( + static_cast(idx.d_dataset), + idx.d_inv_norms, + sc.d_queries_i8); + } else { + launch_kernels.template operator()( + static_cast(idx.d_dataset), + idx.d_inv_norms, + sc.d_queries); } - // Fall through to generic int8 fp32-query path if no i8 queries. - 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); + launch_kernels.template operator()( + static_cast(idx.d_dataset), + idx.d_inv_norms, + sc.d_queries); break; case GpuHnswDatasetType::BF16: - launch_kernels( + launch_kernels.template operator()<__nv_bfloat16, float, false>( static_cast(idx.d_dataset), - idx.d_inv_norms); + idx.d_inv_norms, + sc.d_queries); break; case GpuHnswDatasetType::FP32: default: - launch_kernels( - static_cast(idx.d_dataset), idx.d_inv_norms); + launch_kernels.template operator()( + static_cast(idx.d_dataset), + idx.d_inv_norms, + sc.d_queries); break; } } diff --git a/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearchKernel.cuh b/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearchKernel.cuh index c0a3ae471..620784889 100644 --- a/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearchKernel.cuh +++ b/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearchKernel.cuh @@ -78,12 +78,10 @@ __device__ __forceinline__ float thread_ip_distance( return -sum; } -// DP4A inner-product distance for int8 queries vs 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). +// Native int8 inner-product distance using DP4A (4 int8 MADs/cycle, SM_61+). +// Both query and vec must point at int8 data reinterpreted as packed int32, +// i.e. dim4 == dim / 4 (dim must be a multiple of 4). Returns the negated dot +// product to match thread_ip_distance()'s convention (smaller == closer). __device__ __forceinline__ float thread_ip_distance_dp4a( const int32_t* __restrict__ query_packed, const int32_t* __restrict__ vec_packed, @@ -93,7 +91,58 @@ __device__ __forceinline__ float thread_ip_distance_dp4a( for (int d = 0; d < dim4; d++) { sum = __dp4a(query_packed[d], vec_packed[d], sum); } - return static_cast(-sum); // negate: max-IP convention + return static_cast(-sum); +} + +// Smallest power of two >= x (x >= 1). Used to pad the staging capacity so the +// bitonic sort (which requires a power-of-two capacity) never rejects graphs +// whose search_width * max_degree0 is not already a power of two. +__host__ __device__ __forceinline__ int next_pow2_int(int x) { + int p = 1; + while (p < x) + p <<= 1; + return p; +} + +// Padded staging capacity: power-of-two >= search_width * max_degree0. +__host__ __device__ __forceinline__ int padded_staging_capacity( + int search_width, + int max_degree0) { + return next_pow2_int(search_width * max_degree0); +} + +// Unified layer-0 distance: DP4A for native int8+IP, generic fp32 path +// otherwise. The branch is resolved at compile time so only the relevant +// distance function is instantiated for each specialization. +template +__device__ __forceinline__ float layer0_distance( + const QueryT* __restrict__ query, + const DataT* __restrict__ d_dataset, + const float* __restrict__ d_inv_norms, + uint32_t id, + int dim, + int dim4, + bool use_inner_product) { + if constexpr (USE_DP4A) { + const int32_t* vec_packed = reinterpret_cast( + d_dataset + static_cast(id) * dim); + float dist = thread_ip_distance_dp4a( + reinterpret_cast(query), vec_packed, dim4); + if (d_inv_norms) + dist *= __ldg(&d_inv_norms[id]); + return dist; + } else { + const DataT* vec = d_dataset + static_cast(id) * dim; + float dist; + if (use_inner_product) { + dist = thread_ip_distance(query, vec, dim); + if (d_inv_norms) + dist *= __ldg(&d_inv_norms[id]); + } else { + dist = thread_l2_distance(query, vec, dim); + } + return dist; + } } // ============================================================================ @@ -224,288 +273,150 @@ __global__ void upper_layer_search_kernel( } // ============================================================================ -// Phase 2: Layer-0 parallel beam search kernel +// Phase 1: Parallel merge helpers (bitonic sort + parallel merge) // ============================================================================ -__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; -} - -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 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; - - 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; +__device__ __forceinline__ void bitonic_sort_staging( + uint32_t* ids, float* dists, int active_count, int capacity) { + int i = threadIdx.x; + if (i < capacity && i >= active_count) { + ids[i] = UINT32_MAX; + dists[i] = FLT_MAX; } __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); + for (int k = 2; k <= capacity; k <<= 1) { + for (int j = k >> 1; j > 0; j >>= 1) { + int partner = i ^ j; + if (i < capacity && partner < capacity) { + // Each thread reads both its own slot and its partner's slot, + // then writes only to slot i — no thread writes to partner. + // Both threads in a pair do this symmetrically, so reads and + // writes of slot i are exclusively owned by thread i. + float di = dists[i], dp = dists[partner]; + uint32_t ii = ids[i], ip = ids[partner]; + // Ascending sort when the direction bit (i & k) is 0. + // In ascending order: lower index should hold the smaller value. + bool ascending = ((i & k) == 0); + bool i_is_lower = (i < partner); + // i should hold the smaller value iff: ascending and ipartner. + bool want_smaller = ascending ? i_is_lower : !i_is_lower; + bool i_is_smaller = (di < dp) || (di == dp && ii < ip); + if (want_smaller != i_is_smaller) { + dists[i] = dp; + ids[i] = ip; + } + } + __syncthreads(); } - 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; +__device__ __forceinline__ void parallel_merge_into_result( + uint32_t* result_ids, + float* result_dists, + uint32_t* is_expanded, + uint32_t* staging_ids, + float* staging_dists, + uint32_t* merged_ids, + float* merged_dists, + uint32_t* merged_expanded, + int* meta, + int ef, + int max_staging) { + int sc = min(meta[1], max_staging); + int rc = meta[0]; - 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); + bitonic_sort_staging(staging_ids, staging_dists, sc, max_staging); + + int T = rc + sc; + int merge_limit = min(T, ef); + + for (int pos = threadIdx.x; pos < merge_limit; pos += blockDim.x) { + // Binary search for si: number of staging elements ranked before pos. + // Valid range: si in [max(0,pos-rc), min(pos,sc)]. + // Half-open: lo inclusive, hi exclusive = min(pos,sc)+1. + // ri = pos-si is always >= 0 within this range, but after exit si + // may equal min(pos,sc)+1 (all staging ranked first); guard ri reads. + int lo = max(0, pos - rc); + int hi = min(pos, sc) + 1; + + while (lo < hi) { + int mid = (lo + hi) / 2; + // Clamp ri_mid to 0 to prevent negative-index speculative loads + // under compiler optimization; the predicate discards the value. + int ri_mid = max(0, pos - mid); + float sv = (mid < sc) ? staging_dists[mid] : FLT_MAX; + float rv = (ri_mid < rc) ? result_dists[ri_mid] : FLT_MAX; + uint32_t si_id = (mid < sc) ? staging_ids[mid] : UINT32_MAX; + uint32_t ri_id = (ri_mid < rc) ? result_ids[ri_mid] : UINT32_MAX; + // Override with sentinel if the true ri (pos-mid) would be negative + if (pos < mid) { rv = FLT_MAX; ri_id = UINT32_MAX; } + bool sv_le = (sv < rv) || (sv == rv && si_id < ri_id); + if (sv_le) + lo = mid + 1; + else + hi = mid; } - int slot = atomicAdd(&meta[1], 1); - if (slot < max_staging) { - staging_ids[slot] = nbr; - staging_dists[slot] = dist; + int si = lo; + // Clamp ri to 0 to prevent negative-index speculative loads; + // the ri_valid predicate gates whether the value is actually used. + int ri = pos - si; + bool ri_valid = (ri >= 0 && ri < rc); + int ri_safe = max(0, ri); + + float sv = (si < sc) ? staging_dists[si] : FLT_MAX; + float rv = ri_valid ? result_dists[ri_safe] : FLT_MAX; + uint32_t si_id = (si < sc) ? staging_ids[si] : UINT32_MAX; + uint32_t ri_id = ri_valid ? result_ids[ri_safe] : UINT32_MAX; + bool take_staging = (si < sc) && (!ri_valid || (sv < rv) || (sv == rv && si_id < ri_id)); + + if (take_staging) { + merged_ids[pos] = si_id; + merged_dists[pos] = sv; + merged_expanded[pos] = 0; + } else { + merged_ids[pos] = ri_id; + merged_dists[pos] = rv; + merged_expanded[pos] = ri_valid ? is_expanded[ri_safe] : 0; } } __syncthreads(); + int new_rc = min(T, ef); + for (int i = threadIdx.x; i < new_rc; i += blockDim.x) { + result_ids[i] = merged_ids[i]; + result_dists[i] = merged_dists[i]; + is_expanded[i] = merged_expanded[i]; + } 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; + meta[0] = new_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 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; - } - } } // ============================================================================ -// 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 (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). +// Phase 2: Layer-0 parallel beam search kernel // ============================================================================ -__global__ void layer0_beam_search_kernel_int8( - const int32_t* __restrict__ d_queries_packed, - const int8_t* __restrict__ d_dataset, + +__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; +} + +// Templated on the dataset element type (DataT), the query element type +// (QueryT: float for the generic path, int8_t for the native DP4A path), and +// USE_DP4A which selects the DP4A int8 inner-product distance at compile time. +template +__global__ void layer0_beam_search_kernel( + const QueryT* __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, @@ -525,12 +436,14 @@ __global__ void layer0_beam_search_kernel_int8( if (query_idx >= num_queries) return; + const QueryT* query = d_queries + static_cast(query_idx) * dim; 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; + // Pad to a power of two so the bitonic sort in parallel_merge_into_result + // has a valid capacity for any max_degree0 (e.g. non-power-of-two 2*M). + int max_staging = padded_staging_capacity(search_width, max_degree0); uint32_t* result_ids = reinterpret_cast(smem); float* result_dists = reinterpret_cast(result_ids + ef); @@ -540,6 +453,9 @@ __global__ void layer0_beam_search_kernel_int8( uint32_t* parent_ids = reinterpret_cast(staging_dists + max_staging); int* meta = reinterpret_cast(parent_ids + search_width); + uint32_t* merged_ids = reinterpret_cast(meta + 3); + float* merged_dists = reinterpret_cast(merged_ids + ef); + uint32_t* merged_expanded = reinterpret_cast(merged_dists + ef); int bitmap_words = (N + 31) / 32; uint32_t* visited_bmap = @@ -560,13 +476,9 @@ __global__ void layer0_beam_search_kernel_int8( // --- 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); - // 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 (d_inv_norms) - ep_dist *= __ldg(&d_inv_norms[ep]); + float ep_dist = layer0_distance( + query, d_dataset, d_inv_norms, ep, dim, dim4, + use_inner_product); result_ids[0] = ep; result_dists[0] = ep_dist; is_expanded[0] = 0; @@ -588,11 +500,9 @@ __global__ void layer0_beam_search_kernel_int8( 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]; + float dist = layer0_distance( + query, d_dataset, d_inv_norms, nbr, dim, dim4, + use_inner_product); int slot = atomicAdd(&meta[1], 1); if (slot < max_staging) { @@ -602,44 +512,19 @@ __global__ void layer0_beam_search_kernel_int8( } __syncthreads(); + parallel_merge_into_result( + result_ids, result_dists, is_expanded, + staging_ids, staging_dists, + merged_ids, merged_dists, merged_expanded, + meta, ef, max_staging); 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(); @@ -683,11 +568,9 @@ __global__ void layer0_beam_search_kernel_int8( 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]; + float dist = layer0_distance( + query, d_dataset, d_inv_norms, nbr, dim, dim4, + use_inner_product); int slot = atomicAdd(&meta[1], 1); if (slot < max_staging) { @@ -697,40 +580,13 @@ __global__ void layer0_beam_search_kernel_int8( } __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++; - } + parallel_merge_into_result( + result_ids, result_dists, is_expanded, + staging_ids, staging_dists, + merged_ids, merged_dists, merged_expanded, + meta, ef, max_staging); - meta[0] = rc; - } - __syncthreads(); + // Stagnation detected when num_parents==0 (all expanded) → break above } // --- Copy top-k results to global memory --- @@ -749,16 +605,20 @@ __global__ void layer0_beam_search_kernel_int8( } inline size_t calc_layer0_smem_size(int ef, int search_width, int max_degree0) { - int max_staging = search_width * max_degree0; + // Must match the kernel's padded staging capacity exactly. + int max_staging = padded_staging_capacity(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); + size += ef * sizeof(uint32_t); // result_ids + size += ef * sizeof(float); // result_dists + size += ef * sizeof(uint32_t); // is_expanded + size += max_staging * sizeof(uint32_t); // staging_ids + size += max_staging * sizeof(float); // staging_dists + size += search_width * sizeof(uint32_t); // parent_ids + size += 3 * sizeof(int); // meta + size += ef * sizeof(uint32_t); // merged_ids + size += ef * sizeof(float); // merged_dists + size += ef * sizeof(uint32_t); // merged_expanded return size; } diff --git a/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.h b/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.h index 1f15a3cca..f0d585c62 100644 --- a/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.h +++ b/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.h @@ -96,7 +96,7 @@ 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) + int8_t* d_queries_i8 = nullptr; // int8 queries for the native DP4A path size_t queries_bytes = 0; size_t neighbors_bytes = 0; From 33684d30ecc76969cbb1e24885d281a2ccb1149d Mon Sep 17 00:00:00 2001 From: premal Date: Fri, 17 Jul 2026 03:53:47 +0000 Subject: [PATCH 26/37] fix(gpu_hnsw): re-vendor faiss merge-path off-by-one fix Re-vendors faiss gpu-hnsw-faiss@81c9ff25 GpuHnswSearchKernel.cuh. Corrects the parallel_merge_into_result diagonal binary search (staging[mid] vs result[pos-mid-1], not result[pos-mid]) that caused ~4% recall and a UINT32_MAX-sentinel CUDA illegal memory access on int8/COSINE search. Validated on CPU vs std::merge (200k cases, 0 failures). All 7 vendored faiss/gpu files remain byte-identical to faiss. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../faiss/gpu/impl/GpuHnswSearchKernel.cuh | 31 ++++++++++++++----- 1 file changed, 24 insertions(+), 7 deletions(-) diff --git a/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearchKernel.cuh b/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearchKernel.cuh index 620784889..6cf12fd28 100644 --- a/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearchKernel.cuh +++ b/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearchKernel.cuh @@ -344,15 +344,32 @@ __device__ __forceinline__ void parallel_merge_into_result( while (lo < hi) { int mid = (lo + hi) / 2; - // Clamp ri_mid to 0 to prevent negative-index speculative loads - // under compiler optimization; the predicate discards the value. - int ri_mid = max(0, pos - mid); + // Merge-path diagonal test: staging[mid] is taken ahead of the + // result element it would displace, which is result[pos-mid-1] + // (NOT result[pos-mid]). Getting this index wrong mis-partitions + // every merge (empirically ~4% recall) and can promote a padded + // UINT32_MAX sentinel into the result list, which is later used as + // a graph/dataset index -> illegal memory access. + int ri_mid = pos - mid - 1; + bool ri_mid_valid = (ri_mid >= 0 && ri_mid < rc); + int ri_mid_safe = max(0, ri_mid); float sv = (mid < sc) ? staging_dists[mid] : FLT_MAX; - float rv = (ri_mid < rc) ? result_dists[ri_mid] : FLT_MAX; uint32_t si_id = (mid < sc) ? staging_ids[mid] : UINT32_MAX; - uint32_t ri_id = (ri_mid < rc) ? result_ids[ri_mid] : UINT32_MAX; - // Override with sentinel if the true ri (pos-mid) would be negative - if (pos < mid) { rv = FLT_MAX; ri_id = UINT32_MAX; } + float rv; + uint32_t ri_id; + if (ri_mid < 0) { + // No result element precedes this split: staging[mid] must not + // be taken ahead of it -> compare against -inf (predicate false). + rv = -FLT_MAX; + ri_id = 0u; + } else if (ri_mid_valid) { + rv = result_dists[ri_mid_safe]; + ri_id = result_ids[ri_mid_safe]; + } else { + // ri_mid >= rc: no result element there -> +inf (take staging). + rv = FLT_MAX; + ri_id = UINT32_MAX; + } bool sv_le = (sv < rv) || (sv == rv && si_id < ri_id); if (sv_le) lo = mid + 1; From 5b8baf9883a06ca6866504f17f0d9fef5b29894c Mon Sep 17 00:00:00 2001 From: premal Date: Fri, 17 Jul 2026 05:24:07 +0000 Subject: [PATCH 27/37] fix(gpu_hnsw): re-vendor node-id OOB guards + add full dtype x metric tests Re-vendor faiss GpuHnswSearchKernel.cuh (faiss gpu-hnsw-faiss @46235ee3): guard ep/parent/result node ids < N before graph/dataset indexing (fixes v82 fp16 P1 OOB in layer0_beam_search_kernel<__half,float,false>). Add independently-taggable per-dtype GPU HNSW search tests (nb=4000, dim=64, M=16, ef=200 -- same config as the failing fp16 P1 case), each covering L2/IP/COSINE so a CUDA context crash in one instantiation does not hide the others: - [gpu_hnsw_dtype_fp32] flat -> (config control) - [gpu_hnsw_dtype_fp16] HNSW_SQ fp16 -> <__half,float,false> - [gpu_hnsw_dtype_bf16] HNSW_SQ bf16 -> <__nv_bfloat16,float,false> - [gpu_hnsw_dtype_sq8] HNSW_SQ sq8 -> decoded - [gpu_hnsw_dtype_int8] native int8 (IP/COSINE take DP4A, dim%4==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> --- tests/ut/test_gpu_search.cc | 118 ++++++++++++++++++ .../faiss/gpu/impl/GpuHnswSearchKernel.cuh | 37 ++++-- 2 files changed, 145 insertions(+), 10 deletions(-) diff --git a/tests/ut/test_gpu_search.cc b/tests/ut/test_gpu_search.cc index 62fc13921..63f19c2f3 100644 --- a/tests/ut/test_gpu_search.cc +++ b/tests/ut/test_gpu_search.cc @@ -1486,4 +1486,122 @@ TEST_CASE("Test GPU HNSW Codex P1 Regressions", "[gpu_hnsw_p1]") { REQUIRE(GetKNNRecall(*gt.value(), *r_ok.value()) >= 0.8f); } } + +// ============================================================================ +// Full dtype x metric coverage for the GPU HNSW layer-0 search kernel. +// +// One TEST_CASE per dtype so each runs in its OWN process invocation (a CUDA +// illegal-access in one kernel instantiation corrupts the context and aborts +// the process; separate tags let the harness run every dtype independently and +// still collect per-dtype results instead of losing everything after the first +// crash). Every dtype is exercised for L2, IP and COSINE. +// +// All cases use the SAME config as the fp16 P1 regression that v82 OOBs on +// (nb=4000, dim=64, M=16, ef=200) so the fp32 case below is the control: if +// fp32 (the kernel) is clean here while fp16 +// (<__half,float,false>) faults, the defect is specific to the low-precision +// storage kernel; if fp32 also faults, it is config- rather than dtype-driven. +// dim=64 is divisible by 4, so the int8 IP/COSINE cases take the DP4A path. +namespace { +constexpr int64_t kDtypeNb = 4000; +constexpr int64_t kDtypeNq = 200; +constexpr int64_t kDtypeDim = 64; +constexpr int64_t kDtypeSeed = 42; +constexpr int64_t kDtypeTopk = 10; + +knowhere::Json +dtype_json(const std::string& metric, const std::string& sq_type) { + knowhere::Json json; + json[knowhere::meta::DIM] = kDtypeDim; + json[knowhere::meta::METRIC_TYPE] = metric; + json[knowhere::meta::TOPK] = kDtypeTopk; + json[knowhere::indexparam::HNSW_M] = 16; + json[knowhere::indexparam::EFCONSTRUCTION] = 200; + json[knowhere::indexparam::EF] = 200; + if (!sq_type.empty()) { + json[knowhere::indexparam::SQ_TYPE] = sq_type; + } + return json; +} + +// Build a CPU index of `index_type` on data type T, serialize, deserialize as +// GPU_HNSW, search, and return recall vs a brute-force reference on T. Shared by +// every per-dtype case so the only variables are the data type, the storage +// index type and the metric. +template +float +gpu_dtype_recall(const std::string& index_type, const knowhere::Json& json, int version) { + auto train_ds = knowhere::ConvertToDataTypeIfNeeded(GenDataSet(kDtypeNb, kDtypeDim, kDtypeSeed)); + auto query_ds = knowhere::ConvertToDataTypeIfNeeded(GenDataSet(kDtypeNq, kDtypeDim, kDtypeSeed + 2)); + + auto cpu_idx = knowhere::IndexFactory::Instance().Create(index_type, version).value(); + REQUIRE(cpu_idx.Build(train_ds, json) == knowhere::Status::success); + knowhere::BinarySet bs; + REQUIRE(cpu_idx.Serialize(bs) == knowhere::Status::success); + + 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()); + return GetKNNRecall(*gt.value(), *results.value()); +} +} // namespace + +// FP32 flat storage -> generic kernel. The config control. +TEST_CASE("GPU HNSW dtype: FP32 (L2/IP/COSINE)", "[gpu_hnsw_dtype_fp32]") { + const auto version = GenTestVersionList(); + const auto metric = GENERATE(knowhere::metric::L2, knowhere::metric::IP, knowhere::metric::COSINE); + CAPTURE(metric); + float recall = gpu_dtype_recall(knowhere::IndexEnum::INDEX_HNSW, dtype_json(metric, ""), version); + REQUIRE(recall >= 0.85f); +} + +// FP16 native 2-byte storage -> <__half,float,false> kernel. Built via HNSW_SQ +// on the fp32 node with SQ_TYPE=fp16 so the device selects the native half +// representation (the path the mock-wrapper fp16 tests do NOT reach). +TEST_CASE("GPU HNSW dtype: FP16 native (L2/IP/COSINE)", "[gpu_hnsw_dtype_fp16]") { + const auto version = GenTestVersionList(); + const auto metric = GENERATE(knowhere::metric::L2, knowhere::metric::IP, knowhere::metric::COSINE); + CAPTURE(metric); + float recall = + gpu_dtype_recall(knowhere::IndexEnum::INDEX_HNSW_SQ, dtype_json(metric, "fp16"), version); + REQUIRE(recall >= 0.85f); +} + +// BF16 native 2-byte storage -> <__nv_bfloat16,float,false> kernel. +TEST_CASE("GPU HNSW dtype: BF16 native (L2/IP/COSINE)", "[gpu_hnsw_dtype_bf16]") { + const auto version = GenTestVersionList(); + const auto metric = GENERATE(knowhere::metric::L2, knowhere::metric::IP, knowhere::metric::COSINE); + CAPTURE(metric); + float recall = + gpu_dtype_recall(knowhere::IndexEnum::INDEX_HNSW_SQ, dtype_json(metric, "bf16"), version); + // bf16 has a 7-bit mantissa (coarser than fp16); slightly looser floor. + REQUIRE(recall >= 0.80f); +} + +// SQ8 storage -> decoded to fp32 on upload -> generic +// kernel (distinct from the native half/bf16 paths above). +TEST_CASE("GPU HNSW dtype: SQ8 decoded (L2/IP/COSINE)", "[gpu_hnsw_dtype_sq8]") { + const auto version = GenTestVersionList(); + const auto metric = GENERATE(knowhere::metric::L2, knowhere::metric::IP, knowhere::metric::COSINE); + CAPTURE(metric); + float recall = + gpu_dtype_recall(knowhere::IndexEnum::INDEX_HNSW_SQ, dtype_json(metric, "sq8"), version); + // SQ8 is coarse (8-bit per component); conservative floor. + REQUIRE(recall >= 0.65f); +} + +// INT8 native 1-byte storage. IP/COSINE with dim%4==0 take the DP4A kernel +// (); L2 takes the generic decoded path (). +TEST_CASE("GPU HNSW dtype: INT8 native (L2/IP/COSINE)", "[gpu_hnsw_dtype_int8]") { + const auto version = GenTestVersionList(); + const auto metric = GENERATE(knowhere::metric::L2, knowhere::metric::IP, knowhere::metric::COSINE); + CAPTURE(metric); + float recall = gpu_dtype_recall(knowhere::IndexEnum::INDEX_HNSW, dtype_json(metric, ""), version); + REQUIRE(recall >= 0.85f); +} #endif diff --git a/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearchKernel.cuh b/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearchKernel.cuh index 6cf12fd28..774c4bcf9 100644 --- a/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearchKernel.cuh +++ b/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearchKernel.cuh @@ -492,15 +492,24 @@ __global__ void layer0_beam_search_kernel( // --- Seed with entry point --- uint32_t ep = d_entry_points[query_idx]; + // Defensive: the entry point must be a real node (< N) before it indexes + // the dataset (ep_dist) or the layer-0 graph (neighbor seeding). A bad ep + // would otherwise fault far out of bounds. On the healthy path ep is always + // valid; if it is not, seed nothing and let the search return sentinels. + bool ep_valid = (ep < static_cast(N)); if (threadIdx.x == 0) { - float ep_dist = layer0_distance( - query, d_dataset, d_inv_norms, ep, dim, dim4, - use_inner_product); - result_ids[0] = ep; - result_dists[0] = ep_dist; - is_expanded[0] = 0; - meta[0] = 1; - bitmap_visit(visited_bmap, ep); + if (ep_valid) { + float ep_dist = layer0_distance( + query, d_dataset, d_inv_norms, ep, dim, dim4, + use_inner_product); + result_ids[0] = ep; + result_dists[0] = ep_dist; + is_expanded[0] = 0; + meta[0] = 1; + bitmap_visit(visited_bmap, ep); + } else { + meta[0] = 0; + } } __syncthreads(); @@ -509,7 +518,7 @@ __global__ void layer0_beam_search_kernel( meta[1] = 0; __syncthreads(); - for (int j = threadIdx.x; j < max_degree0; j += blockDim.x) { + for (int j = threadIdx.x; ep_valid && 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)) @@ -576,6 +585,14 @@ __global__ void layer0_beam_search_kernel( int nbr_slot = wi % max_degree0; uint32_t parent = parent_ids[parent_idx]; + // Defensive: a parent id must be a real node (< N). result_ids only + // ever holds the entry point or graph neighbors (both < N) once the + // merge is correct, so this never fires on the healthy path; it caps + // a stray/garbage id at O(N) instead of letting it index the graph + // at parent*max_degree0 and fault far out of bounds. Mirrors the + // neighbor guard below. + if (parent >= static_cast(N)) + continue; uint32_t nbr = d_layer0_graph [static_cast(parent) * max_degree0 + @@ -609,7 +626,7 @@ __global__ void layer0_beam_search_kernel( // --- Copy top-k results to global memory --- int rc = meta[0]; for (int i = threadIdx.x; i < k; i += blockDim.x) { - if (i < rc) { + if (i < rc && result_ids[i] < static_cast(N)) { d_neighbors[static_cast(query_idx) * k + i] = static_cast(result_ids[i]); d_distances[static_cast(query_idx) * k + i] = From 8fbaf1f4746c54bb71a90865e1d45f081426533b Mon Sep 17 00:00:00 2001 From: premal Date: Fri, 17 Jul 2026 05:33:05 +0000 Subject: [PATCH 28/37] fix(gpu_hnsw): re-vendor bitonic_sort_staging shared-memory race fix Re-vendor faiss GpuHnswSearchKernel.cuh (faiss gpu-hnsw-faiss @81e4cc69): split the bitonic compare-exchange into read -> __syncthreads -> write -> __syncthreads phases, removing the cross-warp read/write hazard that produced the v82 fp16 P1 OOB. Vendored file byte-identical to faiss. 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> --- .../faiss/gpu/impl/GpuHnswSearchKernel.cuh | 34 +++++++++++++------ 1 file changed, 24 insertions(+), 10 deletions(-) diff --git a/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearchKernel.cuh b/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearchKernel.cuh index 774c4bcf9..dd0f5ede5 100644 --- a/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearchKernel.cuh +++ b/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearchKernel.cuh @@ -288,13 +288,25 @@ __device__ __forceinline__ void bitonic_sort_staging( for (int k = 2; k <= capacity; k <<= 1) { for (int j = k >> 1; j > 0; j >>= 1) { int partner = i ^ j; + // Compare-exchange in two barrier-separated phases. Each thread + // owns writes to slot i only, but it READS slot `partner` — which + // another thread owns and may overwrite in the SAME phase. When the + // partner is in a different warp (j >= warpSize) the read and that + // write are unsynchronized -> a shared-memory hazard (racecheck + // flagged ~1.6M in the half kernel). So: (1) every thread reads its + // pair into registers and decides whether to swap, (2) __syncthreads + // so all reads complete before any write, (3) threads that must swap + // write their own slot, (4) __syncthreads before the next phase's + // reads. Loop bounds (k, j) are uniform across the block, so every + // thread — including i >= capacity — reaches every barrier. + float dp = 0.0f; + uint32_t ip = 0u; + bool do_swap = false; if (i < capacity && partner < capacity) { - // Each thread reads both its own slot and its partner's slot, - // then writes only to slot i — no thread writes to partner. - // Both threads in a pair do this symmetrically, so reads and - // writes of slot i are exclusively owned by thread i. - float di = dists[i], dp = dists[partner]; - uint32_t ii = ids[i], ip = ids[partner]; + float di = dists[i]; + dp = dists[partner]; + uint32_t ii = ids[i]; + ip = ids[partner]; // Ascending sort when the direction bit (i & k) is 0. // In ascending order: lower index should hold the smaller value. bool ascending = ((i & k) == 0); @@ -303,10 +315,12 @@ __device__ __forceinline__ void bitonic_sort_staging( // or descending and i>partner. bool want_smaller = ascending ? i_is_lower : !i_is_lower; bool i_is_smaller = (di < dp) || (di == dp && ii < ip); - if (want_smaller != i_is_smaller) { - dists[i] = dp; - ids[i] = ip; - } + do_swap = (want_smaller != i_is_smaller); + } + __syncthreads(); + if (do_swap) { + dists[i] = dp; + ids[i] = ip; } __syncthreads(); } From fcdcacf21f53c1f1fa7f6afab30a68bdfbf92738 Mon Sep 17 00:00:00 2001 From: premal Date: Fri, 17 Jul 2026 14:55:53 +0000 Subject: [PATCH 29/37] fix(gpu_hnsw): SQ deserialize alias, lazy metric detection, drop IP/COSINE re-negation; re-vendor faiss Source fixes (src/index/hnsw/faiss_hnsw.cc): - GPU_HNSW_SQ deserialize: alias the incoming payload under this node's own Type() (INDEX_GPU_HNSW_SQ) so the base loader's GetByName(Type()) finds it; the previous hardcoded INDEX_GPU_HNSW alias made every GPU_HNSW_SQ segment unqueryable while leaving base GPU_HNSW behavior intact. - Lazy GPU upload now reuses UploadCpuIndexToGpuLocked() (made const) instead of deriving the metric from the query config, so IP/COSINE/L2 is detected from the FAISS index even when the query omits metric_type. - Remove the IP/COSINE sign-flip loops after searchHost()/searchHostInt8(): faiss now returns true similarity on copy-out, so re-negating would invert it. - Correct the stale native-int8 comments (no +128 shift; signed int8 uploaded verbatim, DP4A kernel). Re-vendored faiss GPU tree (thirdparty/faiss, byte-identical to faiss 48ab2cb7): GpuIndexHNSW.h, impl/GpuHnswSearch.cuh, impl/GpuHnswSearchKernel.cuh, impl/GpuHnswTypes.cu. Tests (tests/ut/test_gpu_search.cc): - [gpu_hnsw_sq_alias] load a CPU HNSW_SQ index into the GPU_HNSW_SQ node. - [gpu_hnsw_ip_sign] IP/COSINE distances are true similarity (sign+magnitude vs recomputed neighbor similarity; descending order). - [gpu_hnsw_k_valid] all k slots are real neighbors (no sentinel) with ef>=k. Signed-off-by: premal Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- src/index/hnsw/faiss_hnsw.cc | 99 ++++++------- tests/ut/test_gpu_search.cc | 140 ++++++++++++++++++ thirdparty/faiss/faiss/gpu/GpuIndexHNSW.h | 16 +- .../faiss/faiss/gpu/impl/GpuHnswSearch.cuh | 52 ++++++- .../faiss/gpu/impl/GpuHnswSearchKernel.cuh | 16 +- .../faiss/faiss/gpu/impl/GpuHnswTypes.cu | 4 + 6 files changed, 261 insertions(+), 66 deletions(-) diff --git a/src/index/hnsw/faiss_hnsw.cc b/src/index/hnsw/faiss_hnsw.cc index 9215ba6a1..3f008c645 100644 --- a/src/index/hnsw/faiss_hnsw.cc +++ b/src/index/hnsw/faiss_hnsw.cc @@ -3472,12 +3472,20 @@ class GpuHnswIndexNode : public BaseFaissRegularIndexHNSWNode { UnpublishGpuIndexLocked(); // Accept CPU-built HNSW (F32) or HNSW_SQ (quantized) binaries. + // The base loader reads the payload via binset.GetByName(Type()), so + // the incoming binary must be filed under *this* node's Type(): + // INDEX_GPU_HNSW for the base node, INDEX_GPU_HNSW_SQ for the SQ + // subclass. Aliasing under a hardcoded INDEX_GPU_HNSW would make + // GPU_HNSW_SQ loads fail — the SQ loader looks up INDEX_GPU_HNSW_SQ and + // would find nothing. Status status; - if (!binset.Contains(IndexEnum::INDEX_GPU_HNSW)) { + const std::string self_key = Type(); + if (!binset.Contains(self_key)) { BinarySet aliased = binset; - for (const char* key : {IndexEnum::INDEX_GPU_HNSW_SQ, IndexEnum::INDEX_HNSW_SQ, IndexEnum::INDEX_HNSW}) { + for (const char* key : {IndexEnum::INDEX_GPU_HNSW, 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)); + aliased.Append(self_key, binset.GetByName(key)); break; } } @@ -3537,32 +3545,25 @@ class GpuHnswIndexNode : public BaseFaissRegularIndexHNSWNode { std::unique_lock lock(gpu_mutex_); gpu_snapshot = gpu_index_; if (!gpu_snapshot) { - const auto* faiss_idx = GetFaissHnswIndex(); - if (!faiss_idx) { + if (GetFaissHnswIndex() == nullptr) { 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::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); - gpu_ntotal_ = faiss_idx->ntotal; - gpu_dim_ = faiss_idx->d; - const_cast&>(indexes[0]).reset(); - gpu_index_ = local; - gpu_ready_.store(true, std::memory_order_release); - gpu_snapshot = local; - } catch (const std::exception& e) { + // Reuse the eager upload path (it holds no lock itself and we + // already hold gpu_mutex_). This keeps metric detection + // identical to Deserialize()/DeserializeFromFile(): the metric + // is derived from the FAISS index type (dynamic_cast for cosine, + // metric_type for IP), NOT from the query config. Detecting from + // config here was a latent bug — a search whose config omits + // metric_type would default to L2 and silently run a cosine/IP + // index with the wrong metric. + Status upload_status = UploadCpuIndexToGpuLocked(); + if (upload_status != Status::success) { return expected::Err(Status::cuvs_inner_error, - std::string("failed to build GPU HNSW index: ") + e.what()); + "failed to build GPU HNSW index"); + } + gpu_snapshot = gpu_index_; + if (!gpu_snapshot) { + return expected::Err(Status::empty_index, "index not loaded"); } } } @@ -3590,9 +3591,11 @@ class GpuHnswIndexNode : public BaseFaissRegularIndexHNSWNode { 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. + // For COSINE, post-scale distances by 1/||q|| so scores are in + // [-1, 1]. The kernel returns dot(shifted_q, shifted_db) * (1/||db||) + // (the true cosine numerator scaled by the db inverse norm, already + // sign-corrected to "larger == closer"); dividing by the query norm + // finalizes the cosine similarity. 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; @@ -3607,14 +3610,10 @@ class GpuHnswIndexNode : public BaseFaissRegularIndexHNSWNode { } } - // 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]; - } - } - + // No sign flip: faiss now returns the true similarity for IP/COSINE + // (larger == closer), negating its internal min-first score on + // copy-out. Previously the kernel returned the negated score and + // this path flipped it back. return GenResultDataSet(nq, k, h_ids.release(), h_dist.release()); } @@ -3647,14 +3646,8 @@ class GpuHnswIndexNode : public BaseFaissRegularIndexHNSWNode { 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]; - } - } - + // No sign flip: faiss now returns the true similarity for IP/COSINE + // (larger == closer), negating its internal min-first score on copy-out. return GenResultDataSet(nq, k, h_ids.release(), h_dist.release()); } @@ -3678,7 +3671,7 @@ class GpuHnswIndexNode : public BaseFaissRegularIndexHNSWNode { // (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() { + UploadCpuIndexToGpuLocked() const { const auto* faiss_idx = GetFaissHnswIndex(); if (faiss_idx == nullptr) { return Status::success; @@ -3704,8 +3697,10 @@ class GpuHnswIndexNode : public BaseFaissRegularIndexHNSWNode { 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(); + // drops the transient host-RAM peak back to ~0. const_cast because + // this helper is const (reused from the const Search() lazy path); + // it only frees the now-uploaded host copy. + const_cast&>(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; @@ -3806,7 +3801,9 @@ KNOWHERE_REGISTER_GLOBAL( GPU_HNSW, [](const int32_t& version, const Object& object) { // Native int8 path: bypass MockWrapper upcast, pass int8 queries directly - // to searchHostInt8() which applies the +128 shift and uses DP4A kernel. + // to searchHostInt8(), which uploads the signed int8 queries verbatim (no + // +128 shift; the dataset upload already reverses FAISS's +128 SQ bias) + // and uses the DP4A kernel. return Index::Create(version, object, DataFormatEnum::int8); }, int8, true, (feature::INT8 | feature::GPU)); @@ -3834,7 +3831,9 @@ KNOWHERE_REGISTER_GLOBAL( GPU_HNSW_SQ, [](const int32_t& version, const Object& object) { // Native int8 path: bypass MockWrapper upcast, pass int8 queries directly - // to searchHostInt8() which applies the +128 shift and uses DP4A kernel. + // to searchHostInt8(), which uploads the signed int8 queries verbatim (no + // +128 shift; the dataset upload already reverses FAISS's +128 SQ bias) + // and uses the DP4A kernel. return Index::Create(version, object, DataFormatEnum::int8); }, int8, true, (feature::INT8 | feature::GPU)); diff --git a/tests/ut/test_gpu_search.cc b/tests/ut/test_gpu_search.cc index 63f19c2f3..2aae65d4e 100644 --- a/tests/ut/test_gpu_search.cc +++ b/tests/ut/test_gpu_search.cc @@ -13,6 +13,7 @@ #include #include #include +#include #include #include #include @@ -1604,4 +1605,143 @@ TEST_CASE("GPU HNSW dtype: INT8 native (L2/IP/COSINE)", "[gpu_hnsw_dtype_int8]") float recall = gpu_dtype_recall(knowhere::IndexEnum::INDEX_HNSW, dtype_json(metric, ""), version); REQUIRE(recall >= 0.85f); } + +// v84 regression: a CPU HNSW_SQ index must load into the *GPU_HNSW_SQ* node. +// GpuHnswSQIndexNode inherits Deserialize from GpuHnswIndexNode, which re-files +// the incoming payload under this node's own Type(). The SQ node's Type() is +// INDEX_GPU_HNSW_SQ and the base loader reads the payload via +// GetByName(Type()); the pre-v84 code aliased under a hardcoded +// INDEX_GPU_HNSW, so the SQ lookup missed and the load failed. Loading into the +// SQ node (not the base GPU_HNSW node) is what exercises the fix. +TEST_CASE("GPU HNSW_SQ node deserialize alias", "[gpu_hnsw_sq_alias]") { + const auto version = GenTestVersionList(); + const auto metric = GENERATE(knowhere::metric::L2, knowhere::metric::COSINE); + CAPTURE(metric); + + auto json = dtype_json(metric, "sq8"); + auto train_ds = GenDataSet(kDtypeNb, kDtypeDim, kDtypeSeed); + auto query_ds = GenDataSet(kDtypeNq, kDtypeDim, kDtypeSeed + 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); + + // Load into the GPU_HNSW_SQ node specifically (not the base GPU_HNSW node). + auto gpu_idx = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_GPU_HNSW_SQ, 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()); + REQUIRE(GetKNNRecall(*gt.value(), *results.value()) >= 0.65f); +} + +// v84: IP/COSINE distances must be returned as the true similarity (larger == +// closer), matching faiss's METRIC_INNER_PRODUCT contract. The GPU kernel keeps +// a negated, min-first score internally and negates it back on copy-out; +// Knowhere no longer re-negates. This checks the numeric distance value (sign + +// magnitude) against the similarity recomputed from the returned neighbor, which +// recall-only checks cannot catch. A re-introduced double negation would flip +// every score negative and fail here. +TEST_CASE("GPU HNSW IP/COSINE distance sign", "[gpu_hnsw_ip_sign]") { + const auto version = GenTestVersionList(); + const auto metric = GENERATE(knowhere::metric::IP, knowhere::metric::COSINE); + CAPTURE(metric); + const bool is_cosine = (metric == knowhere::metric::COSINE); + + auto json = dtype_json(metric, ""); + auto train_ds = GenDataSet(kDtypeNb, kDtypeDim, kDtypeSeed); + auto query_ds = GenDataSet(kDtypeNq, kDtypeDim, kDtypeSeed + 2); + const auto* train = reinterpret_cast(train_ds->GetTensor()); + const auto* query = reinterpret_cast(query_ds->GetTensor()); + + 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; + REQUIRE(cpu_idx.Serialize(bs) == knowhere::Status::success); + + 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()); + const int64_t* ids = results.value()->GetIds(); + const float* dist = results.value()->GetDistance(); + + int positive_top1 = 0; + for (int64_t q = 0; q < kDtypeNq; q++) { + // Distances must be sorted best-first, i.e. descending for a + // larger-is-closer similarity metric. + for (int64_t j = 1; j < kDtypeTopk; j++) { + REQUIRE(dist[q * kDtypeTopk + j] <= dist[q * kDtypeTopk + j - 1] + 1e-4f); + } + // Recompute the true similarity between the query and the returned + // top-1 neighbor and compare to the reported distance. + int64_t nn = ids[q * kDtypeTopk]; + REQUIRE(nn >= 0); + REQUIRE(nn < kDtypeNb); + double dot = 0.0, qn = 0.0, dn = 0.0; + for (int64_t d = 0; d < kDtypeDim; d++) { + double qv = query[q * kDtypeDim + d]; + double dv = train[nn * kDtypeDim + d]; + dot += qv * dv; + qn += qv * qv; + dn += dv * dv; + } + double expected = is_cosine ? dot / (std::sqrt(qn) * std::sqrt(dn) + 1e-12) : dot; + double got = dist[q * kDtypeTopk]; + REQUIRE(std::fabs(got - expected) <= 1e-2 * (std::fabs(expected) + 1.0)); + if (got > 0.0) positive_top1++; + } + // For this data the nearest neighbor's similarity is positive for the vast + // majority of queries; a sign inversion would make all of them negative. + REQUIRE(positive_top1 > kDtypeNq * 0.8); +} + +// v84: every requested result slot must be a real neighbor (no UINT64_MAX/-1 id +// or FLT_MAX sentinel). ef is guaranteed >= k: Knowhere's config bumps an unset +// ef to max(k, 16) and the GPU kernel additionally auto-clamps ef=max(ef,k) as +// defense-in-depth. Sentinel leakage into the last (k-ef) slots was the k>ef +// symptom. ef is intentionally left unset here so the requested topk drives it. +TEST_CASE("GPU HNSW returns k valid results (ef>=k)", "[gpu_hnsw_k_valid]") { + const auto version = GenTestVersionList(); + const int64_t topk = 100; + knowhere::Json json; + json[knowhere::meta::DIM] = kDtypeDim; + json[knowhere::meta::METRIC_TYPE] = knowhere::metric::L2; + json[knowhere::meta::TOPK] = topk; + json[knowhere::indexparam::HNSW_M] = 16; + json[knowhere::indexparam::EFCONSTRUCTION] = 200; + // EF intentionally unset: Knowhere sets ef = max(topk, 16) = topk here. + + auto train_ds = GenDataSet(kDtypeNb, kDtypeDim, kDtypeSeed); + auto query_ds = GenDataSet(kDtypeNq, kDtypeDim, kDtypeSeed + 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; + REQUIRE(cpu_idx.Serialize(bs) == knowhere::Status::success); + + 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()); + const int64_t* ids = results.value()->GetIds(); + const float* dist = results.value()->GetDistance(); + for (int64_t i = 0; i < kDtypeNq * topk; i++) { + REQUIRE(ids[i] >= 0); + REQUIRE(ids[i] < kDtypeNb); + REQUIRE(dist[i] < std::numeric_limits::max()); + } +} #endif diff --git a/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.h b/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.h index b52018a32..f4b87de35 100644 --- a/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.h +++ b/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.h @@ -26,8 +26,8 @@ #include #include -#include #include +#include namespace faiss { @@ -121,10 +121,11 @@ struct GpuIndexHNSW : public GpuIndex { /// over the searchImpl_/GpuIndex::search() path, which round-trips the /// labels D2H then H2D. /// - /// Distance convention: for inner-product and cosine metrics the returned - /// distances are the *negated* dot product (smaller == more similar), so - /// callers that want raw similarity scores must negate them (Knowhere does - /// this on the way out). L2 distances are returned as-is. + /// Distance convention: inner-product and cosine metrics return the true + /// similarity score (larger == more similar), matching faiss's + /// METRIC_INNER_PRODUCT contract. Internally the kernel keeps a negated + /// (smaller == more similar) score for a uniform min-first ordering and + /// negates it back on copy-out. L2 distances are returned as-is. void searchHost( idx_t n, const float* x_host, @@ -138,8 +139,9 @@ struct GpuIndexHNSW : public GpuIndex { /// verbatim — no bias/shift is applied, matching the dataset upload which /// already reverses FAISS's +128 SQ bias. When dim % 4 != 0 (DP4A requires /// groups of four int8 lanes) this transparently falls back to the fp32 - /// searchHost() path. Same negated IP/cosine distance convention as - /// searchHost(). All input/output pointers must be host memory. + /// searchHost() path. Same IP/cosine distance convention as searchHost() + /// (true similarity returned). All input/output pointers must be host + /// memory. void searchHostInt8( idx_t n, const int8_t* x_host, diff --git a/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearch.cuh b/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearch.cuh index f97d93ae4..3bcf2bb27 100644 --- a/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearch.cuh +++ b/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearch.cuh @@ -28,6 +28,7 @@ #include #include +#include #include #include #include @@ -60,6 +61,24 @@ inline void gpu_hnsw_search( int ef = params.ef; int sw = params.search_width; + // Reject degenerate params up front: ef and search_width feed the + // shared-memory sizing and the auto-iteration bound (2*ef/sw below), so a + // value of 0 would divide by zero and a negative value would under-size + // the buffers. + if (ef <= 0 || sw <= 0) { + throw std::runtime_error( + std::string("gpu_hnsw: ef and search_width must be > 0 (ef=") + + std::to_string(ef) + ", search_width=" + std::to_string(sw) + + ")"); + } + // Auto-clamp ef up to k. The kernel tracks exactly ef candidates in the + // result beam, so ef < k could only ever return ef real neighbors with the + // remaining k-ef slots padded with sentinels. Raising ef to k matches the + // CPU / Knowhere HNSW contract (always return k valid results). The + // shared-memory-fit check below still throws if even k slots cannot fit. + if (k > ef) { + ef = k; + } int max_iter = params.max_iterations > 0 ? params.max_iterations : 2 * ef / sw + 10; @@ -160,6 +179,18 @@ inline void gpu_hnsw_search( std::to_string(smem_max) + " bytes); reduce search_width"); } + // ef was auto-clamped up to at least k above. If even k candidate + // slots cannot fit in shared memory, fail loudly rather than + // silently returning fewer than k valid results. + if (k > max_ef) { + throw std::runtime_error( + std::string("gpu_hnsw: k=") + std::to_string(k) + + " needs k candidate slots (ef>=k), only " + + std::to_string(max_ef) + + " fit in device shared memory (" + + std::to_string(smem_max) + + " bytes); reduce k or raise thread_block_size"); + } if (ef > max_ef) { fprintf(stderr, "[gpu_hnsw] warning: ef=%d exceeds the per-block " @@ -175,13 +206,22 @@ inline void gpu_hnsw_search( 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. + // without this the kernel launch would fail for large ef. This is + // cached per kernel instantiation: cudaFuncSetAttribute configures the + // kernel's max dynamic shared memory globally, so it only needs to run + // once per (kernel, high-water size) rather than on every search. The + // static lives in this generic-lambda body, which is instantiated once + // per combination. if (smem_size > 49152) { - GPU_HNSW_CUDA_CHECK(cudaFuncSetAttribute( - hnsw_kernel::layer0_beam_search_kernel< - DataT, QueryT, USE_DP4A>, - cudaFuncAttributeMaxDynamicSharedMemorySize, - static_cast(smem_size))); + static std::atomic configured_smem{0}; + if (smem_size > configured_smem.load(std::memory_order_relaxed)) { + GPU_HNSW_CUDA_CHECK(cudaFuncSetAttribute( + hnsw_kernel::layer0_beam_search_kernel< + DataT, QueryT, USE_DP4A>, + cudaFuncAttributeMaxDynamicSharedMemorySize, + static_cast(smem_size))); + configured_smem.store(smem_size, std::memory_order_relaxed); + } } int N_int = static_cast(idx.n_rows); diff --git a/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearchKernel.cuh b/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearchKernel.cuh index dd0f5ede5..c972d946c 100644 --- a/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearchKernel.cuh +++ b/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearchKernel.cuh @@ -402,7 +402,8 @@ __device__ __forceinline__ void parallel_merge_into_result( float rv = ri_valid ? result_dists[ri_safe] : FLT_MAX; uint32_t si_id = (si < sc) ? staging_ids[si] : UINT32_MAX; uint32_t ri_id = ri_valid ? result_ids[ri_safe] : UINT32_MAX; - bool take_staging = (si < sc) && (!ri_valid || (sv < rv) || (sv == rv && si_id < ri_id)); + bool take_staging = (si < sc) && + (!ri_valid || (sv < rv) || (sv == rv && si_id < ri_id)); if (take_staging) { merged_ids[pos] = si_id; @@ -638,16 +639,25 @@ __global__ void layer0_beam_search_kernel( } // --- Copy top-k results to global memory --- + // The beam search works internally with a "smaller == closer" convention: + // thread_ip_distance() returns the negated dot product so inner-product and + // cosine can share the same min-first ordering as L2. Negate back on + // copy-out for those metrics so callers receive the true similarity score + // (faiss's METRIC_INNER_PRODUCT contract is "larger == closer"). L2 is + // already a genuine distance and is copied verbatim. Empty slots use the + // worst score for the metric (-FLT_MAX for IP/cosine, +FLT_MAX for L2) so + // they always sort last regardless of the caller's ordering. int rc = meta[0]; for (int i = threadIdx.x; i < k; i += blockDim.x) { if (i < rc && result_ids[i] < static_cast(N)) { d_neighbors[static_cast(query_idx) * k + i] = static_cast(result_ids[i]); d_distances[static_cast(query_idx) * k + i] = - result_dists[i]; + use_inner_product ? -result_dists[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; + d_distances[static_cast(query_idx) * k + i] = + use_inner_product ? -FLT_MAX : FLT_MAX; } } } diff --git a/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.cu b/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.cu index 67d00365a..8766cbfb3 100644 --- a/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.cu +++ b/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.cu @@ -118,6 +118,10 @@ GpuHnswSearchScratch::~GpuHnswSearchScratch() { } GpuHnswScratchSlot::~GpuHnswScratchSlot() { + // Set the owning device before destroying the stream so it runs in the + // correct context on multi-GPU systems (scratch.device is assigned when the + // slot is created in GpuHnswScratchPool::init_once()). + cudaSetDevice(scratch.device); if (stream) cudaStreamDestroy(stream); } From d022e89825c499d43822c90041ebb9d4908b6a25 Mon Sep 17 00:00:00 2001 From: premal Date: Fri, 17 Jul 2026 20:44:31 +0000 Subject: [PATCH 30/37] test(gpu_hnsw): assert GPU-vs-CPU recall parity in ef sweep, not absolute floor The [gpu_hnsw_compare] 'Varying ef comparison' section asserted gpu_recall >= 0.5f for every ef, but at ef=16 (~= k) HNSW recall is legitimately low on both engines (CPU HNSW itself scores ~0.41, GPU ~0.47). The absolute floor was meaningless in that regime and the section already computes cpu_recall but never used it. Assert GPU tracks the CPU oracle within a small tolerance instead, which is the correct invariant for a GPU-vs-CPU comparison test and catches real regressions at every ef. Signed-off-by: premal Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/ut/test_gpu_search.cc | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/tests/ut/test_gpu_search.cc b/tests/ut/test_gpu_search.cc index 2aae65d4e..dd71ff120 100644 --- a/tests/ut/test_gpu_search.cc +++ b/tests/ut/test_gpu_search.cc @@ -1039,8 +1039,11 @@ TEST_CASE("Test CPU vs GPU HNSW Comparison", "[gpu_hnsw_compare]") { 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); + // This is a GPU-vs-CPU comparison: GPU HNSW must track the CPU HNSW oracle at + // every ef, not hit an arbitrary absolute floor. At small ef (e.g. ef=16, ~= k) + // HNSW recall is legitimately low on both engines (CPU itself is < 0.5 here), so + // an absolute floor is meaningless; parity with CPU is the correct invariant. + REQUIRE(gpu_recall >= cpu_recall - 0.05f); } fprintf(stderr, "===\n\n"); } From 898a93f2bd36641488145659782110fa0108c250 Mon Sep 17 00:00:00 2001 From: premal Date: Sat, 18 Jul 2026 07:03:38 +0000 Subject: [PATCH 31/37] fix(gpu_hnsw): derive query-side metric from GPU index, not search config I1: GpuHnswIndexNode::Search() detected COSINE from hnsw_cfg.metric_type to decide fp32 query normalization and int8 1/||q|| distance scaling. A search config that omits metric_type defaults to L2, so a cosine index would skip normalization and silently return wrong scores. Derive the metric from the built GPU index via the new GpuIndexHNSW::isCosine() accessor (captured at upload time from the FAISS index type), matching the index-side detection already used on the upload path. Re-vendor the faiss GPU tree (byte-identical to 6si/faiss@24be1b66): - mutex-guarded, monotonic cudaFuncSetAttribute smem configuration (TOCTOU fix) - entry_point bounds validation in extract_hnsw_layers - GpuIndexHNSW isCosine()/useInnerProduct() accessors Add a regression test that searches a COSINE index with metric_type omitted from the config and asserts identical results to the explicit-COSINE search. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Signed-off-by: premal Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- src/index/hnsw/faiss_hnsw.cc | 10 +++- tests/ut/test_gpu_search.cc | 52 +++++++++++++++++++ thirdparty/faiss/faiss/gpu/GpuIndexHNSW.cu | 6 +++ thirdparty/faiss/faiss/gpu/GpuIndexHNSW.h | 16 ++++++ .../faiss/faiss/gpu/impl/GpuHnswBuild.cuh | 10 ++++ .../faiss/faiss/gpu/impl/GpuHnswSearch.cuh | 21 ++++++-- 6 files changed, 108 insertions(+), 7 deletions(-) diff --git a/src/index/hnsw/faiss_hnsw.cc b/src/index/hnsw/faiss_hnsw.cc index 3f008c645..0457f286f 100644 --- a/src/index/hnsw/faiss_hnsw.cc +++ b/src/index/hnsw/faiss_hnsw.cc @@ -3574,6 +3574,12 @@ class GpuHnswIndexNode : public BaseFaissRegularIndexHNSWNode { auto dim = dataset->GetDim(); auto ef = hnsw_cfg.ef.value_or(200); + // Derive the metric from the built GPU index (captured at upload time + // from the FAISS index type), NOT from the search config. A config that + // omits metric_type defaults to L2, which would skip cosine query + // normalization and silently return wrong scores for a cosine index. + const bool is_cosine = gpu_snapshot->isCosine(); + auto h_ids = std::make_unique(nq * k); auto h_dist = std::make_unique(nq * k); @@ -3596,7 +3602,7 @@ class GpuHnswIndexNode : public BaseFaissRegularIndexHNSWNode { // (the true cosine numerator scaled by the db inverse norm, already // sign-corrected to "larger == closer"); dividing by the query norm // finalizes the cosine similarity. - if (IsMetricType(hnsw_cfg.metric_type.value(), metric::COSINE)) { + if (is_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; @@ -3623,7 +3629,7 @@ class GpuHnswIndexNode : public BaseFaissRegularIndexHNSWNode { // 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)) { + if (is_cosine) { normalized_queries = std::make_unique(nq * dim); for (int64_t i = 0; i < nq; i++) { const float* src = h_queries_raw + i * dim; diff --git a/tests/ut/test_gpu_search.cc b/tests/ut/test_gpu_search.cc index dd71ff120..f472a2ccd 100644 --- a/tests/ut/test_gpu_search.cc +++ b/tests/ut/test_gpu_search.cc @@ -538,6 +538,58 @@ TEST_CASE("Test All GPU Index", "[search]") { REQUIRE(recall >= 0.80f); } + SECTION("Test GPU HNSW Cosine Metric Omitted From Search Config") { + // Regression for I1: the GPU node must derive the metric from the built + // GPU index (via GpuIndexHNSW::isCosine()), NOT from the search config. + // A search config that omits metric_type defaults to L2, which would + // skip cosine query normalization and silently return wrong scores. + // Here the index is COSINE but the search json carries no metric_type; + // results must match a search whose json explicitly sets COSINE. + knowhere::Json build_json; + build_json[knowhere::meta::DIM] = dim; + build_json[knowhere::meta::METRIC_TYPE] = knowhere::metric::COSINE; + build_json[knowhere::meta::TOPK] = 10; + build_json[knowhere::indexparam::HNSW_M] = 16; + build_json[knowhere::indexparam::EFCONSTRUCTION] = 200; + build_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(); + 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); + + // Baseline: search with metric_type explicitly set to COSINE. + auto with_metric = gpu_idx.Search(query_ds, build_json, nullptr); + REQUIRE(with_metric.has_value()); + + // Same search, but metric_type omitted from the config. + knowhere::Json no_metric_json = build_json; + no_metric_json.erase(knowhere::meta::METRIC_TYPE); + auto without_metric = gpu_idx.Search(query_ds, no_metric_json, nullptr); + REQUIRE(without_metric.has_value()); + + // The metric is index-derived, so both searches must be identical. + const int64_t topk = build_json[knowhere::meta::TOPK].get(); + const auto* ids_a = with_metric.value()->GetIds(); + const auto* ids_b = without_metric.value()->GetIds(); + const auto* dist_a = with_metric.value()->GetDistance(); + const auto* dist_b = without_metric.value()->GetDistance(); + for (int64_t i = 0; i < nq * topk; ++i) { + REQUIRE(ids_a[i] == ids_b[i]); + REQUIRE(dist_a[i] == Approx(dist_b[i]).epsilon(1e-5)); + } + } + SECTION("Test GPU HNSW Search TopK") { knowhere::Json hnsw_json; hnsw_json[knowhere::meta::DIM] = dim; diff --git a/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.cu b/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.cu index 63eb18e91..287066e21 100644 --- a/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.cu +++ b/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.cu @@ -68,6 +68,9 @@ void GpuIndexHNSW::copyFrom( bool use_ip = is_cosine || (index->metric_type == faiss::METRIC_INNER_PRODUCT); + is_cosine_ = is_cosine; + use_ip_ = use_ip; + if (dynamic_cast(index->storage)) { deviceIndex_ = from_faiss_hnsw_sq(*index, use_ip, is_cosine, config_.device); @@ -92,6 +95,9 @@ void GpuIndexHNSW::copyFromWithMetric( this->metric_type = index->metric_type; this->ntotal = index->ntotal; + is_cosine_ = is_cosine; + use_ip_ = use_ip; + if (dynamic_cast(index->storage)) { deviceIndex_ = from_faiss_hnsw_sq(*index, use_ip, is_cosine, config_.device); diff --git a/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.h b/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.h index f4b87de35..f374e3fe0 100644 --- a/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.h +++ b/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.h @@ -150,6 +150,18 @@ struct GpuIndexHNSW : public GpuIndex { idx_t* labels_host, const GpuHnswSearchParams& params) const; + /// Metric interpretation captured at copyFrom()/copyFromWithMetric() time + /// from the FAISS index type, so callers can post-process results (cosine + /// query normalization, IP handling) without re-deriving the metric from a + /// per-search config — the config may omit metric_type and default to L2, + /// which would silently mishandle a cosine/IP index. + bool isCosine() const { + return is_cosine_; + } + bool useInnerProduct() const { + return use_ip_; + } + protected: bool addImplRequiresIDs_() const override; @@ -171,6 +183,10 @@ struct GpuIndexHNSW : public GpuIndex { mutable std::mutex searchParamsMutex_; mutable GpuHnswSearchParams directSearchParams_; mutable bool hasDirectSearchParams_ = false; + + // Metric interpretation, set at copy time (see isCosine()/useInnerProduct()). + bool is_cosine_ = false; + bool use_ip_ = false; }; } // namespace gpu diff --git a/thirdparty/faiss/faiss/gpu/impl/GpuHnswBuild.cuh b/thirdparty/faiss/faiss/gpu/impl/GpuHnswBuild.cuh index f051bcd73..7d2305d72 100644 --- a/thirdparty/faiss/faiss/gpu/impl/GpuHnswBuild.cuh +++ b/thirdparty/faiss/faiss/gpu/impl/GpuHnswBuild.cuh @@ -95,6 +95,16 @@ inline void extract_hnsw_layers( const int maxM = hnsw.nb_neighbors(1); const int max_lv = hnsw.max_level; + // The kernels dereference d_dataset + entry_point * dim directly, so a + // malformed CPU index carrying an invalid/sentinel entry_point (e.g. -1 + // cast to UINT32_MAX) would be an out-of-bounds device read. Validate here + // rather than faulting on the GPU. + if (hnsw.entry_point < 0 || hnsw.entry_point >= n_rows) { + throw std::runtime_error( + std::string("gpu_hnsw: invalid HNSW entry_point ") + + std::to_string(static_cast(hnsw.entry_point)) + + " (n_rows=" + std::to_string(n_rows) + ")"); + } entry_point = static_cast(hnsw.entry_point); M = maxM; max_degree0 = maxM0; diff --git a/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearch.cuh b/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearch.cuh index 3bcf2bb27..78b75af94 100644 --- a/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearch.cuh +++ b/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearch.cuh @@ -28,8 +28,8 @@ #include #include -#include #include +#include #include #include #include @@ -210,17 +210,28 @@ inline void gpu_hnsw_search( // cached per kernel instantiation: cudaFuncSetAttribute configures the // kernel's max dynamic shared memory globally, so it only needs to run // once per (kernel, high-water size) rather than on every search. The - // static lives in this generic-lambda body, which is instantiated once + // statics live in this generic-lambda body, which is instantiated once // per combination. + // + // The mutex makes the check-set-store atomic and monotonic: without it, + // two concurrent searches with different smem sizes can interleave so a + // smaller-ef search's cudaFuncSetAttribute runs *after* a larger one, + // downgrading the kernel's global attribute below what a recorded + // high-water mark implies, and a later intermediate search then skips + // the set (thinks it's configured) and fails to launch. Under the lock + // the attribute only ever grows, and is always >= the current launch's + // requirement before we proceed. if (smem_size > 49152) { - static std::atomic configured_smem{0}; - if (smem_size > configured_smem.load(std::memory_order_relaxed)) { + static std::mutex configured_smem_mutex; + static size_t configured_smem = 0; + std::lock_guard lock(configured_smem_mutex); + if (smem_size > configured_smem) { GPU_HNSW_CUDA_CHECK(cudaFuncSetAttribute( hnsw_kernel::layer0_beam_search_kernel< DataT, QueryT, USE_DP4A>, cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast(smem_size))); - configured_smem.store(smem_size, std::memory_order_relaxed); + configured_smem = smem_size; } } From da32be6ec2b3f364e88f91e656afdce30c9c2204 Mon Sep 17 00:00:00 2001 From: premal Date: Sat, 18 Jul 2026 14:53:51 +0000 Subject: [PATCH 32/37] fix(gpu_hnsw): re-vendor faiss GpuIndexHNSW.cu (reset clears metric flags) Byte-identical re-vendor of 6si/faiss@a43d3afe: reset() now clears is_cosine_/use_ip_ so a reused GpuIndexHNSW can't inherit a stale cosine/IP interpretation. Verified with cmp against the faiss source. Signed-off-by: premal Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- thirdparty/faiss/faiss/gpu/GpuIndexHNSW.cu | 2 ++ 1 file changed, 2 insertions(+) diff --git a/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.cu b/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.cu index 287066e21..440df00c0 100644 --- a/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.cu +++ b/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.cu @@ -113,6 +113,8 @@ void GpuIndexHNSW::reset() { deviceIndex_.reset(); this->ntotal = 0; this->is_trained = false; + is_cosine_ = false; + use_ip_ = false; } void GpuIndexHNSW::setSearchParams(const GpuHnswSearchParams& params) const { From e8138631c325c1ea8f394bee65ce219761ed0f95 Mon Sep 17 00:00:00 2001 From: premal Date: Sat, 18 Jul 2026 18:33:47 +0000 Subject: [PATCH 33/37] doc(gpu_hnsw): document faiss_gpu_hnsw as sole GPU compile path (no dup) The faiss_gpu_hnsw OBJECT library is the only place the faiss GPU sources are compiled in the knowhere build: FAISS_SRCS does not glob faiss/gpu/*, and the vendored thirdparty/faiss/faiss/gpu/CMakeLists.txt (faiss_gpu_objs) is never add_subdirectory()'d. Comment-only; no compiled objects change. 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> --- cmake/libs/libfaiss.cmake | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/cmake/libs/libfaiss.cmake b/cmake/libs/libfaiss.cmake index 828a47817..b20dbe261 100644 --- a/cmake/libs/libfaiss.cmake +++ b/cmake/libs/libfaiss.cmake @@ -594,7 +594,22 @@ if(__PPC64) target_compile_definitions(faiss PRIVATE FINTEGER=int) endif() -# GPU HNSW CUDA sources — compiled when WITH_CUVS is enabled +# GPU HNSW CUDA sources — compiled when WITH_CUVS is enabled. +# +# This `faiss_gpu_hnsw` OBJECT library is the ONLY place the faiss GPU sources +# are compiled in the knowhere build: +# * The main `faiss` STATIC target above is built from FAISS_SRCS, which globs +# only thirdparty/faiss/faiss/{*.cpp,impl,utils,...} — it deliberately does +# NOT include thirdparty/faiss/faiss/gpu/*, so none of the files below are +# also compiled into `faiss`. +# * Upstream faiss ships its own GPU build file +# (thirdparty/faiss/faiss/gpu/CMakeLists.txt → the `faiss_gpu_objs` target +# over FAISS_GPU_SRC). Knowhere assembles faiss manually via this .cmake and +# never `add_subdirectory()`s the vendored faiss tree, so that file (and +# `faiss_gpu_objs`) is never part of the knowhere build graph. +# Therefore the two GPU build paths are mutually exclusive here — there is no +# duplicate compilation or ODR hazard. Keep it that way: add GPU sources to the +# list below, NOT by pulling in the vendored gpu/CMakeLists.txt. if(WITH_CUVS) set(FAISS_GPU_HNSW_SRCS thirdparty/faiss/faiss/gpu/GpuIndexHNSW.cu From fb261257865b365b63ea52ecc08b75c9a7edac11 Mon Sep 17 00:00:00 2001 From: premal Date: Sat, 18 Jul 2026 21:00:44 +0000 Subject: [PATCH 34/37] feat(gpu_hnsw): GPU-side filtered search (deletes/TTL/partition bitset) Re-vendor the faiss GPU tree with CPU-parity filtered search and wire the delete/TTL/partition BitsetView through GpuHnswIndexNode::Search() instead of rejecting a non-empty bitset. - faiss thirdparty: byte-identical copy of the filtered-search kernel, GPU brute-force header (GpuHnswBruteForce.cuh, added to gpu/CMakeLists.txt), host launcher and scratch types. - Search(): remove the blanket filtered-search rejection; pass bitset ptr/nbits/filtered-count and disable_fallback_brute_force into faiss for both native int8 and fp32 paths. Still reject id-remapped / id-offset bitsets (has_out_ids()/id_offset()!=0) since the kernel tests by raw node id (== row offset == bit index), and guard a bitset shorter than ntotal. Snapshot/reload concurrency, CPU-index release after upload, and cosine/int8 post-scaling preserved. - bitsetview.h: add id_offset() accessor for the shape guard. - tests: gpu_hnsw_id_mapping (blocking G-ID proof: node id == row offset == bitset index, via unfiltered self-query and filtered self-query) and gpu_hnsw_filtered_recall (no returned id filtered; recall vs brute force across delete ratios 0.1-0.99). Signed-off-by: premal Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- include/knowhere/bitsetview.h | 5 + src/index/hnsw/faiss_hnsw.cc | 49 ++- tests/ut/test_gpu_search.cc | 140 +++++++ thirdparty/faiss/faiss/gpu/CMakeLists.txt | 1 + thirdparty/faiss/faiss/gpu/GpuIndexHNSW.cu | 32 ++ .../faiss/gpu/impl/GpuHnswBruteForce.cuh | 162 ++++++++ .../faiss/faiss/gpu/impl/GpuHnswSearch.cuh | 211 ++++++++-- .../faiss/gpu/impl/GpuHnswSearchKernel.cuh | 380 ++++++++++++++++-- .../faiss/faiss/gpu/impl/GpuHnswTypes.cu | 30 ++ .../faiss/faiss/gpu/impl/GpuHnswTypes.h | 43 ++ 10 files changed, 983 insertions(+), 70 deletions(-) create mode 100644 thirdparty/faiss/faiss/gpu/impl/GpuHnswBruteForce.cuh diff --git a/include/knowhere/bitsetview.h b/include/knowhere/bitsetview.h index ee6db3a47..84eddcb08 100644 --- a/include/knowhere/bitsetview.h +++ b/include/knowhere/bitsetview.h @@ -69,6 +69,11 @@ class BitsetView { return out_ids_ != nullptr; } + size_t + id_offset() const { + return id_offset_; + } + void set_out_ids(const uint32_t* out_ids, size_t num_internal_ids, std::optional num_filtered_out_ids = std::nullopt) { diff --git a/src/index/hnsw/faiss_hnsw.cc b/src/index/hnsw/faiss_hnsw.cc index 0457f286f..7617c5817 100644 --- a/src/index/hnsw/faiss_hnsw.cc +++ b/src/index/hnsw/faiss_hnsw.cc @@ -3522,12 +3522,22 @@ class GpuHnswIndexNode : public BaseFaissRegularIndexHNSWNode { expected Search(const DataSetPtr dataset, std::unique_ptr cfg, const BitsetView& bitset, milvus::OpContext* op_context) const override { - // 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"); + // Filtered search (deletes / TTL / partition / visibility bitset) is + // applied inside the GPU kernel: filtered rows stay graph waypoints but + // never enter the top-k, and a brute-force fallback guarantees k live + // results under heavy deletes (see faiss GpuHnswSearch.cuh / + // GpuHnswSearchKernel.cuh). The kernel tests the bitset by raw FAISS + // node id, which equals the segment row offset == BitsetView index + // (proven by the gpu_hnsw_id_mapping test). That direct mapping only + // holds when the view has no id remapping and no id offset, so those + // rarer shapes (never produced for a sealed GPU segment's delete + // bitset) are still rejected rather than silently mis-filtered. + const bool has_filter = !bitset.empty() && bitset.data() != nullptr; + if (has_filter && (bitset.has_out_ids() || bitset.id_offset() != 0)) { + return expected::Err( + Status::invalid_args, + "GPU_HNSW filtered search does not support id-remapped or " + "id-offset bitsets"); } // Take a shared-ownership snapshot of the GPU index under gpu_mutex_. The @@ -3580,6 +3590,31 @@ class GpuHnswIndexNode : public BaseFaissRegularIndexHNSWNode { // normalization and silently return wrong scores for a cosine index. const bool is_cosine = gpu_snapshot->isCosine(); + // The kernel indexes the bitset by node id in [0, ntotal); a bitset + // shorter than ntotal would read out of bounds. Milvus sizes the delete + // bitset to the segment row count (== ntotal), so this only guards + // against a malformed caller. + if (has_filter && + static_cast(bitset.size()) < gpu_snapshot->ntotal) { + return expected::Err( + Status::invalid_args, + "GPU_HNSW filtered search: bitset shorter than index size"); + } + + // Fill the filtered-search fields of a search-params struct from the + // incoming bitset. A null/empty bitset leaves them defaulted, taking + // the unfiltered fast path. + auto set_filter_params = [&](faiss::gpu::GpuHnswSearchParams& gsp) { + if (!has_filter) { + return; + } + gsp.bitset_data = bitset.data(); + gsp.bitset_nbits = static_cast(bitset.size()); + gsp.bitset_filtered_count = static_cast(bitset.count()); + gsp.disable_fallback_brute_force = + hnsw_cfg.disable_fallback_brute_force.value_or(false); + }; + auto h_ids = std::make_unique(nq * k); auto h_dist = std::make_unique(nq * k); @@ -3590,6 +3625,7 @@ class GpuHnswIndexNode : public BaseFaissRegularIndexHNSWNode { try { faiss::gpu::GpuHnswSearchParams gsp; gsp.ef = ef; + set_filter_params(gsp); 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(); @@ -3645,6 +3681,7 @@ class GpuHnswIndexNode : public BaseFaissRegularIndexHNSWNode { try { faiss::gpu::GpuHnswSearchParams gsp; gsp.ef = ef; + set_filter_params(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(); diff --git a/tests/ut/test_gpu_search.cc b/tests/ut/test_gpu_search.cc index f472a2ccd..76ca55169 100644 --- a/tests/ut/test_gpu_search.cc +++ b/tests/ut/test_gpu_search.cc @@ -1799,4 +1799,144 @@ TEST_CASE("GPU HNSW returns k valid results (ef>=k)", "[gpu_hnsw_k_valid]") { REQUIRE(dist[i] < std::numeric_limits::max()); } } + +// Blocking prerequisite for GPU filtered search: the kernel tests the delete +// bitset by raw FAISS node id, so correctness hinges on +// GPU node id == storage row offset == BitsetView bit index. +// This proves that mapping empirically with known ordered data and known +// filters, independent of recall. If this ever fails, filtered search silently +// excludes the wrong rows and every downstream filtered test is meaningless. +TEST_CASE("GPU HNSW G-ID mapping: node id == row offset == bitset index", "[gpu_hnsw_id_mapping]") { + const auto version = GenTestVersionList(); + const int64_t nb = 2000; + const int64_t dim = 32; + + knowhere::Json json; + json[knowhere::meta::DIM] = dim; + json[knowhere::meta::METRIC_TYPE] = knowhere::metric::L2; + json[knowhere::meta::TOPK] = 1; + json[knowhere::indexparam::HNSW_M] = 16; + json[knowhere::indexparam::EFCONSTRUCTION] = 200; + json[knowhere::indexparam::EF] = 200; + + // Ordered training data; each row is its own exact nearest neighbor (L2=0), + // so querying with row i must return id i. + auto train_ds = GenDataSet(nb, dim, 42); + + 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; + REQUIRE(cpu_idx.Serialize(bs) == knowhere::Status::success); + + auto gpu_idx = + knowhere::IndexFactory::Instance().Create(knowhere::IndexEnum::INDEX_GPU_HNSW, version).value(); + REQUIRE(gpu_idx.Deserialize(bs) == knowhere::Status::success); + REQUIRE(gpu_idx.Count() == nb); + + // 1) Unfiltered self-query: node id == storage row offset. + { + auto res = gpu_idx.Search(train_ds, json, nullptr); + REQUIRE(res.has_value()); + const int64_t* ids = res.value()->GetIds(); + for (int64_t i = 0; i < nb; i++) { + REQUIRE(ids[i] == i); + } + } + + // 2) Filtered self-query: bit index == storage row offset. Filter every + // third row; a returned id must never be a filtered bit, and a filtered + // row's own id must be excluded from its self-query. + { + std::vector bitset_data((nb + 7) / 8, 0); + int64_t filtered_count = 0; + for (int64_t r = 0; r < nb; r += 3) { + bitset_data[r / 8] |= static_cast(1 << (r % 8)); + filtered_count++; + } + knowhere::BitsetView bitset(bitset_data.data(), nb, filtered_count); + + auto res = gpu_idx.Search(train_ds, json, bitset); + REQUIRE(res.has_value()); + const int64_t* ids = res.value()->GetIds(); + for (int64_t i = 0; i < nb; i++) { + const int64_t id = ids[i]; + if (id == -1) { + continue; + } + const bool id_filtered = (bitset_data[id / 8] & (1 << (id % 8))) != 0; + REQUIRE_FALSE(id_filtered); + const bool query_filtered = (bitset_data[i / 8] & (1 << (i % 8))) != 0; + if (query_filtered) { + REQUIRE(id != i); + } + } + } +} + +// Filtered-search recall/exclusion parity vs brute force across delete ratios. +// Exercises the two-tier beam + alpha gate + brute-force fallback: no returned +// id may be filtered, and recall of the live top-k must track brute force. +TEST_CASE("GPU HNSW filtered search recall vs brute force", "[gpu_hnsw_filtered_recall]") { + const auto version = GenTestVersionList(); + const int64_t nb = 8000; + const int64_t dim = 64; + const int64_t nq = 200; + const int64_t topk = 10; + + knowhere::Json json; + json[knowhere::meta::DIM] = dim; + json[knowhere::meta::METRIC_TYPE] = knowhere::metric::L2; + json[knowhere::meta::TOPK] = topk; + json[knowhere::indexparam::HNSW_M] = 16; + json[knowhere::indexparam::EFCONSTRUCTION] = 200; + json[knowhere::indexparam::EF] = 200; + + auto train_ds = GenDataSet(nb, dim, 42); + auto query_ds = GenDataSet(nq, dim, 44); + + 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; + REQUIRE(cpu_idx.Serialize(bs) == knowhere::Status::success); + + auto gpu_idx = + knowhere::IndexFactory::Instance().Create(knowhere::IndexEnum::INDEX_GPU_HNSW, version).value(); + REQUIRE(gpu_idx.Deserialize(bs) == knowhere::Status::success); + + std::mt19937 rng(123); + for (const float ratio : {0.1f, 0.5f, 0.9f, 0.95f, 0.99f}) { + std::vector bitset_data((nb + 7) / 8, 0); + int64_t filtered_count = 0; + std::uniform_real_distribution u(0.0f, 1.0f); + for (int64_t r = 0; r < nb; r++) { + if (u(rng) < ratio) { + bitset_data[r / 8] |= static_cast(1 << (r % 8)); + filtered_count++; + } + } + knowhere::BitsetView bitset(bitset_data.data(), nb, filtered_count); + + auto res = gpu_idx.Search(query_ds, json, bitset); + REQUIRE(res.has_value()); + const int64_t* ids = res.value()->GetIds(); + + // No returned id may be filtered out. + for (int64_t i = 0; i < nq * topk; i++) { + const int64_t id = ids[i]; + if (id == -1) { + continue; + } + REQUIRE_FALSE((bitset_data[id / 8] & (1 << (id % 8))) != 0); + } + + auto gt = knowhere::BruteForce::Search(train_ds, query_ds, json, bitset); + REQUIRE(gt.has_value()); + float recall = GetKNNRecall(*gt.value(), *res.value()); + // Parity is asserted by recall, not bit-identical traversal (the alpha + // gate is deterministic but not CPU encounter-order identical). + REQUIRE(recall >= 0.85f); + } +} #endif diff --git a/thirdparty/faiss/faiss/gpu/CMakeLists.txt b/thirdparty/faiss/faiss/gpu/CMakeLists.txt index 59c7e7eab..7ab9ffd39 100644 --- a/thirdparty/faiss/faiss/gpu/CMakeLists.txt +++ b/thirdparty/faiss/faiss/gpu/CMakeLists.txt @@ -103,6 +103,7 @@ set(FAISS_GPU_HEADERS GpuIndicesOptions.h GpuResources.h StandardGpuResources.h + impl/GpuHnswBruteForce.cuh impl/GpuHnswBuild.cuh impl/GpuHnswSearch.cuh impl/GpuHnswSearchKernel.cuh diff --git a/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.cu b/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.cu index 440df00c0..0f7b5a0c2 100644 --- a/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.cu +++ b/thirdparty/faiss/faiss/gpu/GpuIndexHNSW.cu @@ -37,6 +37,34 @@ namespace faiss { namespace gpu { +namespace { + +// Upload the filtered-search bitset (deletes / TTL / partition) into the +// scratch slot's device buffer on the search stream. gpu_hnsw_search reads +// sp.bitset_data only to decide the filtered path; the bytes themselves come +// from sc.d_bitset. The host BitsetView bytes referenced by sp.bitset_data +// must stay alive until the stream is synchronized (the caller synchronizes +// before returning, which happens in every search path here). +inline void upload_bitset_if_needed( + GpuHnswSearchScratch& sc, + const GpuHnswSearchParams& sp, + int nq, + cudaStream_t stream) { + if (sp.bitset_data == nullptr || sp.bitset_nbits <= 0) { + return; + } + size_t bytes = static_cast((sp.bitset_nbits + 7) / 8); + sc.ensure_filter(nq, bytes); + GPU_HNSW_CUDA_CHECK(cudaMemcpyAsync( + sc.d_bitset, + sp.bitset_data, + bytes, + cudaMemcpyHostToDevice, + stream)); +} + +} // namespace + GpuIndexHNSW::GpuIndexHNSW( GpuResourcesProvider* provider, int dims, @@ -270,6 +298,8 @@ void GpuIndexHNSW::searchHost( cudaMemcpyDefault, stream)); + upload_bitset_if_needed(sc, sp, nq, stream); + gpu_hnsw_search(stream, sp, idx, sc, nq, k); GPU_HNSW_CUDA_CHECK(cudaMemcpyAsync( @@ -360,6 +390,8 @@ void GpuIndexHNSW::searchHostInt8( // while the async copies are in flight. GPU_HNSW_CUDA_CHECK(cudaStreamSynchronize(stream)); + upload_bitset_if_needed(sc, sp, nq, stream); + gpu_hnsw_search(stream, sp, idx, sc, nq, k); GPU_HNSW_CUDA_CHECK(cudaMemcpyAsync( diff --git a/thirdparty/faiss/faiss/gpu/impl/GpuHnswBruteForce.cuh b/thirdparty/faiss/faiss/gpu/impl/GpuHnswBruteForce.cuh new file mode 100644 index 000000000..0f0045db3 --- /dev/null +++ b/thirdparty/faiss/faiss/gpu/impl/GpuHnswBruteForce.cuh @@ -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 + +namespace faiss { +namespace gpu { +namespace hnsw_kernel { + +// Fixed upper bound on the launch block size for the brute-force kernel; the +// tree reduction requires a power-of-two block, so callers must launch with a +// power-of-two blockDim <= kBruteForceMaxBlock. +constexpr int kBruteForceMaxBlock = 256; + +// Brute-force top-k over the live (non-filtered) rows for a set of queries, +// operating entirely from GPU-resident vectors (Knowhere frees the CPU HNSW +// index after upload, so the fallback cannot call it). Distances use the same +// "smaller == closer" internal convention as the graph kernel (IP negated), +// and are negated back on copy-out so callers receive the true similarity. +// +// One block per work item. The work item maps to a query index via d_worklist, +// or identity when d_worklist == nullptr (the up-front all-queries path). Rows +// filtered out by the bitset are skipped. Selection is iterative in +// (distance, id) lexicographic order: k block-wide reductions, each finding the +// smallest live candidate strictly after the previous winner. This needs no +// per-query top-k buffer and is duplicate-free because (dist, id) is a total +// order. Cost is O(k * N / blockDim) per query. When fewer than k live rows +// exist the remaining slots are filled with sentinels (-1 id, worst score). +// +// Requires: blockDim.x is a power of two, <= kBruteForceMaxBlock. For the DP4A +// int8 path (USE_DP4A) dim must be a multiple of 4, matching the graph kernel. +template +__global__ void brute_force_topk_kernel( + const QueryT* __restrict__ d_queries, + const DataT* __restrict__ d_dataset, + const float* __restrict__ d_inv_norms, + const uint8_t* __restrict__ d_bitset, + const uint32_t* __restrict__ d_worklist, // null => identity mapping + int num_items, + const int* __restrict__ d_num_items, // null => use num_items directly + int N, + int dim, + int k, + bool use_inner_product, + uint64_t* __restrict__ d_neighbors, + float* __restrict__ d_distances) { + // The per-query fallback launches a full num_queries grid and reads the + // device-side worklist length here, so no host sync is needed between the + // graph kernel and this pass. The up-front path passes num_items directly. + int count = (d_num_items != nullptr) ? *d_num_items : num_items; + int item = blockIdx.x; + if (item >= count) + return; + int query_idx = (d_worklist != nullptr) + ? static_cast(d_worklist[item]) + : item; + + const QueryT* query = d_queries + static_cast(query_idx) * dim; + int dim4 = dim / 4; + int tid = threadIdx.x; + + __shared__ float s_dist[kBruteForceMaxBlock]; + __shared__ uint32_t s_id[kBruteForceMaxBlock]; + __shared__ float win_dist_s; + __shared__ uint32_t win_id_s; + + // (-inf, 0): the first round selects the global minimum since any finite + // distance is strictly greater than -FLT_MAX. + float prev_dist = -FLT_MAX; + uint32_t prev_id = 0; + + for (int t = 0; t < k; t++) { + float local_best = FLT_MAX; + uint32_t local_id = UINT32_MAX; + for (int row = tid; row < N; row += blockDim.x) { + uint32_t r = static_cast(row); + if (is_bitset_filtered(d_bitset, r)) + continue; + float d = layer0_distance( + query, d_dataset, d_inv_norms, r, dim, dim4, + use_inner_product); + bool after = (d > prev_dist) || (d == prev_dist && r > prev_id); + if (!after) + continue; + if (d < local_best || (d == local_best && r < local_id)) { + local_best = d; + local_id = r; + } + } + s_dist[tid] = local_best; + s_id[tid] = local_id; + __syncthreads(); + for (int stride = blockDim.x >> 1; stride > 0; stride >>= 1) { + if (tid < stride) { + float od = s_dist[tid + stride]; + uint32_t oi = s_id[tid + stride]; + if (od < s_dist[tid] || + (od == s_dist[tid] && oi < s_id[tid])) { + s_dist[tid] = od; + s_id[tid] = oi; + } + } + __syncthreads(); + } + if (tid == 0) { + win_dist_s = s_dist[0]; + win_id_s = s_id[0]; + } + __syncthreads(); + float win_dist = win_dist_s; + uint32_t win_id = win_id_s; + + int64_t out = static_cast(query_idx) * k + t; + if (win_id == UINT32_MAX) { + // No live candidate remains: sentinel-fill this and the rest. + for (int tt = t + tid; tt < k; tt += blockDim.x) { + int64_t o2 = static_cast(query_idx) * k + tt; + d_neighbors[o2] = UINT64_MAX; + d_distances[o2] = use_inner_product ? -FLT_MAX : FLT_MAX; + } + break; + } + if (tid == 0) { + d_neighbors[out] = static_cast(win_id); + d_distances[out] = use_inner_product ? -win_dist : win_dist; + } + prev_dist = win_dist; + prev_id = win_id; + __syncthreads(); + } +} + +} // namespace hnsw_kernel +} // namespace gpu +} // namespace faiss diff --git a/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearch.cuh b/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearch.cuh index 78b75af94..afab6fecf 100644 --- a/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearch.cuh +++ b/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearch.cuh @@ -34,8 +34,9 @@ #include #include -#include #include +#include +#include #define GPU_HNSW_CUDA_CHECK(expr) \ do { \ @@ -85,6 +86,37 @@ inline void gpu_hnsw_search( int dim = static_cast(idx.dim); int num_upper_layers = idx.num_upper_layers_built; + // --- Filtered-search setup (deletes / TTL / partition bitset) --- + // The bitset bytes are uploaded into sc.d_bitset by the caller (searchHost/ + // searchHostInt8) on this same stream before gpu_hnsw_search runs. When + // bitset_data is null we take the unfiltered fast path. + bool has_filter = (params.bitset_data != nullptr); + int64_t nbits = params.bitset_nbits; + int64_t filtered_count = params.bitset_filtered_count; + int64_t live_count = has_filter ? (nbits - filtered_count) + : static_cast(idx.n_rows); + if (live_count < 0) { + live_count = 0; + } + // kAlpha rate-limits how often filtered nodes become waypoints; it scales + // with the delete ratio, matching CPU HNSW (kAlpha = filter_ratio * 0.7). + float filter_ratio = (has_filter && nbits > 0) + ? static_cast(filtered_count) / static_cast(nbits) + : 0.0f; + float kAlpha = filter_ratio * 0.7f; + bool disable_bf = params.disable_fallback_brute_force; + + // Up-front brute force: when almost everything is filtered, or k is a large + // fraction of the live set, the graph walk cannot beat a full scan and + // risks returning fewer than k live results. Mirrors CPU HNSW's + // kHnswSearchKnnBFFilterThreshold (0.93) and kHnswSearchBFTopkThreshold + // (0.5, applied to the live count). Disabled when the caller disables the + // brute-force fallback. + bool up_front_bf = has_filter && !disable_bf && nbits > 0 && + (static_cast(filtered_count) >= + 0.93 * static_cast(nbits) || + static_cast(k) >= 0.5 * static_cast(live_count)); + // Layer-0 is templated on the dataset type (DataT), the layer-0 query type // (QueryT: float generic, int8_t for the native DP4A path) and USE_DP4A. // The upper-layer greedy descent always uses the fp32 queries (sc.d_queries). @@ -92,6 +124,41 @@ inline void gpu_hnsw_search( const DataT* d_data, const float* d_inv_norms, const QueryT* d_layer0_queries) { + int N_int = static_cast(idx.n_rows); + + // Brute-force launcher (shares the current dtype specialization). The + // block must be a power of two <= kBruteForceMaxBlock for the tree + // reduction. Grid is always num_queries: the up-front path uses the + // identity mapping, the per-query fallback reads its worklist length + // from d_num on the device (no host sync between graph and BF). + int bf_block = 128; + auto launch_bf = [&](const uint32_t* worklist, + const int* d_num, + int num_items) { + hnsw_kernel::brute_force_topk_kernel + <<>>( + d_layer0_queries, + d_data, + d_inv_norms, + sc.d_bitset, + worklist, + num_items, + d_num, + N_int, + dim, + k, + idx.use_ip, + sc.d_neighbors, + sc.d_distances); + GPU_HNSW_CUDA_CHECK(cudaGetLastError()); + }; + + // Up-front brute force skips the graph search entirely. + if (up_front_bf) { + launch_bf(nullptr, nullptr, num_queries); + return; + } + if (num_upper_layers > 0) { auto* d_layer_ptrs = static_cast( idx.d_upper_layer_ptrs); @@ -167,10 +234,15 @@ inline void gpu_hnsw_search( std::to_string(idx.max_degree0) + ")"); } // Fixed overhead: staging (max_staging*8) + parent_ids (sw*4) - // + meta (12) - int smem_overhead = max_staging * 8 + sw * 4 + 12; - // Per-ef cost: 3 result arrays + 3 merge arrays = 6 × 4 = 24 bytes/slot - int max_ef = (smem_max - smem_overhead) / 24; + // + meta (12, or 24 on the filtered path with 6 slots). + int meta_bytes = has_filter ? 24 : 12; + int smem_overhead = max_staging * 8 + sw * 4 + meta_bytes; + // Per-ef cost: 3 result arrays + 3 merge arrays = 6 × 4 = 24 + // bytes/slot. The filtered path adds the invalid frontier (3 more + // arrays); with ef_inv capped at ef (the default) that is a further + // 12 bytes/slot, so divide by 36 to bound ef conservatively. + int per_ef = has_filter ? 36 : 24; + int max_ef = (smem_max - smem_overhead) / per_ef; if (max_ef < 1) { throw std::runtime_error( std::string("gpu_hnsw: search_width=") + @@ -202,8 +274,17 @@ inline void gpu_hnsw_search( } } + // Invalid-frontier capacity: default to the (now finalized) ef, but + // allow the caller to cap it smaller to save shared memory. Capped at + // ef so the /36 clamp above stays conservative. Only used on the + // filtered path. + int ef_inv = 0; + if (has_filter) { + ef_inv = (params.ef_inv > 0) ? std::min(params.ef_inv, ef) : ef; + } + size_t smem_size = hnsw_kernel::calc_layer0_smem_size( - ef, sw, idx.max_degree0); + ef, sw, idx.max_degree0, has_filter ? ef_inv : 0); // Opt into >48 KiB dynamic shared memory when the device supports it; // without this the kernel launch would fail for large ef. This is @@ -221,47 +302,95 @@ inline void gpu_hnsw_search( // the set (thinks it's configured) and fails to launch. Under the lock // the attribute only ever grows, and is always >= the current launch's // requirement before we proceed. - if (smem_size > 49152) { - static std::mutex configured_smem_mutex; - static size_t configured_smem = 0; - std::lock_guard lock(configured_smem_mutex); - if (smem_size > configured_smem) { - GPU_HNSW_CUDA_CHECK(cudaFuncSetAttribute( - hnsw_kernel::layer0_beam_search_kernel< - DataT, QueryT, USE_DP4A>, - cudaFuncAttributeMaxDynamicSharedMemorySize, - static_cast(smem_size))); - configured_smem = 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 - <<>>( - d_layer0_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); - GPU_HNSW_CUDA_CHECK(cudaGetLastError()); + // The graph kernel is templated on HAS_FILTER; the filtered and + // unfiltered instantiations are distinct __global__ functions, so each + // needs its own cudaFuncSetAttribute high-water tracking. run_graph is + // a nested generic lambda so its statics are per-<...,HAS_FILTER>. + auto run_graph = [&]() { + if (smem_size > 49152) { + static std::mutex configured_smem_mutex; + static size_t configured_smem = 0; + std::lock_guard lock(configured_smem_mutex); + if (smem_size > configured_smem) { + GPU_HNSW_CUDA_CHECK(cudaFuncSetAttribute( + hnsw_kernel::layer0_beam_search_kernel< + DataT, + QueryT, + USE_DP4A, + HAS_FILTER>, + cudaFuncAttributeMaxDynamicSharedMemorySize, + static_cast(smem_size))); + configured_smem = smem_size; + } + } + + // Zero the per-query BF worklist counter before the graph launch + // (same stream, so ordered ahead of the kernel that appends to it). + if constexpr (HAS_FILTER) { + if (!disable_bf) { + GPU_HNSW_CUDA_CHECK(cudaMemsetAsync( + sc.d_needs_bf_count, 0, sizeof(int), stream)); + } + } + + const uint8_t* d_bitset = HAS_FILTER ? sc.d_bitset : nullptr; + uint32_t* d_needs_bf = HAS_FILTER ? sc.d_needs_bf : nullptr; + int* d_needs_bf_count = HAS_FILTER ? sc.d_needs_bf_count : nullptr; + + hnsw_kernel::layer0_beam_search_kernel< + DataT, + QueryT, + USE_DP4A, + HAS_FILTER><<>>( + d_layer0_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, + d_bitset, + ef_inv, + kAlpha, + static_cast(live_count), + disable_bf, + d_needs_bf, + d_needs_bf_count); + GPU_HNSW_CUDA_CHECK(cudaGetLastError()); + + // Per-query brute-force fallback: queries whose graph search fell + // short (fewer than k live results) were appended to sc.d_needs_bf + // by the graph kernel. Launch a full-width grid; each block reads + // the device worklist length and early-exits when out of range, so + // no host sync is needed between the two kernels. + if constexpr (HAS_FILTER) { + if (!disable_bf) { + launch_bf(sc.d_needs_bf, sc.d_needs_bf_count, num_queries); + } + } + }; + + if (has_filter) { + run_graph.template operator()(); + } else { + run_graph.template operator()(); + } }; switch (idx.dataset_type) { diff --git a/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearchKernel.cuh b/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearchKernel.cuh index c972d946c..b04cb2654 100644 --- a/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearchKernel.cuh +++ b/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearchKernel.cuh @@ -442,10 +442,168 @@ __device__ __forceinline__ bool bitmap_visit( return (old & bit) == 0; } +// True if row `id` is filtered OUT (deleted / not visible). Matches Knowhere's +// BitsetView byte/bit order: byte id/8, bit id%8, LSB-first. A null bitset (the +// unfiltered fast path) is never consulted — callers gate on HAS_FILTER. +__device__ __forceinline__ bool is_bitset_filtered( + const uint8_t* __restrict__ bitset, + uint32_t id) { + return (bitset[id >> 3] >> (id & 7)) & 1u; +} + +// Serial insert of (id, dist) into a distance-ascending beam of current size +// *count and capacity cap. Evicts the worst (last) element when full and the +// new distance is better; sets expanded[pos]=0 for the freshly inserted node +// and shifts existing expanded flags with their elements so an already-expanded +// node keeps its flag. MUST be called from a single thread (thread 0) — it is +// the serialized post-sort section that mirrors CPU NeighborSetPopList::insert +// (the sequential encounter order the alpha gate depends on). Returns true if +// the node was inserted. +__device__ __forceinline__ bool beam_insert_serial( + uint32_t* ids, + float* dists, + uint32_t* expanded, + int* count, + int cap, + uint32_t id, + float dist) { + int c = *count; + if (c >= cap && !(dist < dists[cap - 1])) { + return false; // full and not better than the current worst + } + // Number of elements retained ahead of the insert: drop the worst if full. + int n = (c < cap) ? c : cap - 1; + int p = n; + while (p > 0 && dists[p - 1] > dist) { + p--; + } + for (int i = n; i > p; i--) { + ids[i] = ids[i - 1]; + dists[i] = dists[i - 1]; + expanded[i] = expanded[i - 1]; + } + ids[p] = id; + dists[p] = dist; + expanded[p] = 0; + *count = (c < cap) ? (c + 1) : cap; + return true; +} + +// Two-tier merge for the filtered path (thread 0 only). Walks the already +// bitonic-sorted staging candidates in ascending-distance order and routes +// each one: +// - not filtered -> result beam (valid results, ef capacity) +// - filtered -> alpha gate; if it fires, admit to the invalid frontier +// only while closer than the valid beam's worst +// (at_search_back_dist: FLT_MAX until the valid beam is +// full, else result_dists[ef-1]). +// The alpha gate (accumulated_alpha, carried per query in meta[4]) rate-limits +// how often filtered nodes become waypoints, matching CPU +// NeighborSetDoublePopList. Deterministic and rate-limiting, but not +// bit-identical to CPU graph-neighbor encounter order — parity is asserted by +// the recall gate, not identical traversal. +__device__ __forceinline__ void filtered_merge_serial( + uint32_t* result_ids, + float* result_dists, + uint32_t* is_expanded, + uint32_t* invalid_ids, + float* invalid_dists, + uint32_t* invalid_exp, + uint32_t* staging_ids, + float* staging_dists, + int* meta, + const uint8_t* __restrict__ bitset, + int ef, + int ef_inv, + int max_staging, + float kAlpha) { + int sc = min(meta[1], max_staging); + int rc = meta[0]; + int ic = meta[3]; + float* alpha_ptr = reinterpret_cast(&meta[4]); + float accumulated_alpha = *alpha_ptr; + + for (int i = 0; i < sc; i++) { + uint32_t id = staging_ids[i]; + if (id == UINT32_MAX) + continue; + float d = staging_dists[i]; + if (!is_bitset_filtered(bitset, id)) { + beam_insert_serial( + result_ids, result_dists, is_expanded, &rc, ef, id, d); + } else { + accumulated_alpha += kAlpha; + if (accumulated_alpha < 1.0f) + continue; + accumulated_alpha -= 1.0f; + float valid_worst = (rc >= ef) ? result_dists[ef - 1] : FLT_MAX; + if (d < valid_worst) { + beam_insert_serial( + invalid_ids, + invalid_dists, + invalid_exp, + &ic, + ef_inv, + id, + d); + } + } + } + meta[0] = rc; + meta[3] = ic; + *alpha_ptr = accumulated_alpha; +} + +// Cross-beam parent selection for the filtered path (thread 0 only). Repeatedly +// takes the globally-closest UNEXPANDED node across the valid result beam and +// the invalid frontier (both distance-ascending), up to search_width parents. +// This is the GPU analog of CPU NeighborSetDoublePopList::pop_based_on_distance +// + has_next: an invalid node is only eligible while it is closer than the +// valid beam's worst (FLT_MAX until the valid beam fills). Returns the number +// of parents chosen; 0 means the search has converged. Filtered parents are +// expanded as waypoints but never enter the top-k (they live only in the +// invalid frontier). +__device__ __forceinline__ int select_parents_cross_beam( + uint32_t* result_ids, + float* result_dists, + uint32_t* is_expanded, + int rc, + uint32_t* invalid_ids, + float* invalid_dists, + uint32_t* invalid_exp, + int ic, + uint32_t* parent_ids, + int ef, + int search_width) { + int num_parents = 0; + float valid_worst = (rc >= ef) ? result_dists[ef - 1] : FLT_MAX; + int ri = 0, ii = 0; + while (num_parents < search_width) { + while (ri < rc && is_expanded[ri]) + ri++; + while (ii < ic && invalid_exp[ii]) + ii++; + bool hasR = (ri < rc); + bool hasI = (ii < ic) && (invalid_dists[ii] < valid_worst); + if (!hasR && !hasI) + break; + if (hasR && (!hasI || result_dists[ri] <= invalid_dists[ii])) { + parent_ids[num_parents++] = result_ids[ri]; + is_expanded[ri] = 1; + ri++; + } else { + parent_ids[num_parents++] = invalid_ids[ii]; + invalid_exp[ii] = 1; + ii++; + } + } + return num_parents; +} + // Templated on the dataset element type (DataT), the query element type // (QueryT: float for the generic path, int8_t for the native DP4A path), and // USE_DP4A which selects the DP4A int8 inner-product distance at compile time. -template +template __global__ void layer0_beam_search_kernel( const QueryT* __restrict__ d_queries, const DataT* __restrict__ d_dataset, @@ -463,7 +621,17 @@ __global__ void layer0_beam_search_kernel( int ef, int search_width, int max_iterations, - bool use_inner_product) { + bool use_inner_product, + // --- Filtered-search args (ignored when HAS_FILTER is false; the + // branches are compiled out via `if constexpr` so the unfiltered + // instantiation is codegen-identical to the append-only kernel). --- + const uint8_t* __restrict__ d_bitset, + int ef_inv, + float kAlpha, + int live_count, + bool disable_bf, + uint32_t* __restrict__ d_needs_bf, + int* __restrict__ d_needs_bf_count) { int query_idx = blockIdx.x; if (query_idx >= num_queries) return; @@ -477,6 +645,11 @@ __global__ void layer0_beam_search_kernel( // has a valid capacity for any max_degree0 (e.g. non-power-of-two 2*M). int max_staging = padded_staging_capacity(search_width, max_degree0); + // The filtered path carries 3 extra meta slots (invalid count + alpha + + // spare) and appends an ef_inv-sized invalid frontier after the merge + // scratch. The unfiltered path keeps exactly today's layout. + constexpr int kMetaSlots = HAS_FILTER ? 6 : 3; + uint32_t* result_ids = reinterpret_cast(smem); float* result_dists = reinterpret_cast(result_ids + ef); uint32_t* is_expanded = reinterpret_cast(result_dists + ef); @@ -485,9 +658,15 @@ __global__ void layer0_beam_search_kernel( uint32_t* parent_ids = reinterpret_cast(staging_dists + max_staging); int* meta = reinterpret_cast(parent_ids + search_width); - uint32_t* merged_ids = reinterpret_cast(meta + 3); + uint32_t* merged_ids = reinterpret_cast(meta + kMetaSlots); float* merged_dists = reinterpret_cast(merged_ids + ef); uint32_t* merged_expanded = reinterpret_cast(merged_dists + ef); + // Invalid frontier (filtered path only): filtered nodes retained as + // waypoints. Never emitted; expanded when they are the globally-closest + // unexpanded node. + uint32_t* invalid_ids = merged_expanded + ef; + float* invalid_dists = reinterpret_cast(invalid_ids + ef_inv); + uint32_t* invalid_exp = reinterpret_cast(invalid_dists + ef_inv); int bitmap_words = (N + 31) / 32; uint32_t* visited_bmap = @@ -502,6 +681,13 @@ __global__ void layer0_beam_search_kernel( meta[0] = 0; meta[1] = 0; meta[2] = 0; + if constexpr (HAS_FILTER) { + meta[3] = 0; // invalid frontier count + // accumulated_alpha, carried per query in meta[4], init 1.0f to + // match CPU NeighborSetDoublePopList (first filtered node always + // admitted once the gate accumulates). + *reinterpret_cast(&meta[4]) = 1.0f; + } } __syncthreads(); @@ -517,10 +703,28 @@ __global__ void layer0_beam_search_kernel( float ep_dist = layer0_distance( query, d_dataset, d_inv_norms, ep, dim, dim4, use_inner_product); - result_ids[0] = ep; - result_dists[0] = ep_dist; - is_expanded[0] = 0; - meta[0] = 1; + if constexpr (HAS_FILTER) { + // A filtered entry point is a waypoint only: seed it into the + // invalid frontier so it is never emitted but can still be + // expanded. Otherwise seed the valid result beam as usual. + if (is_bitset_filtered(d_bitset, ep)) { + invalid_ids[0] = ep; + invalid_dists[0] = ep_dist; + invalid_exp[0] = 0; + meta[3] = 1; + meta[0] = 0; + } else { + result_ids[0] = ep; + result_dists[0] = ep_dist; + is_expanded[0] = 0; + meta[0] = 1; + } + } else { + result_ids[0] = ep; + result_dists[0] = ep_dist; + is_expanded[0] = 0; + meta[0] = 1; + } bitmap_visit(visited_bmap, ep); } else { meta[0] = 0; @@ -553,19 +757,68 @@ __global__ void layer0_beam_search_kernel( } __syncthreads(); - parallel_merge_into_result( - result_ids, result_dists, is_expanded, - staging_ids, staging_dists, - merged_ids, merged_dists, merged_expanded, - meta, ef, max_staging); + if constexpr (HAS_FILTER) { + bitonic_sort_staging( + staging_ids, + staging_dists, + min(meta[1], max_staging), + max_staging); + if (threadIdx.x == 0) { + filtered_merge_serial( + result_ids, + result_dists, + is_expanded, + invalid_ids, + invalid_dists, + invalid_exp, + staging_ids, + staging_dists, + meta, + d_bitset, + ef, + ef_inv, + max_staging, + kAlpha); + } + } else { + parallel_merge_into_result( + result_ids, + result_dists, + is_expanded, + staging_ids, + staging_dists, + merged_ids, + merged_dists, + merged_expanded, + meta, + ef, + max_staging); + } + __syncthreads(); + // Mark the entry point expanded (its neighbors were already seeded above). + // It lives in the valid beam, or in the invalid frontier if it was + // filtered. if (threadIdx.x == 0) { int rc = meta[0]; + bool found = false; for (int i = 0; i < rc; i++) { if (result_ids[i] == ep) { is_expanded[i] = 1; + found = true; break; } } + if constexpr (HAS_FILTER) { + if (!found) { + int ic = meta[3]; + for (int i = 0; i < ic; i++) { + if (invalid_ids[i] == ep) { + invalid_exp[i] = 1; + break; + } + } + } + } } __syncthreads(); @@ -575,10 +828,25 @@ __global__ void layer0_beam_search_kernel( 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 constexpr (HAS_FILTER) { + num_parents = select_parents_cross_beam( + result_ids, + result_dists, + is_expanded, + rc, + invalid_ids, + invalid_dists, + invalid_exp, + meta[3], + parent_ids, + ef, + search_width); + } else { + 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; + } } } @@ -629,11 +897,44 @@ __global__ void layer0_beam_search_kernel( } __syncthreads(); - parallel_merge_into_result( - result_ids, result_dists, is_expanded, - staging_ids, staging_dists, - merged_ids, merged_dists, merged_expanded, - meta, ef, max_staging); + if constexpr (HAS_FILTER) { + bitonic_sort_staging( + staging_ids, + staging_dists, + min(meta[1], max_staging), + max_staging); + if (threadIdx.x == 0) { + filtered_merge_serial( + result_ids, + result_dists, + is_expanded, + invalid_ids, + invalid_dists, + invalid_exp, + staging_ids, + staging_dists, + meta, + d_bitset, + ef, + ef_inv, + max_staging, + kAlpha); + } + __syncthreads(); + } else { + parallel_merge_into_result( + result_ids, + result_dists, + is_expanded, + staging_ids, + staging_dists, + merged_ids, + merged_dists, + merged_expanded, + meta, + ef, + max_staging); + } // Stagnation detected when num_parents==0 (all expanded) → break above } @@ -660,12 +961,40 @@ __global__ void layer0_beam_search_kernel( use_inner_product ? -FLT_MAX : FLT_MAX; } } + + // --- Per-query brute-force fallback (filtered path) --- + // If the filtered graph search returned fewer than k valid results while + // more than that many live rows exist, this query short-fell (its near + // neighbors were filtered out). Append it to the device worklist for a + // batched BF pass, matching CPU HNSW's per-query fallback. Skipped when + // disable_bf mirrors hnsw_cfg.disable_fallback_brute_force (short queries + // then keep their padded sentinels). + if constexpr (HAS_FILTER) { + if (threadIdx.x == 0 && !disable_bf) { + int rc = meta[0]; + int real_topk = min(rc, k); + if (real_topk < k && real_topk < live_count) { + int pos = atomicAdd(d_needs_bf_count, 1); + d_needs_bf[pos] = static_cast(query_idx); + } + } + } } -inline size_t calc_layer0_smem_size(int ef, int search_width, int max_degree0) { +inline size_t calc_layer0_smem_size( + int ef, + int search_width, + int max_degree0, + int ef_inv = 0) { // Must match the kernel's padded staging capacity exactly. int max_staging = padded_staging_capacity(search_width, max_degree0); + // ef_inv > 0 selects the filtered layout: 3 extra meta slots and an + // ef_inv-sized invalid frontier (12 B / ef_inv). ef_inv == 0 is the + // unfiltered layout, byte-identical to the append-only kernel. + bool has_filter = (ef_inv > 0); + int meta_slots = has_filter ? 6 : 3; + size_t size = 0; size += ef * sizeof(uint32_t); // result_ids size += ef * sizeof(float); // result_dists @@ -673,10 +1002,15 @@ inline size_t calc_layer0_smem_size(int ef, int search_width, int max_degree0) { size += max_staging * sizeof(uint32_t); // staging_ids size += max_staging * sizeof(float); // staging_dists size += search_width * sizeof(uint32_t); // parent_ids - size += 3 * sizeof(int); // meta + size += meta_slots * sizeof(int); // meta size += ef * sizeof(uint32_t); // merged_ids size += ef * sizeof(float); // merged_dists size += ef * sizeof(uint32_t); // merged_expanded + if (has_filter) { + size += ef_inv * sizeof(uint32_t); // invalid_ids + size += ef_inv * sizeof(float); // invalid_dists + size += ef_inv * sizeof(uint32_t); // invalid_exp + } return size; } diff --git a/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.cu b/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.cu index 8766cbfb3..75f057edf 100644 --- a/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.cu +++ b/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.cu @@ -99,6 +99,30 @@ void GpuHnswSearchScratch::ensure( } } +void GpuHnswSearchScratch::ensure_filter(int nq, size_t bitset_bytes_needed) { + // Device bitset: reallocate when the segment row count (hence the byte + // count) grows. ensure() is called per search with the current n_rows, so + // a segment that gains rows via reload re-sizes the bitset here. + if (bitset_bytes_needed > bitset_bytes) { + if (d_bitset) + cudaFree(d_bitset); + SCRATCH_CUDA_CHECK(cudaMalloc(&d_bitset, bitset_bytes_needed)); + bitset_bytes = bitset_bytes_needed; + } + // needs_bf worklist: one uint32 query index per query, plus a single + // atomic counter. Sized to nq (grow-only). + if (nq > needs_bf_cap) { + if (d_needs_bf) + cudaFree(d_needs_bf); + SCRATCH_CUDA_CHECK(cudaMalloc( + &d_needs_bf, static_cast(nq) * sizeof(uint32_t))); + needs_bf_cap = nq; + } + if (d_needs_bf_count == nullptr) { + SCRATCH_CUDA_CHECK(cudaMalloc(&d_needs_bf_count, sizeof(int))); + } +} + GpuHnswSearchScratch::~GpuHnswSearchScratch() { // Set the owning device before freeing so cudaFree runs in the correct // context on multi-GPU systems. @@ -115,6 +139,12 @@ GpuHnswSearchScratch::~GpuHnswSearchScratch() { cudaFree(d_visited_bitmaps); if (d_queries_i8) cudaFree(d_queries_i8); + if (d_bitset) + cudaFree(d_bitset); + if (d_needs_bf) + cudaFree(d_needs_bf); + if (d_needs_bf_count) + cudaFree(d_needs_bf_count); } GpuHnswScratchSlot::~GpuHnswScratchSlot() { diff --git a/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.h b/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.h index f0d585c62..2658abde8 100644 --- a/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.h +++ b/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.h @@ -81,6 +81,31 @@ struct GpuHnswSearchParams { int search_width = 4; int max_iterations = 0; int thread_block_size = 0; + + // --- Filtered search (deletes / TTL / partition / visibility bitset) --- + // + // When bitset_data == nullptr the search runs the unfiltered fast path, + // which is codegen-identical to the append-only kernel (the filter + // branches are compiled out via `if constexpr`). When non-null it enables + // CPU-HNSW-parity filtered search: filtered nodes stay graph waypoints but + // are never emitted, an alpha gate rate-limits their expansion, and a + // brute-force fallback guarantees k live results under heavy deletes. + // + // Semantics match Knowhere's BitsetView: a set bit means the row at that + // index is filtered OUT (deleted / not visible). The bit order is + // LSB-first per byte: row r is byte r/8, bit r%8. The index space is the + // storage add-order row id, which the GPU returns directly (see the design + // doc's ID-mapping section and the gpu_hnsw_id_mapping test). + const uint8_t* bitset_data = nullptr; // host ptr, uploaded per search batch + int64_t bitset_nbits = 0; // rows the bitset covers (== n_rows) + int64_t bitset_filtered_count = 0; // popcount: rows filtered out + // Invalid-frontier capacity. 0 => default to ef. Cappable to bound the + // extra shared memory (12 B / ef_inv on top of the 24 B / ef base). + int ef_inv = 0; + // Mirrors Knowhere's hnsw_cfg.disable_fallback_brute_force. When true the + // short-result BF fallback is skipped and short queries keep padded + // sentinels, matching CPU HNSW with fallback disabled. + bool disable_fallback_brute_force = false; }; struct GpuHnswDeviceUpperLayer { @@ -98,12 +123,23 @@ struct GpuHnswSearchScratch { uint32_t* d_visited_bitmaps = nullptr; int8_t* d_queries_i8 = nullptr; // int8 queries for the native DP4A path + // Filtered-search scratch. d_bitset holds the uploaded BitsetView bytes + // (ceil(nbits/8)); d_needs_bf is the device-side worklist of query indices + // whose graph search returned fewer than k live results (per-query BF + // fallback), and d_needs_bf_count is its atomic length. All allocated + // lazily and only when a filtered search runs. + uint8_t* d_bitset = nullptr; + uint32_t* d_needs_bf = nullptr; + int* d_needs_bf_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 queries_i8_bytes = 0; + size_t bitset_bytes = 0; + int needs_bf_cap = 0; // Device this scratch's allocations live on; used to set the CUDA device // context before freeing in the destructor (multi-GPU correctness). @@ -111,6 +147,13 @@ struct GpuHnswSearchScratch { void ensure(int nq, int k, int dim, int N, bool use_i8_queries = false); + // Ensure the per-search filter scratch is large enough: a device bitset of + // `bitset_bytes_needed` bytes (reallocated when n_rows grows) and a + // needs_bf worklist sized for `nq` queries plus its counter. Separate from + // ensure() because filtering is optional and the bitset size tracks the + // segment row count, not nq/k/dim. + void ensure_filter(int nq, size_t bitset_bytes_needed); + ~GpuHnswSearchScratch(); GpuHnswSearchScratch() = default; From 45e2ba26683ec1a62692d7c55dc47da7648e121f Mon Sep 17 00:00:00 2001 From: premal Date: Sat, 18 Jul 2026 21:23:01 +0000 Subject: [PATCH 35/37] test(gpu_hnsw): filtered recall IP/COSINE/int8 + CPU-HNSW baseline + edge cases; vendor ensure_filter device fix Signed-off-by: premal Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/ut/test_gpu_search.cc | 199 ++++++++++++++---- .../faiss/faiss/gpu/impl/GpuHnswTypes.cu | 5 + 2 files changed, 160 insertions(+), 44 deletions(-) diff --git a/tests/ut/test_gpu_search.cc b/tests/ut/test_gpu_search.cc index 76ca55169..db6bf8111 100644 --- a/tests/ut/test_gpu_search.cc +++ b/tests/ut/test_gpu_search.cc @@ -14,8 +14,10 @@ #include #include #include +#include #include #include +#include #include #include "catch2/catch_approx.hpp" @@ -1604,6 +1606,63 @@ gpu_dtype_recall(const std::string& index_type, const knowhere::Json& json, int REQUIRE(gt.has_value()); return GetKNNRecall(*gt.value(), *results.value()); } + +// Filtered-search recall helper. Builds a CPU index of `index_type` on T, +// uploads to GPU_HNSW, then for a given filter `ratio` searches with a random +// delete bitset. Returns {gpu_recall, cpu_recall} where both are measured vs a +// brute-force reference that uses the SAME bitset, so CPU HNSW is the parity +// baseline (not brute force alone). Also asserts no returned GPU id is a +// filtered bit. ratio == 0 exercises the HAS_FILTER=false fast path (empty +// view => data()==nullptr => unfiltered). +template +std::pair +gpu_filtered_recall(const std::string& index_type, const knowhere::Json& json, int version, float ratio, + uint32_t bitset_seed) { + auto train_ds = knowhere::ConvertToDataTypeIfNeeded(GenDataSet(kDtypeNb, kDtypeDim, kDtypeSeed)); + auto query_ds = knowhere::ConvertToDataTypeIfNeeded(GenDataSet(kDtypeNq, kDtypeDim, kDtypeSeed + 2)); + + auto cpu_idx = knowhere::IndexFactory::Instance().Create(index_type, version).value(); + REQUIRE(cpu_idx.Build(train_ds, json) == knowhere::Status::success); + knowhere::BinarySet bs; + REQUIRE(cpu_idx.Serialize(bs) == knowhere::Status::success); + + auto gpu_idx = knowhere::IndexFactory::Instance().Create(knowhere::IndexEnum::INDEX_GPU_HNSW, version).value(); + REQUIRE(gpu_idx.Deserialize(bs) == knowhere::Status::success); + + std::vector bitset_data((kDtypeNb + 7) / 8, 0); + int64_t filtered_count = 0; + std::mt19937 rng(bitset_seed); + std::uniform_real_distribution u(0.0f, 1.0f); + for (int64_t r = 0; r < kDtypeNb; r++) { + if (u(rng) < ratio) { + bitset_data[r / 8] |= static_cast(1 << (r % 8)); + filtered_count++; + } + } + const bool has_filter = filtered_count > 0; + knowhere::BitsetView bitset = + has_filter ? knowhere::BitsetView(bitset_data.data(), kDtypeNb, filtered_count) : knowhere::BitsetView(); + + auto gpu_res = gpu_idx.Search(query_ds, json, bitset); + REQUIRE(gpu_res.has_value()); + const int64_t* gids = gpu_res.value()->GetIds(); + for (int64_t i = 0; i < kDtypeNq * kDtypeTopk; i++) { + const int64_t id = gids[i]; + if (id < 0) { + continue; + } + REQUIRE(id < kDtypeNb); + if (has_filter) { + REQUIRE_FALSE((bitset_data[id / 8] & (1 << (id % 8))) != 0); + } + } + + auto cpu_res = cpu_idx.Search(query_ds, json, bitset); + REQUIRE(cpu_res.has_value()); + auto gt = knowhere::BruteForce::Search(train_ds, query_ds, json, bitset); + REQUIRE(gt.has_value()); + return {GetKNNRecall(*gt.value(), *gpu_res.value()), GetKNNRecall(*gt.value(), *cpu_res.value())}; +} } // namespace // FP32 flat storage -> generic kernel. The config control. @@ -1753,7 +1812,8 @@ TEST_CASE("GPU HNSW IP/COSINE distance sign", "[gpu_hnsw_ip_sign]") { double expected = is_cosine ? dot / (std::sqrt(qn) * std::sqrt(dn) + 1e-12) : dot; double got = dist[q * kDtypeTopk]; REQUIRE(std::fabs(got - expected) <= 1e-2 * (std::fabs(expected) + 1.0)); - if (got > 0.0) positive_top1++; + if (got > 0.0) + positive_top1++; } // For this data the nearest neighbor's similarity is positive for the vast // majority of queries; a sign inversion would make all of them negative. @@ -1819,8 +1879,9 @@ TEST_CASE("GPU HNSW G-ID mapping: node id == row offset == bitset index", "[gpu_ json[knowhere::indexparam::EFCONSTRUCTION] = 200; json[knowhere::indexparam::EF] = 200; - // Ordered training data; each row is its own exact nearest neighbor (L2=0), - // so querying with row i must return id i. + // Randomly generated training data; regardless of order, every row is + // trivially its own exact nearest neighbor (L2 self-distance 0), so an + // unfiltered self-query for row i must return id i. auto train_ds = GenDataSet(nb, dim, 42); auto cpu_idx = @@ -1874,24 +1935,73 @@ TEST_CASE("GPU HNSW G-ID mapping: node id == row offset == bitset index", "[gpu_ } } -// Filtered-search recall/exclusion parity vs brute force across delete ratios. -// Exercises the two-tier beam + alpha gate + brute-force fallback: no returned -// id may be filtered, and recall of the live top-k must track brute force. -TEST_CASE("GPU HNSW filtered search recall vs brute force", "[gpu_hnsw_filtered_recall]") { +// Filtered-search recall/exclusion parity across delete ratios and metrics. +// Exercises the two-tier beam + alpha gate + brute-force fallback. For every +// (metric, ratio) the helper asserts no returned id is a filtered bit; here we +// assert GPU recall tracks the CPU HNSW baseline (both vs the SAME-bitset brute +// force), which is the design's parity requirement. Parity is asserted by +// recall, not bit-identical traversal (the alpha gate is deterministic but not +// CPU encounter-order identical). ratio 0.0 exercises the HAS_FILTER=false fast +// path; 0.3/0.7 cover the mid-range where the alpha-gate order divergence is +// most pronounced; 0.95/0.99 trip the up-front brute-force fallback. +namespace { +void +check_filtered_ratios(const std::string& index_type, const knowhere::Json& json) { const auto version = GenTestVersionList(); - const int64_t nb = 8000; - const int64_t dim = 64; - const int64_t nq = 200; - const int64_t topk = 10; + uint32_t seed = 123; + for (const float ratio : {0.0f, 0.1f, 0.3f, 0.5f, 0.7f, 0.9f, 0.95f, 0.99f}) { + CAPTURE(ratio); + auto [gpu_recall, cpu_recall] = gpu_filtered_recall(index_type, json, version, ratio, seed++); + CAPTURE(gpu_recall); + CAPTURE(cpu_recall); + // Primary parity gate: GPU must not lag CPU HNSW by more than 0.05. + REQUIRE(gpu_recall >= cpu_recall - 0.05f); + if (ratio >= 0.93f) { + // Up-front brute force is exact -> recall should be ~1.0. + REQUIRE(gpu_recall >= 0.90f); + } + } +} +} // namespace - knowhere::Json json; - json[knowhere::meta::DIM] = dim; - json[knowhere::meta::METRIC_TYPE] = knowhere::metric::L2; - json[knowhere::meta::TOPK] = topk; - json[knowhere::indexparam::HNSW_M] = 16; - json[knowhere::indexparam::EFCONSTRUCTION] = 200; - json[knowhere::indexparam::EF] = 200; +TEST_CASE("GPU HNSW filtered search recall (FP32 L2/IP/COSINE)", "[gpu_hnsw_filtered_recall]") { + const auto metric = GENERATE(knowhere::metric::L2, knowhere::metric::IP, knowhere::metric::COSINE); + CAPTURE(metric); + check_filtered_ratios(knowhere::IndexEnum::INDEX_HNSW, dtype_json(metric, "")); +} + +// INT8 native storage: IP/COSINE with dim%4==0 take the DP4A brute-force + +// graph path; L2 takes the generic decoded path. Confirms the filter/BF kernels +// are threaded through searchHostInt8, not just the fp32 searchHost. +TEST_CASE("GPU HNSW filtered search recall (INT8 native L2/IP/COSINE)", "[gpu_hnsw_filtered_recall_int8]") { + const auto version = GenTestVersionList(); + const auto metric = GENERATE(knowhere::metric::L2, knowhere::metric::IP, knowhere::metric::COSINE); + CAPTURE(metric); + uint32_t seed = 777; + for (const float ratio : {0.0f, 0.5f, 0.9f, 0.99f}) { + CAPTURE(ratio); + auto [gpu_recall, cpu_recall] = gpu_filtered_recall( + knowhere::IndexEnum::INDEX_HNSW, dtype_json(metric, ""), version, ratio, seed++); + CAPTURE(gpu_recall); + CAPTURE(cpu_recall); + REQUIRE(gpu_recall >= cpu_recall - 0.05f); + if (ratio >= 0.93f) { + REQUIRE(gpu_recall >= 0.90f); + } + } +} +// Edge cases the ratio sweep does not cover: everything filtered (no live rows) +// and k larger than the live-row count. Both must return sentinels for the +// missing slots rather than a filtered id or an error. +TEST_CASE("GPU HNSW filtered search edge cases", "[gpu_hnsw_filtered_edge]") { + const auto version = GenTestVersionList(); + const int64_t nb = 4000; + const int64_t dim = 64; + const int64_t nq = 32; + + auto json = dtype_json(knowhere::metric::L2, ""); + json[knowhere::meta::TOPK] = 10; auto train_ds = GenDataSet(nb, dim, 42); auto query_ds = GenDataSet(nq, dim, 44); @@ -1900,43 +2010,44 @@ TEST_CASE("GPU HNSW filtered search recall vs brute force", "[gpu_hnsw_filtered_ REQUIRE(cpu_idx.Build(train_ds, json) == knowhere::Status::success); knowhere::BinarySet bs; REQUIRE(cpu_idx.Serialize(bs) == knowhere::Status::success); - auto gpu_idx = knowhere::IndexFactory::Instance().Create(knowhere::IndexEnum::INDEX_GPU_HNSW, version).value(); REQUIRE(gpu_idx.Deserialize(bs) == knowhere::Status::success); - std::mt19937 rng(123); - for (const float ratio : {0.1f, 0.5f, 0.9f, 0.95f, 0.99f}) { - std::vector bitset_data((nb + 7) / 8, 0); - int64_t filtered_count = 0; - std::uniform_real_distribution u(0.0f, 1.0f); - for (int64_t r = 0; r < nb; r++) { - if (u(rng) < ratio) { - bitset_data[r / 8] |= static_cast(1 << (r % 8)); - filtered_count++; - } + // 1) All rows filtered: every result slot must be a sentinel (-1), never a + // filtered id, and the search must not error. + { + std::vector bitset_data((nb + 7) / 8, 0xFF); + knowhere::BitsetView bitset(bitset_data.data(), nb, nb); + auto res = gpu_idx.Search(query_ds, json, bitset); + REQUIRE(res.has_value()); + const int64_t* ids = res.value()->GetIds(); + for (int64_t i = 0; i < nq * 10; i++) { + REQUIRE(ids[i] == -1); } - knowhere::BitsetView bitset(bitset_data.data(), nb, filtered_count); + } + // 2) k > live_count: only a handful of rows survive; the first `live` slots + // per query must be live ids, the remainder sentinels. + { + const int64_t live = 3; + std::vector bitset_data((nb + 7) / 8, 0xFF); + for (int64_t r = 0; r < live; r++) { + bitset_data[r / 8] &= static_cast(~(1 << (r % 8))); + } + knowhere::BitsetView bitset(bitset_data.data(), nb, nb - live); auto res = gpu_idx.Search(query_ds, json, bitset); REQUIRE(res.has_value()); const int64_t* ids = res.value()->GetIds(); - - // No returned id may be filtered out. - for (int64_t i = 0; i < nq * topk; i++) { - const int64_t id = ids[i]; - if (id == -1) { - continue; + for (int64_t q = 0; q < nq; q++) { + for (int64_t j = 0; j < 10; j++) { + const int64_t id = ids[q * 10 + j]; + if (id == -1) { + continue; + } + REQUIRE(id < live); } - REQUIRE_FALSE((bitset_data[id / 8] & (1 << (id % 8))) != 0); } - - auto gt = knowhere::BruteForce::Search(train_ds, query_ds, json, bitset); - REQUIRE(gt.has_value()); - float recall = GetKNNRecall(*gt.value(), *res.value()); - // Parity is asserted by recall, not bit-identical traversal (the alpha - // gate is deterministic but not CPU encounter-order identical). - REQUIRE(recall >= 0.85f); } } #endif diff --git a/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.cu b/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.cu index 75f057edf..d35e20eb7 100644 --- a/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.cu +++ b/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.cu @@ -100,6 +100,11 @@ void GpuHnswSearchScratch::ensure( } void GpuHnswSearchScratch::ensure_filter(int nq, size_t bitset_bytes_needed) { + // Bind this scratch slot's owning device before allocating so cudaMalloc + // lands on the right GPU even if the active device was changed elsewhere + // (matches the destructor). searchHost already sets it, but keep the + // allocation self-contained on multi-GPU systems. + cudaSetDevice(device); // Device bitset: reallocate when the segment row count (hence the byte // count) grows. ensure() is called per search with the current n_rows, so // a segment that gains rows via reload re-sizes the bitset here. From 407b60958cc9c1f85d41441313fa352fe06a184b Mon Sep 17 00:00:00 2001 From: premal Date: Sun, 19 Jul 2026 23:52:12 +0000 Subject: [PATCH 36/37] feat(gpu_hnsw): re-vendor faiss GPU tree (warp-cooperative layer-0 distance, H1) Re-vendors faiss gpu-hnsw-faiss tip f22116e8 (merge of faiss PR #8): the layer-0 neighbor expansion is now warp-per-candidate cooperative distance (32 lanes stride one vector, warp-reduce) instead of thread-per-candidate, coalescing the dominant dataset loads. Measured on the L40S build box (counter deltas transfer to Blackwell): sectors/request 12.1 -> 3.10 (3.9x), throughput 1.57x int8-DP4A / 1.85x fp32, recall identical to baseline (7/7, max delta -0.0025), compute-sanitizer clean (memcheck/racecheck/initcheck). No public API or semantic change. Signed-off-by: premal Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../faiss/faiss/gpu/impl/GpuHnswSearch.cuh | 7 + .../faiss/gpu/impl/GpuHnswSearchKernel.cuh | 178 ++++++++++++++---- 2 files changed, 151 insertions(+), 34 deletions(-) diff --git a/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearch.cuh b/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearch.cuh index afab6fecf..2c418b036 100644 --- a/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearch.cuh +++ b/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearch.cuh @@ -225,6 +225,13 @@ inline void gpu_hnsw_search( if (block_size < max_staging) { block_size = max_staging; } + // The layer-0 expansion is warp-cooperative (one warp per candidate + // edge, 32 lanes striding the vector for coalesced loads), so the + // block must be a whole number of warps for the full-mask __shfl + // reductions to be well-defined. The default (128) and the + // power-of-two staging bump are already multiples of 32; this only + // rounds up a caller-supplied non-multiple thread_block_size. + block_size = ((block_size + 31) / 32) * 32; if (block_size > 1024) { throw std::runtime_error( std::string("gpu_hnsw: padded staging capacity ") + diff --git a/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearchKernel.cuh b/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearchKernel.cuh index b04cb2654..66ab00709 100644 --- a/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearchKernel.cuh +++ b/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearchKernel.cuh @@ -145,6 +145,75 @@ __device__ __forceinline__ float layer0_distance( } } +__device__ __forceinline__ float warp_reduce_sum(float v) { +#pragma unroll + for (int off = 16; off > 0; off >>= 1) + v += __shfl_down_sync(0xffffffff, v, off); + return v; +} + +__device__ __forceinline__ int warp_reduce_sum(int v) { +#pragma unroll + for (int off = 16; off > 0; off >>= 1) + v += __shfl_down_sync(0xffffffff, v, off); + return v; +} + +// Warp-cooperative layer-0 distance. The 32 lanes stride the vector (lane d +// reads elements d, d+32, ...), so the dataset loads coalesce into 1 +// sector/request instead of the ~12 the thread-per-candidate path produced +// (see the perf analysis: uncoalesced loads were 49% of warp stalls / DRAM at +// 38% of peak). The per-lane partial sums are warp-reduced; the returned +// distance is valid on lane 0 only (the reduction accumulates there), so +// callers that need it on other lanes must broadcast. Same convention as +// layer0_distance() (smaller == closer, IP negated, COSINE scaled by inv_norm). +template +__device__ __forceinline__ float layer0_distance_warp( + const QueryT* __restrict__ query, + const DataT* __restrict__ d_dataset, + const float* __restrict__ d_inv_norms, + uint32_t id, + int dim, + int dim4, + bool use_inner_product, + int lane) { + if constexpr (USE_DP4A) { + const int32_t* vec_packed = reinterpret_cast( + d_dataset + static_cast(id) * dim); + const int32_t* query_packed = + reinterpret_cast(query); + int sum = 0; + for (int d = lane; d < dim4; d += 32) { + sum = __dp4a(query_packed[d], vec_packed[d], sum); + } + sum = warp_reduce_sum(sum); + float dist = static_cast(-sum); + if (d_inv_norms) + dist *= __ldg(&d_inv_norms[id]); + return dist; + } else { + const DataT* vec = d_dataset + static_cast(id) * dim; + if (use_inner_product) { + float sum = 0.0f; + for (int d = lane; d < dim; d += 32) { + sum += query[d] * load_elem(vec, d); + } + sum = warp_reduce_sum(sum); + float dist = -sum; + if (d_inv_norms) + dist *= __ldg(&d_inv_norms[id]); + return dist; + } else { + float sum = 0.0f; + for (int d = lane; d < dim; d += 32) { + float diff = query[d] - load_elem(vec, d); + sum += diff * diff; + } + return warp_reduce_sum(sum); + } + } +} + // ============================================================================ // Phase 1: Upper-layer greedy search // ============================================================================ @@ -639,6 +708,15 @@ __global__ void layer0_beam_search_kernel( const QueryT* query = d_queries + static_cast(query_idx) * dim; int dim4 = dim / 4; + // Warp-cooperative neighbor expansion: one warp owns one candidate edge and + // its 32 lanes compute the distance together (coalesced loads). The host + // rounds block_size up to a whole number of warps so every lane below is + // part of a full 32-lane warp (required for the full-mask __shfl reductions + // in layer0_distance_warp). + const int lane = threadIdx.x & 31; + const int warp_in_block = threadIdx.x >> 5; + const int num_warps = blockDim.x >> 5; + extern __shared__ char smem[]; // Pad to a power of two so the bitonic sort in parallel_merge_into_result @@ -737,22 +815,37 @@ __global__ void layer0_beam_search_kernel( meta[1] = 0; __syncthreads(); - for (int j = threadIdx.x; ep_valid && 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)) + for (int j = warp_in_block; ep_valid && j < max_degree0; j += num_warps) { + // Lane 0 resolves the neighbor and claims the first-visit; the go/no-go + // and id are broadcast so the whole warp stays convergent for the + // cooperative distance below. `proceed` is warp-uniform, so the + // `continue` never splits the warp (no __shfl hazard). + uint32_t nbr = UINT32_MAX; + bool proceed = false; + if (lane == 0) { + uint32_t cand = + d_layer0_graph[static_cast(ep) * max_degree0 + j]; + if (cand != UINT32_MAX && cand < static_cast(N) && + bitmap_visit(visited_bmap, cand)) { + nbr = cand; + proceed = true; + } + } + proceed = __shfl_sync(0xffffffff, proceed, 0); + if (!proceed) continue; + nbr = __shfl_sync(0xffffffff, nbr, 0); - float dist = layer0_distance( + float dist = layer0_distance_warp( query, d_dataset, d_inv_norms, nbr, dim, dim4, - use_inner_product); + use_inner_product, lane); - int slot = atomicAdd(&meta[1], 1); - if (slot < max_staging) { - staging_ids[slot] = nbr; - staging_dists[slot] = dist; + if (lane == 0) { + int slot = atomicAdd(&meta[1], 1); + if (slot < max_staging) { + staging_ids[slot] = nbr; + staging_dists[slot] = dist; + } } } __syncthreads(); @@ -863,36 +956,53 @@ __global__ void layer0_beam_search_kernel( __syncthreads(); int total_work = num_parents * max_degree0; - for (int wi = threadIdx.x; wi < total_work; wi += blockDim.x) { + for (int wi = warp_in_block; wi < total_work; wi += num_warps) { int parent_idx = wi / max_degree0; int nbr_slot = wi % max_degree0; uint32_t parent = parent_ids[parent_idx]; - // Defensive: a parent id must be a real node (< N). result_ids only - // ever holds the entry point or graph neighbors (both < N) once the - // merge is correct, so this never fires on the healthy path; it caps - // a stray/garbage id at O(N) instead of letting it index the graph - // at parent*max_degree0 and fault far out of bounds. Mirrors the - // neighbor guard below. - if (parent >= static_cast(N)) - continue; - 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)) + + // Lane 0 resolves the neighbor id, bounds-checks it, and claims the + // first-visit; the go/no-go and id are broadcast so the warp stays + // convergent for the cooperative distance. All predicates below are + // warp-uniform (`parent` is identical across lanes), so the + // `continue` never splits the warp and the __shfl calls in + // layer0_distance_warp always see a full 32-lane mask. + uint32_t nbr = UINT32_MAX; + bool proceed = false; + if (lane == 0) { + // Defensive: a parent id must be a real node (< N). result_ids + // only ever holds the entry point or graph neighbors (both < N) + // once the merge is correct, so this never fires on the healthy + // path; it caps a stray/garbage id at O(N) instead of letting it + // index the graph at parent*max_degree0 and fault far out of + // bounds. Mirrors the neighbor guard below. + if (parent < static_cast(N)) { + uint32_t cand = d_layer0_graph + [static_cast(parent) * max_degree0 + + nbr_slot]; + if (cand != UINT32_MAX && cand < static_cast(N) && + bitmap_visit(visited_bmap, cand)) { + nbr = cand; + proceed = true; + } + } + } + proceed = __shfl_sync(0xffffffff, proceed, 0); + if (!proceed) continue; + nbr = __shfl_sync(0xffffffff, nbr, 0); - float dist = layer0_distance( + float dist = layer0_distance_warp( query, d_dataset, d_inv_norms, nbr, dim, dim4, - use_inner_product); + use_inner_product, lane); - int slot = atomicAdd(&meta[1], 1); - if (slot < max_staging) { - staging_ids[slot] = nbr; - staging_dists[slot] = dist; + if (lane == 0) { + int slot = atomicAdd(&meta[1], 1); + if (slot < max_staging) { + staging_ids[slot] = nbr; + staging_dists[slot] = dist; + } } } __syncthreads(); From 2b52090d70c5d61081e23438cea6107cff633a62 Mon Sep 17 00:00:00 2001 From: premal Date: Mon, 20 Jul 2026 04:57:32 +0000 Subject: [PATCH 37/37] feat(gpu_hnsw): re-vendor faiss GPU tree (nq-chunked visited bitmap, OOM fix) Re-vendors the faiss GPU HNSW fix that bounds the grow-only visited bitmap by processing queries in nq-chunks (byte-identical to 6si/faiss devin/gpu-hnsw-oom-chunk). The bitmap was nq * ceil(N/32) * 4 bytes per scratch slot and never released, so high search concurrency on large segments exhausted device VRAM. Chunking caps it regardless of batch size / concurrency; small segments still run a single pass. Adds tests/ut/test_gpu_search.cc [gpu_hnsw_chunk_parity]: runs each search twice (default cap => single pass, vs GPU_HNSW_BITMAP_BYTES forced tiny => many chunks) and asserts identical ids/distances across fp32 & int8-DP4A, unfiltered, mid-ratio filtered, and up-front-BF (99%) paths for L2/IP/COSINE. Signed-off-by: premal Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/ut/test_gpu_search.cc | 99 ++++++++ .../faiss/faiss/gpu/impl/GpuHnswSearch.cuh | 240 ++++++++++-------- .../faiss/faiss/gpu/impl/GpuHnswTypes.cu | 9 +- .../faiss/faiss/gpu/impl/GpuHnswTypes.h | 56 ++++ 4 files changed, 304 insertions(+), 100 deletions(-) diff --git a/tests/ut/test_gpu_search.cc b/tests/ut/test_gpu_search.cc index db6bf8111..adaf63ff7 100644 --- a/tests/ut/test_gpu_search.cc +++ b/tests/ut/test_gpu_search.cc @@ -2050,4 +2050,103 @@ TEST_CASE("GPU HNSW filtered search edge cases", "[gpu_hnsw_filtered_edge]") { } } } + +// Visited-bitmap OOM fix: the layer-0 search processes queries in nq-chunks so +// the grow-only visited bitmap (nq * ceil(N/32) * 4 bytes per scratch slot) is +// bounded regardless of batch size / search concurrency. Chunking must be +// result-invariant: each query is independent, and every per-query buffer +// (queries, entry points, neighbors, distances, visited bitmap, BF worklist) is +// indexed by the chunk-local block index against a caller-offset base pointer. +// This runs the SAME search twice -- once with the default cap (single pass) and +// once with GPU_HNSW_BITMAP_BYTES forced small enough to split the batch into +// many chunks -- and asserts the ids and distances match element-for-element. +namespace { +template +void +check_chunk_parity(const std::string& index_type, const knowhere::Json& json, const knowhere::BitsetView& bitset) { + using Catch::Approx; + const auto version = GenTestVersionList(); + auto train_ds = knowhere::ConvertToDataTypeIfNeeded(GenDataSet(kDtypeNb, kDtypeDim, kDtypeSeed)); + auto query_ds = knowhere::ConvertToDataTypeIfNeeded(GenDataSet(kDtypeNq, kDtypeDim, kDtypeSeed + 2)); + + auto cpu_idx = knowhere::IndexFactory::Instance().Create(index_type, version).value(); + REQUIRE(cpu_idx.Build(train_ds, json) == knowhere::Status::success); + knowhere::BinarySet bs; + REQUIRE(cpu_idx.Serialize(bs) == knowhere::Status::success); + auto gpu_idx = knowhere::IndexFactory::Instance().Create(knowhere::IndexEnum::INDEX_GPU_HNSW, version).value(); + REQUIRE(gpu_idx.Deserialize(bs) == knowhere::Status::success); + + // Baseline: default cap => chunk == nq (single pass). + ::unsetenv("GPU_HNSW_BITMAP_BYTES"); + ::unsetenv("GPU_HNSW_BITMAP_MB"); + auto base = gpu_idx.Search(query_ds, json, bitset); + REQUIRE(base.has_value()); + const int64_t n_out = static_cast(kDtypeNq) * kDtypeTopk; + std::vector base_ids(base.value()->GetIds(), base.value()->GetIds() + n_out); + std::vector base_dist(base.value()->GetDistance(), base.value()->GetDistance() + n_out); + + // Force chunking: cap the bitmap to ~2 queries' worth so the batch is split + // into ~nq/2 chunks, exercising the multi-launch offset math and the + // per-chunk bitmap/worklist resets. + const long per_query = static_cast((kDtypeNb + 31) / 32) * 4; + ::setenv("GPU_HNSW_BITMAP_BYTES", std::to_string(2 * per_query).c_str(), 1); + auto chunked = gpu_idx.Search(query_ds, json, bitset); + ::unsetenv("GPU_HNSW_BITMAP_BYTES"); + REQUIRE(chunked.has_value()); + const int64_t* cids = chunked.value()->GetIds(); + const float* cdist = chunked.value()->GetDistance(); + + for (int64_t i = 0; i < n_out; i++) { + REQUIRE(cids[i] == base_ids[i]); + REQUIRE(cdist[i] == Approx(base_dist[i]).epsilon(1e-5)); + } +} + +std::vector +make_random_bitset(float ratio, uint32_t seed, int64_t& filtered_count) { + std::vector bd((kDtypeNb + 7) / 8, 0); + filtered_count = 0; + std::mt19937 rng(seed); + std::uniform_real_distribution u(0.0f, 1.0f); + for (int64_t r = 0; r < kDtypeNb; r++) { + if (u(rng) < ratio) { + bd[r / 8] |= static_cast(1 << (r % 8)); + filtered_count++; + } + } + return bd; +} +} // namespace + +TEST_CASE("GPU HNSW nq-chunking is result-invariant (visited-bitmap OOM fix)", "[gpu_hnsw_chunk_parity]") { + const auto metric = GENERATE(knowhere::metric::L2, knowhere::metric::IP, knowhere::metric::COSINE); + CAPTURE(metric); + + // Unfiltered (HAS_FILTER=false fast path). + SECTION("fp32 unfiltered") { + check_chunk_parity(knowhere::IndexEnum::INDEX_HNSW, dtype_json(metric, ""), + knowhere::BitsetView()); + } + // int8: IP/COSINE with dim%4==0 take the DP4A path, where the layer-0 (int8) + // and upper-layer (fp32) query buffers are distinct, so BOTH chunk offsets + // must be correct; L2 takes the generic decoded path. + SECTION("int8 unfiltered") { + check_chunk_parity(knowhere::IndexEnum::INDEX_HNSW, dtype_json(metric, ""), + knowhere::BitsetView()); + } + // Mid-ratio filter: HAS_FILTER=true two-tier beam + per-chunk worklist reset. + SECTION("fp32 filtered (~30%)") { + int64_t fc = 0; + auto bd = make_random_bitset(0.3f, 7, fc); + knowhere::BitsetView bitset(bd.data(), kDtypeNb, fc); + check_chunk_parity(knowhere::IndexEnum::INDEX_HNSW, dtype_json(metric, ""), bitset); + } + // High-ratio filter: trips the up-front brute-force fallback (also chunked). + SECTION("fp32 filtered (~99% -> up-front BF)") { + int64_t fc = 0; + auto bd = make_random_bitset(0.99f, 9, fc); + knowhere::BitsetView bitset(bd.data(), kDtypeNb, fc); + check_chunk_parity(knowhere::IndexEnum::INDEX_HNSW, dtype_json(metric, ""), bitset); + } +} #endif diff --git a/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearch.cuh b/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearch.cuh index 2c418b036..df128d4c8 100644 --- a/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearch.cuh +++ b/thirdparty/faiss/faiss/gpu/impl/GpuHnswSearch.cuh @@ -126,18 +126,24 @@ inline void gpu_hnsw_search( const QueryT* d_layer0_queries) { int N_int = static_cast(idx.n_rows); - // Brute-force launcher (shares the current dtype specialization). The - // block must be a power of two <= kBruteForceMaxBlock for the tree - // reduction. Grid is always num_queries: the up-front path uses the - // identity mapping, the per-query fallback reads its worklist length - // from d_num on the device (no host sync between graph and BF). + // Brute-force launcher (shares the current dtype specialization). It + // operates on a chunk of `num_items` queries whose scratch pointers + // (q0/nb0/ds0) are already offset by the caller; when non-null the + // worklist holds chunk-local query indices, matching the offset outputs. + // Grid is num_items: the up-front path uses the identity mapping, the + // per-query fallback reads its worklist length from d_num on the device + // (no host sync between graph and BF). Block must be a power of two <= + // kBruteForceMaxBlock for the tree reduction. int bf_block = 128; - auto launch_bf = [&](const uint32_t* worklist, + auto launch_bf = [&](const QueryT* q0, + uint64_t* nb0, + float* ds0, + const uint32_t* worklist, const int* d_num, int num_items) { hnsw_kernel::brute_force_topk_kernel - <<>>( - d_layer0_queries, + <<>>( + q0, d_data, d_inv_norms, sc.d_bitset, @@ -148,53 +154,14 @@ inline void gpu_hnsw_search( dim, k, idx.use_ip, - sc.d_neighbors, - sc.d_distances); + nb0, + ds0); GPU_HNSW_CUDA_CHECK(cudaGetLastError()); }; - // Up-front brute force skips the graph search entirely. - if (up_front_bf) { - launch_bf(nullptr, nullptr, num_queries); - return; - } - - 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)); - // 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)); - } - + // --- Shared-memory sizing / ef finalization --- + // Query-count independent (depends only on ef/search_width/M and the + // device smem limit), so compute it once before the chunk loop. int block_size = params.thread_block_size > 0 ? params.thread_block_size : 128; @@ -215,13 +182,13 @@ inline void gpu_hnsw_search( } } + // The bitonic sort in the parallel merge needs a power-of-two staging + // capacity, one thread per slot. Pad up to the next power of two + // (handles non-power-of-two 2*M) and grow the block so every staging + // slot is owned by a thread. + int max_staging = + hnsw_kernel::padded_staging_capacity(sw, idx.max_degree0); { - // The bitonic sort in the parallel merge needs a power-of-two - // staging capacity, one thread per slot. Pad up to the next power - // of two (handles non-power-of-two 2*M) and grow the block so every - // staging slot is owned by a thread. - int max_staging = hnsw_kernel::padded_staging_capacity( - sw, idx.max_degree0); if (block_size < max_staging) { block_size = max_staging; } @@ -293,33 +260,35 @@ inline void gpu_hnsw_search( size_t smem_size = hnsw_kernel::calc_layer0_smem_size( ef, sw, idx.max_degree0, has_filter ? ef_inv : 0); - // Opt into >48 KiB dynamic shared memory when the device supports it; - // without this the kernel launch would fail for large ef. This is - // cached per kernel instantiation: cudaFuncSetAttribute configures the - // kernel's max dynamic shared memory globally, so it only needs to run - // once per (kernel, high-water size) rather than on every search. The - // statics live in this generic-lambda body, which is instantiated once - // per combination. - // - // The mutex makes the check-set-store atomic and monotonic: without it, - // two concurrent searches with different smem sizes can interleave so a - // smaller-ef search's cudaFuncSetAttribute runs *after* a larger one, - // downgrading the kernel's global attribute below what a recorded - // high-water mark implies, and a later intermediate search then skips - // the set (thinks it's configured) and fails to launch. Under the lock - // the attribute only ever grows, and is always >= the current launch's - // requirement before we proceed. - 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)); - - // The graph kernel is templated on HAS_FILTER; the filtered and - // unfiltered instantiations are distinct __global__ functions, so each - // needs its own cudaFuncSetAttribute high-water tracking. run_graph is - // a nested generic lambda so its statics are per-<...,HAS_FILTER>. - auto run_graph = [&]() { + // Graph search + per-query brute-force fallback for one chunk of `cnq` + // queries. The scratch pointers (q0/ep0/nb0/ds0) are offset by the + // caller so the kernels' chunk-local blockIdx.x maps to the right + // global query; the visited bitmap is the base buffer, indexed by the + // same chunk-local index. The graph kernel is templated on HAS_FILTER; + // the filtered and unfiltered instantiations are distinct __global__ + // functions, so each needs its own cudaFuncSetAttribute high-water + // tracking (the statics below are per-<...,HAS_FILTER>). + auto run_graph = [&]( + const QueryT* q0, + uint32_t* ep0, + uint64_t* nb0, + float* ds0, + int cnq) { + // Opt into >48 KiB dynamic shared memory when the device supports + // it; without this the kernel launch would fail for large ef. This + // is cached per kernel instantiation: cudaFuncSetAttribute + // configures the kernel's max dynamic shared memory globally, so it + // only needs to run once per (kernel, high-water size). + // + // The mutex makes the check-set-store atomic and monotonic: without + // it, two concurrent searches with different smem sizes can + // interleave so a smaller-ef search's cudaFuncSetAttribute runs + // *after* a larger one, downgrading the kernel's global attribute + // below what a recorded high-water mark implies, and a later + // intermediate search then skips the set (thinks it's configured) + // and fails to launch. Under the lock the attribute only ever grows, + // and is always >= the current launch's requirement before we + // proceed. if (smem_size > 49152) { static std::mutex configured_smem_mutex; static size_t configured_smem = 0; @@ -354,16 +323,16 @@ inline void gpu_hnsw_search( DataT, QueryT, USE_DP4A, - HAS_FILTER><<>>( - d_layer0_queries, + HAS_FILTER><<>>( + q0, d_data, d_inv_norms, idx.d_layer0_graph, - sc.d_entry_points, + ep0, sc.d_visited_bitmaps, - sc.d_neighbors, - sc.d_distances, - num_queries, + nb0, + ds0, + cnq, N_int, dim, idx.max_degree0, @@ -383,20 +352,93 @@ inline void gpu_hnsw_search( // Per-query brute-force fallback: queries whose graph search fell // short (fewer than k live results) were appended to sc.d_needs_bf - // by the graph kernel. Launch a full-width grid; each block reads - // the device worklist length and early-exits when out of range, so - // no host sync is needed between the two kernels. + // (chunk-local indices) by the graph kernel. Launch a full-width + // grid; each block reads the device worklist length and early-exits + // when out of range, so no host sync is needed between the two + // kernels. if constexpr (HAS_FILTER) { if (!disable_bf) { - launch_bf(sc.d_needs_bf, sc.d_needs_bf_count, num_queries); + launch_bf(q0, nb0, ds0, sc.d_needs_bf, + sc.d_needs_bf_count, cnq); } } }; - if (has_filter) { - run_graph.template operator()(); - } else { - run_graph.template operator()(); + // --- Query chunking to bound the visited-bitmap VRAM --- + // Process the batch in chunks of at most chunk_nq queries so the visited + // bitmap (sized in GpuHnswSearchScratch::ensure with the same chunk) + // stays within the VRAM cap regardless of nq / search concurrency. For + // small segments chunk_nq == num_queries -> a single pass, identical to + // the pre-chunk behavior. + int chunk_nq = gpu_hnsw_bitmap_chunk(num_queries, N_int); + for (int c0 = 0; c0 < num_queries; c0 += chunk_nq) { + int cnq = std::min(chunk_nq, num_queries - c0); + // Chunk-offset scratch pointers. q0 is the layer-0 query buffer + // (int8 on the DP4A path, else fp32); fq0 is always the fp32 buffer + // used by the upper-layer greedy descent. + const QueryT* q0 = + d_layer0_queries + static_cast(c0) * dim; + const float* fq0 = sc.d_queries + static_cast(c0) * dim; + uint32_t* ep0 = sc.d_entry_points + c0; + uint64_t* nb0 = sc.d_neighbors + static_cast(c0) * k; + float* ds0 = sc.d_distances + static_cast(c0) * k; + + // Up-front brute force skips the graph search entirely (does not + // touch the visited bitmap, but still runs per chunk so its outputs + // land at the right offsets). + if (up_front_bf) { + launch_bf(q0, nb0, ds0, nullptr, nullptr, cnq); + continue; + } + + 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 = + (cnq + warps_per_block - 1) / warps_per_block; + + hnsw_kernel::upper_layer_search_kernel + <<>>( + fq0, + d_data, + d_inv_norms, + d_layer_ptrs, + ep0, + idx.entry_point, + cnq, + dim, + num_upper_layers, + idx.use_ip); + GPU_HNSW_CUDA_CHECK(cudaGetLastError()); + } else { + std::vector h_eps(cnq, idx.entry_point); + GPU_HNSW_CUDA_CHECK(cudaMemcpyAsync( + ep0, + h_eps.data(), + static_cast(cnq) * 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)); + } + + // Zero this chunk's visited bitmap (reused across chunks). + size_t bitmap_bytes = + hnsw_kernel::calc_visited_bitmap_size(cnq, N_int); + GPU_HNSW_CUDA_CHECK(cudaMemsetAsync( + sc.d_visited_bitmaps, 0, bitmap_bytes, stream)); + + if (has_filter) { + run_graph.template operator()(q0, ep0, nb0, ds0, cnq); + } else { + run_graph.template operator()(q0, ep0, nb0, ds0, cnq); + } } }; diff --git a/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.cu b/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.cu index d35e20eb7..1d6ca4d2b 100644 --- a/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.cu +++ b/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.cu @@ -79,9 +79,16 @@ void GpuHnswSearchScratch::ensure( static_cast(nq) * sizeof(uint32_t))); entry_cap = nq; } + // The visited bitmap is only ever indexed by the chunk-local query index + // (see the launch loop in gpu_hnsw_search), so size it for one chunk rather + // than the full batch. This is the fix for the grow-only OOM: without the + // chunk cap, a large nq (or high-concurrency high-water mark) sized this at + // nq * ceil(N/32) * 4 and never released it. Must match the chunk used at + // launch time. + int bm_nq = gpu_hnsw_bitmap_chunk(nq, N); int bitmap_words = (N + 31) / 32; size_t need_bm = - static_cast(nq) * bitmap_words * sizeof(uint32_t); + static_cast(bm_nq) * bitmap_words * sizeof(uint32_t); if (need_bm > bitmap_bytes) { if (d_visited_bitmaps) cudaFree(d_visited_bitmaps); diff --git a/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.h b/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.h index 2658abde8..3778da7ba 100644 --- a/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.h +++ b/thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.h @@ -25,9 +25,11 @@ #include +#include #include #include #include +#include #include #include #include @@ -76,6 +78,60 @@ enum class GpuHnswDatasetType { BF16 = 3, }; +// Cap (bytes) on the per-search visited bitmap for a single scratch slot. The +// bitmap is nq * ceil(N/32) * 4 bytes and is grow-only per slot, so a large +// query batch or high search concurrency (one bitmap per pool slot) can grow it +// until it exhausts device memory. Observed regression: 16 concurrent batch=512 +// searches on a 538M-row segment grew the pool to ~97 of ~98 GB and OOM'd every +// subsequent allocation. Bounding the bitmap and processing queries in +// nq-chunks caps it regardless of batch size / concurrency. +// +// Tunable via environment (re-read on each call so tests can toggle it; getenv +// is negligible next to a GPU search): GPU_HNSW_BITMAP_BYTES sets the cap in +// bytes (used by tests to force the multi-chunk path on tiny inputs) and takes +// precedence over GPU_HNSW_BITMAP_MB (megabytes). Default 256 MiB. The value is +// only a chunking bound, so any positive value is safe (the chunk is clamped to +// >= 1 query regardless). +inline size_t gpu_hnsw_bitmap_cap_bytes() { + if (const char* eb = std::getenv("GPU_HNSW_BITMAP_BYTES")) { + long long v = std::atoll(eb); + if (v > 0) { + return static_cast(v); + } + } + size_t mb = 256; + if (const char* e = std::getenv("GPU_HNSW_BITMAP_MB")) { + long v = std::atol(e); + if (v > 0) { + mb = static_cast(v); + } + } + return mb * (static_cast(1) << 20); +} + +// Number of queries to process per launch so the visited bitmap stays within +// gpu_hnsw_bitmap_cap_bytes(). Always in [1, nq]; returns nq when the whole +// batch already fits, so small segments are chunked into a single pass and see +// no behavior change. Queries are independent in HNSW search, so chunking does +// not change per-query results. Must stay in lockstep with the bitmap sizing in +// GpuHnswSearchScratch::ensure() and the launch loop in gpu_hnsw_search(). +inline int gpu_hnsw_bitmap_chunk(int nq, int N) { + if (nq <= 0) { + return nq; + } + size_t per_query = + static_cast((N + 31) / 32) * sizeof(uint32_t); + if (per_query == 0) { + return nq; + } + size_t cap_q = gpu_hnsw_bitmap_cap_bytes() / per_query; + if (cap_q < 1) { + cap_q = 1; + } + return (cap_q >= static_cast(nq)) ? nq + : static_cast(cap_q); +} + struct GpuHnswSearchParams { int ef = 200; int search_width = 4;