@@ -3286,4 +3286,248 @@ KNOWHERE_SIMPLE_REGISTER_DENSE_FLOAT_ALL_GLOBAL(HNSW_PRQ, BaseFaissRegularIndexH
32863286KNOWHERE_SIMPLE_REGISTER_DENSE_INT_GLOBAL (HNSW_PRQ , BaseFaissRegularIndexHNSWPRQNodeTemplate,
32873287 knowhere::feature::MMAP | knowhere::feature::MV | knowhere::feature::EMB_LIST )
32883288
3289+ #ifdef KNOWHERE_WITH_CUVS
3290+ } // namespace knowhere — temporarily close to include GPU headers at file scope
3291+ // ── GPU HNSW ─────────────────────────────────────────────────────────────────
3292+ #include < faiss/gpu/GpuIndexHNSW.h>
3293+ #include < faiss/gpu/StandardGpuResources.h>
3294+ namespace knowhere { // reopen namespace knowhere
3295+
3296+ // Process-global StandardGpuResources shared by all GpuHnswIndexNode instances.
3297+ // Avoids per-segment 256 MiB pinned memory, cuBLAS handle, and CUDA stream
3298+ // allocations that accumulate to tens of GiB with many segments.
3299+ static std::shared_ptr<faiss::gpu::StandardGpuResources>&
3300+ GetSharedGpuResources () {
3301+ static std::once_flag flag;
3302+ static std::shared_ptr<faiss::gpu::StandardGpuResources> instance;
3303+ std::call_once (flag, [] {
3304+ instance = std::make_shared<faiss::gpu::StandardGpuResources>();
3305+ instance->setTempMemory (0 );
3306+ instance->setPinnedMemory (0 );
3307+ });
3308+ return instance;
3309+ }
3310+
3311+ // Serialize GpuIndexHNSW construction across segments.
3312+ // StandardGpuResourcesImpl::initializeForDevice is not thread-safe;
3313+ // concurrent constructors race on the allocs_ map assertion.
3314+ static std::mutex&
3315+ GetGpuConstructionMutex () {
3316+ static std::mutex mtx;
3317+ return mtx;
3318+ }
3319+
3320+ // Single GPU HNSW index node that handles all CPU storage formats (F32, SQ8,
3321+ // FP16, BF16) transparently. Uses faiss::gpu::GpuIndexHNSW for GPU search.
3322+ // Accepts CPU-serialized HNSW or HNSW_SQ binaries at load time.
3323+ class GpuHnswIndexNode : public BaseFaissRegularIndexHNSWNode {
3324+ public:
3325+ GpuHnswIndexNode (const int32_t & version, const Object& object)
3326+ : BaseFaissRegularIndexHNSWNode(version, object, DataFormatEnum::fp32) {
3327+ }
3328+
3329+ static std::unique_ptr<BaseConfig>
3330+ StaticCreateConfig () {
3331+ return std::make_unique<FaissHnswConfig>();
3332+ }
3333+
3334+ static bool
3335+ StaticHasRawData (const knowhere::BaseConfig& config, const IndexVersion& version) {
3336+ return true ;
3337+ }
3338+
3339+ static expected<Resource>
3340+ StaticEstimateLoadResource (const uint64_t file_size_in_bytes, const int64_t num_rows, const int64_t dim,
3341+ const knowhere::BaseConfig& config, const IndexVersion& version) {
3342+ // GPU HNSW stores vectors and graph in VRAM; the CPU copy is freed
3343+ // after upload in Deserialize(). Report zero CPU memory cost so the
3344+ // Milvus segment loader does not over-commit host RAM reservations.
3345+ return Resource{.memoryCost = 0 , .diskCost = 0 };
3346+ }
3347+
3348+ std::unique_ptr<BaseConfig>
3349+ CreateConfig () const override {
3350+ return StaticCreateConfig ();
3351+ }
3352+
3353+ std::string
3354+ Type () const override {
3355+ return IndexEnum::INDEX_GPU_HNSW ;
3356+ }
3357+
3358+ protected:
3359+ Status
3360+ TrainInternal (const DataSetPtr /* dataset*/ , const Config& /* cfg*/ ) override {
3361+ return Status::not_implemented;
3362+ }
3363+
3364+ public:
3365+ Status
3366+ Deserialize (const BinarySet& binset, std::shared_ptr<Config> cfg) override {
3367+ std::unique_lock lock (gpu_mutex_);
3368+ gpu_index_.reset ();
3369+
3370+ // Accept CPU-built HNSW (F32) or HNSW_SQ (quantized) binaries.
3371+ Status status;
3372+ if (!binset.Contains (IndexEnum::INDEX_GPU_HNSW )) {
3373+ BinarySet aliased = binset;
3374+ for (const char * key : {IndexEnum::INDEX_HNSW_SQ , IndexEnum::INDEX_HNSW }) {
3375+ if (binset.Contains (key)) {
3376+ aliased.Append (IndexEnum::INDEX_GPU_HNSW , binset.GetByName (key));
3377+ break ;
3378+ }
3379+ }
3380+ status = BaseFaissRegularIndexHNSWNode::Deserialize (aliased, cfg);
3381+ } else {
3382+ status = BaseFaissRegularIndexHNSWNode::Deserialize (binset, cfg);
3383+ }
3384+ if (status != Status::success) {
3385+ return status;
3386+ }
3387+
3388+ // Eager GPU upload via faiss::gpu::GpuIndexHNSW.
3389+ const auto * faiss_idx = GetFaissHnswIndex ();
3390+ if (faiss_idx) {
3391+ try {
3392+ // Detect metric from the FAISS index type rather than config,
3393+ // because Deserialize may be called without metric_type in the config
3394+ // (e.g. empty json defaults metric_type to L2).
3395+ bool is_cosine =
3396+ dynamic_cast <const ::faiss::cppcontrib::knowhere::HasInverseL2Norms*>(faiss_idx) != nullptr ;
3397+ bool use_ip = is_cosine || (faiss_idx->metric_type == ::faiss::METRIC_INNER_PRODUCT );
3398+
3399+ {
3400+ std::lock_guard<std::mutex> gpu_ctor_lock (GetGpuConstructionMutex ());
3401+ gpu_resources_ = GetSharedGpuResources ();
3402+ gpu_index_ = std::make_unique<faiss::gpu::GpuIndexHNSW>(gpu_resources_.get (), faiss_idx->d ,
3403+ faiss_idx->metric_type );
3404+ }
3405+ gpu_index_->copyFromWithMetric (faiss_idx, use_ip, is_cosine);
3406+ // Release CPU copy — vectors and graph are now on GPU.
3407+ indexes[0 ].reset ();
3408+ } catch (const std::exception& e) {
3409+ fprintf (stderr, " [gpu_hnsw] eager GPU upload failed: %s\n " , e.what ());
3410+ gpu_index_.reset ();
3411+ }
3412+ }
3413+ return Status::success;
3414+ }
3415+
3416+ expected<DataSetPtr>
3417+ Search (const DataSetPtr dataset, std::unique_ptr<Config> cfg, const BitsetView& bitset,
3418+ milvus::OpContext* op_context) const override {
3419+ if (!bitset.empty () && bitset.count () > 0 ) {
3420+ return expected<DataSetPtr>::Err (Status::invalid_args, " GPU_HNSW does not support filtered search" );
3421+ }
3422+
3423+ // Fast path: gpu_index_ is set during Deserialize and never cleared.
3424+ if (!gpu_index_) {
3425+ std::unique_lock lock (gpu_mutex_);
3426+ if (!gpu_index_) {
3427+ const auto * faiss_idx = GetFaissHnswIndex ();
3428+ if (!faiss_idx) {
3429+ return expected<DataSetPtr>::Err (Status::empty_index, " index not loaded" );
3430+ }
3431+ try {
3432+ const auto & hnsw_cfg = static_cast <const FaissHnswConfig&>(*cfg);
3433+ bool is_cosine = IsMetricType (hnsw_cfg.metric_type .value (), metric::COSINE );
3434+ bool use_ip = IsMetricType (hnsw_cfg.metric_type .value (), metric::IP ) || is_cosine;
3435+
3436+ {
3437+ std::lock_guard<std::mutex> gpu_ctor_lock (GetGpuConstructionMutex ());
3438+ gpu_resources_ = GetSharedGpuResources ();
3439+ gpu_index_ = std::make_unique<faiss::gpu::GpuIndexHNSW>(gpu_resources_.get (), faiss_idx->d ,
3440+ faiss_idx->metric_type );
3441+ }
3442+ gpu_index_->copyFromWithMetric (faiss_idx, use_ip, is_cosine);
3443+ const_cast <std::shared_ptr<faiss::Index>&>(indexes[0 ]).reset ();
3444+ } catch (const std::exception& e) {
3445+ return expected<DataSetPtr>::Err (Status::cuvs_inner_error,
3446+ std::string (" failed to build GPU HNSW index: " ) + e.what ());
3447+ }
3448+ }
3449+ }
3450+
3451+ const auto & hnsw_cfg = static_cast <const FaissHnswConfig&>(*cfg);
3452+ auto k = hnsw_cfg.k .value ();
3453+ auto nq = dataset->GetRows ();
3454+ auto dim = dataset->GetDim ();
3455+ auto ef = hnsw_cfg.ef .value_or (200 );
3456+ const auto * h_queries_raw = reinterpret_cast <const float *>(dataset->GetTensor ());
3457+
3458+ // For COSINE metric, normalize queries to unit length.
3459+ const float * h_queries = h_queries_raw;
3460+ std::unique_ptr<float []> normalized_queries;
3461+ if (IsMetricType (hnsw_cfg.metric_type .value (), metric::COSINE )) {
3462+ normalized_queries = std::make_unique<float []>(nq * dim);
3463+ for (int64_t i = 0 ; i < nq; i++) {
3464+ const float * src = h_queries_raw + i * dim;
3465+ float * dst = normalized_queries.get () + i * dim;
3466+ float sq_norm = 0 .0f ;
3467+ for (int64_t d = 0 ; d < dim; d++) sq_norm += src[d] * src[d];
3468+ float inv = (sq_norm > 0 .0f ) ? (1 .0f / std::sqrt (sq_norm)) : 1 .0f ;
3469+ for (int64_t d = 0 ; d < dim; d++) dst[d] = src[d] * inv;
3470+ }
3471+ h_queries = normalized_queries.get ();
3472+ }
3473+
3474+ auto h_ids = std::make_unique<int64_t []>(nq * k);
3475+ auto h_dist = std::make_unique<float []>(nq * k);
3476+
3477+ try {
3478+ faiss::gpu::GpuHnswSearchParams gsp;
3479+ gsp.ef = ef;
3480+ gpu_index_->searchHost (nq, h_queries, k, h_dist.get (), h_ids.get (), gsp);
3481+ } catch (const std::exception& e) {
3482+ LOG_KNOWHERE_ERROR_ << " GPU_HNSW search failed: " << e.what ();
3483+ return expected<DataSetPtr>::Err (Status::cuvs_inner_error,
3484+ std::string (" GPU HNSW search failed: " ) + e.what ());
3485+ }
3486+
3487+ // Negate back to positive for IP and COSINE.
3488+ if (IsMetricType (hnsw_cfg.metric_type .value (), metric::IP ) ||
3489+ IsMetricType (hnsw_cfg.metric_type .value (), metric::COSINE )) {
3490+ for (int64_t i = 0 ; i < static_cast <int64_t >(nq * k); i++) {
3491+ h_dist[i] = -h_dist[i];
3492+ }
3493+ }
3494+
3495+ return GenResultDataSet (nq, k, h_ids.release (), h_dist.release ());
3496+ }
3497+
3498+ ~GpuHnswIndexNode () override = default ;
3499+
3500+ private:
3501+ const ::faiss::cppcontrib::knowhere::IndexHNSW*
3502+ GetFaissHnswIndex () const {
3503+ if (indexes.empty () || !indexes[0 ])
3504+ return nullptr ;
3505+ return dynamic_cast <const ::faiss::cppcontrib::knowhere::IndexHNSW*>(indexes[0 ].get ());
3506+ }
3507+
3508+ mutable std::mutex gpu_mutex_;
3509+ mutable std::shared_ptr<faiss::gpu::StandardGpuResources> gpu_resources_;
3510+ mutable std::unique_ptr<faiss::gpu::GpuIndexHNSW> gpu_index_;
3511+ };
3512+
3513+ // Register GPU_HNSW in the static config map at process startup.
3514+ __attribute__ ((constructor)) static void
3515+ register_gpu_hnsw_static_config () {
3516+ IndexStaticFaced<fp32>::Instance ().RegisterStaticFunc <GpuHnswIndexNode>(IndexEnum::INDEX_GPU_HNSW );
3517+ IndexStaticFaced<int8>::Instance ().RegisterStaticFunc <GpuHnswIndexNode>(IndexEnum::INDEX_GPU_HNSW );
3518+ }
3519+
3520+ KNOWHERE_REGISTER_GLOBAL (
3521+ GPU_HNSW ,
3522+ [](const int32_t & version, const Object& object) { return Index<GpuHnswIndexNode>::Create (version, object); }, fp32,
3523+ true , feature::GPU_ANN_FLOAT_INDEX );
3524+
3525+ KNOWHERE_REGISTER_GLOBAL (
3526+ GPU_HNSW ,
3527+ [](const int32_t & version, const Object& object) {
3528+ return Index<IndexNodeDataMockWrapper<int8>>::Create (std::make_unique<GpuHnswIndexNode>(version, object));
3529+ },
3530+ int8, true , (feature::INT8 | feature::GPU ));
3531+ #endif // KNOWHERE_WITH_CUVS
3532+
32893533} // namespace knowhere
0 commit comments