feat: GPU_HNSW index using faiss::gpu::GpuIndexHNSW (slim, feature-only)

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)

cmake/libs/libfaiss.cmake

@@ -531,3 +531,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()

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";

include/knowhere/index/index_table.h

@@ -82,6 +82,8 @@ static std::set<std::pair<std::string, VecType>> 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},

src/index/hnsw/faiss_hnsw.cc

@@ -3283,4 +3283,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 <faiss/gpu/GpuIndexHNSW.h>
+#include <faiss/gpu/StandardGpuResources.h>
+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<faiss::gpu::StandardGpuResources>&
+GetSharedGpuResources() {
+    static std::once_flag flag;
+    static std::shared_ptr<faiss::gpu::StandardGpuResources> instance;
+    std::call_once(flag, [] {
+        instance = std::make_shared<faiss::gpu::StandardGpuResources>();
+        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<BaseConfig>
+    StaticCreateConfig() {
+        return std::make_unique<FaissHnswConfig>();
+    }
+
+    static bool
+    StaticHasRawData(const knowhere::BaseConfig& config, const IndexVersion& version) {
+        return true;
+    }
+
+    static expected<Resource>
+    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<BaseConfig>
+    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<Config> 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<const ::faiss::cppcontrib::knowhere::HasInverseL2Norms*>(faiss_idx) != nullptr;
+                bool use_ip = is_cosine || (faiss_idx->metric_type == ::faiss::METRIC_INNER_PRODUCT);
+
+                {
+                    std::lock_guard<std::mutex> gpu_ctor_lock(GetGpuConstructionMutex());
+                    gpu_resources_ = GetSharedGpuResources();
+                    gpu_index_ = std::make_unique<faiss::gpu::GpuIndexHNSW>(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<DataSetPtr>
+    Search(const DataSetPtr dataset, std::unique_ptr<Config> cfg, const BitsetView& bitset,
+           milvus::OpContext* op_context) const override {
+        if (!bitset.empty() && bitset.count() > 0) {
+            return expected<DataSetPtr>::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<DataSetPtr>::Err(Status::empty_index, "index not loaded");
+                }
+                try {
+                    const auto& hnsw_cfg = static_cast<const FaissHnswConfig&>(*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<std::mutex> gpu_ctor_lock(GetGpuConstructionMutex());
+                        gpu_resources_ = GetSharedGpuResources();
+                        gpu_index_ = std::make_unique<faiss::gpu::GpuIndexHNSW>(gpu_resources_.get(), faiss_idx->d,
+                                                                                faiss_idx->metric_type);
+                    }
+                    gpu_index_->copyFromWithMetric(faiss_idx, use_ip, is_cosine);
+                    const_cast<std::shared_ptr<faiss::Index>&>(indexes[0]).reset();
+                } catch (const std::exception& e) {
+                    return expected<DataSetPtr>::Err(Status::cuvs_inner_error,
+                                                     std::string("failed to build GPU HNSW index: ") + e.what());
+                }
+            }
+        }
+
+        const auto& hnsw_cfg = static_cast<const FaissHnswConfig&>(*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<const float*>(dataset->GetTensor());
+
+        // For COSINE metric, normalize queries to unit length.
+        const float* h_queries = h_queries_raw;
+        std::unique_ptr<float[]> normalized_queries;
+        if (IsMetricType(hnsw_cfg.metric_type.value(), metric::COSINE)) {
+            normalized_queries = std::make_unique<float[]>(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<int64_t[]>(nq * k);
+        auto h_dist = std::make_unique<float[]>(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<DataSetPtr>::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<int64_t>(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<const ::faiss::cppcontrib::knowhere::IndexHNSW*>(indexes[0].get());
+    }
+
+    mutable std::mutex gpu_mutex_;
+    mutable std::shared_ptr<faiss::gpu::StandardGpuResources> gpu_resources_;
+    mutable std::unique_ptr<faiss::gpu::GpuIndexHNSW> gpu_index_;
+};
+
+// Register GPU_HNSW in the static config map at process startup.
+__attribute__((constructor)) static void
+register_gpu_hnsw_static_config() {
+    IndexStaticFaced<fp32>::Instance().RegisterStaticFunc<GpuHnswIndexNode>(IndexEnum::INDEX_GPU_HNSW);
+    IndexStaticFaced<int8>::Instance().RegisterStaticFunc<GpuHnswIndexNode>(IndexEnum::INDEX_GPU_HNSW);
+}
+
+KNOWHERE_REGISTER_GLOBAL(
+    GPU_HNSW,
+    [](const int32_t& version, const Object& object) { return Index<GpuHnswIndexNode>::Create(version, object); }, fp32,
+    true, feature::GPU_ANN_FLOAT_INDEX);
+
+KNOWHERE_REGISTER_GLOBAL(
+    GPU_HNSW,
+    [](const int32_t& version, const Object& object) {
+        return Index<IndexNodeDataMockWrapper<int8>>::Create(std::make_unique<GpuHnswIndexNode>(version, object));
+    },
+    int8, true, (feature::INT8 | feature::GPU));
+#endif  // KNOWHERE_WITH_CUVS
+
 }  // namespace knowhere