Skip to content

Commit 8e24da1

Browse files
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 <devin-ai-integration[bot]@users.noreply.github.com>
1 parent 7964964 commit 8e24da1

14 files changed

Lines changed: 2499 additions & 2 deletions

File tree

cmake/libs/libdiskann.cmake

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
add_definitions(-DKNOWHERE_WITH_DISKANN)
2-
find_package(Boost REQUIRED COMPONENTS program_options)
3-
include_directories(${Boost_INCLUDE_DIR})
2+
find_package(Boost REQUIRED COMPONENTS program_options CONFIG)
3+
include_directories(${Boost_INCLUDE_DIRS})
44
find_package(aio REQUIRED)
55
include_directories(${AIO_INCLUDE})
66
find_package(fmt REQUIRED)

cmake/libs/libfaiss.cmake

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -593,3 +593,26 @@ if(__PPC64)
593593
knowhere_utils)
594594
target_compile_definitions(faiss PRIVATE FINTEGER=int)
595595
endif()
596+
597+
# GPU HNSW CUDA sources — compiled when WITH_CUVS is enabled
598+
if(WITH_CUVS)
599+
set(FAISS_GPU_HNSW_SRCS
600+
thirdparty/faiss/faiss/gpu/GpuIndexHNSW.cu
601+
thirdparty/faiss/faiss/gpu/GpuIndex.cu
602+
thirdparty/faiss/faiss/gpu/GpuResources.cpp
603+
thirdparty/faiss/faiss/gpu/StandardGpuResources.cpp
604+
thirdparty/faiss/faiss/gpu/impl/GpuHnswTypes.cu
605+
thirdparty/faiss/faiss/gpu/impl/IndexUtils.cu
606+
thirdparty/faiss/faiss/gpu/utils/DeviceUtils.cu
607+
thirdparty/faiss/faiss/gpu/utils/StackDeviceMemory.cpp
608+
thirdparty/faiss/faiss/gpu/utils/Timer.cpp
609+
)
610+
add_library(faiss_gpu_hnsw OBJECT ${FAISS_GPU_HNSW_SRCS})
611+
target_include_directories(faiss_gpu_hnsw PRIVATE
612+
${CMAKE_CURRENT_SOURCE_DIR}/thirdparty/faiss
613+
${Boost_INCLUDE_DIRS}
614+
)
615+
target_compile_definitions(faiss_gpu_hnsw PRIVATE FINTEGER=int)
616+
target_link_libraries(faiss_gpu_hnsw PRIVATE CUDA::cudart)
617+
target_link_libraries(faiss PUBLIC faiss_gpu_hnsw)
618+
endif()

include/knowhere/comp/index_param.h

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,7 @@ constexpr const char* INDEX_GPU_BRUTEFORCE = "GPU_BRUTE_FORCE";
5353
constexpr const char* INDEX_GPU_IVFFLAT = "GPU_IVF_FLAT";
5454
constexpr const char* INDEX_GPU_IVFPQ = "GPU_IVF_PQ";
5555
constexpr const char* INDEX_GPU_CAGRA = "GPU_CAGRA";
56+
constexpr const char* INDEX_GPU_HNSW = "GPU_HNSW";
5657

5758
constexpr const char* INDEX_HNSW = "HNSW";
5859
constexpr const char* INDEX_HNSW_SQ = "HNSW_SQ";

include/knowhere/index/index_table.h

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -82,6 +82,8 @@ static std::set<std::pair<std::string, VecType>> legal_knowhere_index = {
8282
{IndexEnum::INDEX_GPU_CAGRA, VecType::VECTOR_FLOAT16},
8383
{IndexEnum::INDEX_GPU_CAGRA, VecType::VECTOR_INT8},
8484
{IndexEnum::INDEX_GPU_CAGRA, VecType::VECTOR_BINARY},
85+
{IndexEnum::INDEX_GPU_HNSW, VecType::VECTOR_FLOAT},
86+
{IndexEnum::INDEX_GPU_HNSW, VecType::VECTOR_INT8},
8587

8688
// hnsw
8789
{IndexEnum::INDEX_HNSW, VecType::VECTOR_FLOAT},

src/index/hnsw/faiss_hnsw.cc

Lines changed: 244 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3286,4 +3286,248 @@ KNOWHERE_SIMPLE_REGISTER_DENSE_FLOAT_ALL_GLOBAL(HNSW_PRQ, BaseFaissRegularIndexH
32863286
KNOWHERE_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

Comments
 (0)