Skip to content

Commit b5a3c86

Browse files
tarang-jaintarang-jain
committed
ai slop
1 parent 613e0ac commit b5a3c86

8 files changed

Lines changed: 549 additions & 0 deletions

File tree

‎faiss/gpu/CMakeLists.txt‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -254,22 +254,26 @@ if(FAISS_ENABLE_CUVS)
254254
list(APPEND FAISS_GPU_HEADERS
255255
GpuIndexBinaryCagra.h
256256
GpuIndexCagra.h
257+
GpuIndexIVFRaBitQ.h
257258
impl/BinaryCuvsCagra.cuh
258259
impl/CuvsCagra.cuh
259260
impl/CuvsFlatIndex.cuh
260261
impl/CuvsIVFFlat.cuh
261262
impl/CuvsIVFPQ.cuh
263+
impl/CuvsIVFRaBitQ.cuh
262264
impl/CuvsIVFSQ.cuh
263265
utils/CuvsFilterConvert.h
264266
utils/CuvsUtils.h)
265267
list(APPEND FAISS_GPU_SRC
266268
GpuIndexBinaryCagra.cu
267269
GpuIndexCagra.cu
270+
GpuIndexIVFRaBitQ.cu
268271
impl/BinaryCuvsCagra.cu
269272
impl/CuvsCagra.cu
270273
impl/CuvsFlatIndex.cu
271274
impl/CuvsIVFFlat.cu
272275
impl/CuvsIVFPQ.cu
276+
impl/CuvsIVFRaBitQ.cu
273277
impl/CuvsIVFSQ.cu
274278
utils/CuvsFilterConvert.cu
275279
utils/CuvsUtils.cu)
@@ -306,12 +310,14 @@ if(FAISS_ENABLE_CUVS)
306310
GpuDistance.cu
307311
GpuIndexIVFFlat.cu
308312
GpuIndexIVFPQ.cu
313+
GpuIndexIVFRaBitQ.cu
309314
GpuIndexFlat.cu
310315
StandardGpuResources.cpp
311316
impl/CuvsCagra.cu
312317
impl/CuvsFlatIndex.cu
313318
impl/CuvsIVFFlat.cu
314319
impl/CuvsIVFPQ.cu
320+
impl/CuvsIVFRaBitQ.cu
315321
impl/CuvsIVFSQ.cu
316322
utils/CuvsFilterConvert.cu
317323
utils/CuvsUtils.cu

‎faiss/gpu/GpuIndexIVFRaBitQ.cu‎

Lines changed: 137 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,137 @@
1+
// @lint-ignore-every LICENSELINT
2+
/*
3+
* Copyright (c) Meta Platforms, Inc. and affiliates.
4+
*
5+
* This source code is licensed under the MIT license found in the
6+
* LICENSE file in the root directory of this source tree.
7+
*/
8+
/*
9+
* Copyright (c) 2026, NVIDIA CORPORATION.
10+
*
11+
* Licensed under the Apache License, Version 2.0 (the "License");
12+
* you may not use this file except in compliance with the License.
13+
* You may obtain a copy of the License at
14+
*
15+
* http://www.apache.org/licenses/LICENSE-2.0
16+
*
17+
* Unless required by applicable law or agreed to in writing, software
18+
* distributed under the License is distributed on an "AS IS" BASIS,
19+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
20+
* See the License for the specific language governing permissions and
21+
* limitations under the License.
22+
*/
23+
24+
#include <faiss/gpu/GpuIndexIVFRaBitQ.h>
25+
26+
#include <faiss/gpu/utils/DeviceUtils.h>
27+
#include <faiss/gpu/impl/CuvsIVFRaBitQ.cuh>
28+
29+
#include <algorithm>
30+
#include <limits>
31+
32+
namespace faiss {
33+
namespace gpu {
34+
35+
GpuIndexIVFRaBitQ::GpuIndexIVFRaBitQ(
36+
GpuResourcesProvider* provider,
37+
int dims,
38+
idx_t nlist,
39+
faiss::MetricType metric,
40+
GpuIndexIVFRaBitQConfig config)
41+
: GpuIndex(provider->getResources(), dims, metric, 0.0f, config),
42+
nlist_(nlist),
43+
ivfRabitqConfig_(config) {
44+
FAISS_THROW_IF_NOT_MSG(nlist > 0, "GpuIndexIVFRaBitQ requires nlist > 0");
45+
FAISS_THROW_IF_NOT_MSG(
46+
nlist <= std::numeric_limits<uint32_t>::max(),
47+
"GpuIndexIVFRaBitQ nlist must fit in uint32_t");
48+
FAISS_THROW_IF_NOT_MSG(
49+
config.bitsPerDim >= 1 && config.bitsPerDim <= 9,
50+
"GpuIndexIVFRaBitQ bitsPerDim must be in [1, 9]");
51+
FAISS_THROW_IF_NOT_MSG(
52+
metric == METRIC_L2, "GpuIndexIVFRaBitQ supports METRIC_L2 only");
53+
FAISS_THROW_IF_NOT_MSG(
54+
should_use_cuvs(config),
55+
"GpuIndexIVFRaBitQ requires a supported GPU and "
56+
"GpuIndexIVFRaBitQConfig::use_cuvs = true");
57+
this->is_trained = false;
58+
}
59+
60+
GpuIndexIVFRaBitQ::~GpuIndexIVFRaBitQ() = default;
61+
62+
void GpuIndexIVFRaBitQ::train(idx_t n, const float* x) {
63+
DeviceScope scope(config_.device);
64+
FAISS_THROW_IF_NOT_MSG(
65+
n > 0, "GpuIndexIVFRaBitQ cannot train on an empty dataset");
66+
67+
if (is_trained) {
68+
return;
69+
}
70+
71+
index_ = std::make_shared<CuvsIVFRaBitQ>(
72+
resources_.get(), d, nlist_, metric_type, ivfRabitqConfig_);
73+
index_->train(n, x);
74+
ntotal = n;
75+
is_trained = true;
76+
}
77+
78+
void GpuIndexIVFRaBitQ::add(idx_t n, const float* x) {
79+
FAISS_THROW_IF_MSG(
80+
is_trained,
81+
"GpuIndexIVFRaBitQ does not support incremental additions; "
82+
"call reset() before building a new index");
83+
train(n, x);
84+
}
85+
86+
bool GpuIndexIVFRaBitQ::addImplRequiresIDs_() const {
87+
return false;
88+
}
89+
90+
void GpuIndexIVFRaBitQ::addImpl_(idx_t, const float*, const idx_t*) {
91+
FAISS_THROW_MSG(
92+
"GpuIndexIVFRaBitQ does not support incremental additions; "
93+
"build the index with train() or the first add() call");
94+
}
95+
96+
void GpuIndexIVFRaBitQ::searchImpl_(
97+
idx_t n,
98+
const float* x,
99+
int k,
100+
float* distances,
101+
idx_t* labels,
102+
const SearchParameters* search_params) const {
103+
FAISS_ASSERT(is_trained && index_);
104+
FAISS_THROW_IF_NOT_MSG(k > 0, "GpuIndexIVFRaBitQ requires k > 0");
105+
106+
const auto* params = search_params
107+
? dynamic_cast<const SearchParametersIVFRaBitQ*>(search_params)
108+
: nullptr;
109+
FAISS_THROW_IF_NOT_MSG(
110+
!search_params || params,
111+
"GpuIndexIVFRaBitQ requires SearchParametersIVFRaBitQ");
112+
FAISS_THROW_IF_MSG(
113+
params && params->sel,
114+
"GpuIndexIVFRaBitQ does not support IDSelector filtering");
115+
116+
uint32_t nprobe = params ? params->nprobe : ivfRabitqConfig_.nprobe;
117+
const auto searchMode =
118+
params ? params->searchMode : ivfRabitqConfig_.searchMode;
119+
FAISS_THROW_IF_NOT_MSG(
120+
nprobe > 0, "GpuIndexIVFRaBitQ nprobe must be greater than zero");
121+
nprobe = std::min(nprobe, static_cast<uint32_t>(nlist_));
122+
123+
Tensor<float, 2, true> queries(const_cast<float*>(x), {n, d});
124+
Tensor<float, 2, true> outDistances(distances, {n, k});
125+
Tensor<idx_t, 2, true> outLabels(labels, {n, k});
126+
index_->search(queries, k, outDistances, outLabels, nprobe, searchMode);
127+
}
128+
129+
void GpuIndexIVFRaBitQ::reset() {
130+
DeviceScope scope(config_.device);
131+
index_.reset();
132+
ntotal = 0;
133+
is_trained = false;
134+
}
135+
136+
} // namespace gpu
137+
} // namespace faiss

‎faiss/gpu/GpuIndexIVFRaBitQ.h‎

Lines changed: 122 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,122 @@
1+
// @lint-ignore-every LICENSELINT
2+
/*
3+
* Copyright (c) Meta Platforms, Inc. and affiliates.
4+
*
5+
* This source code is licensed under the MIT license found in the
6+
* LICENSE file in the root directory of this source tree.
7+
*/
8+
/*
9+
* Copyright (c) 2026, NVIDIA CORPORATION.
10+
*
11+
* Licensed under the Apache License, Version 2.0 (the "License");
12+
* you may not use this file except in compliance with the License.
13+
* You may obtain a copy of the License at
14+
*
15+
* http://www.apache.org/licenses/LICENSE-2.0
16+
*
17+
* Unless required by applicable law or agreed to in writing, software
18+
* distributed under the License is distributed on an "AS IS" BASIS,
19+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
20+
* See the License for the specific language governing permissions and
21+
* limitations under the License.
22+
*/
23+
24+
#pragma once
25+
26+
#include <faiss/gpu/GpuIndex.h>
27+
28+
#include <cstddef>
29+
#include <cstdint>
30+
#include <memory>
31+
32+
namespace faiss {
33+
namespace gpu {
34+
35+
class CuvsIVFRaBitQ;
36+
37+
/// Search implementation used by cuVS IVF-RaBitQ.
38+
enum class IVFRaBitQSearchMode {
39+
LUT16 = 0,
40+
LUT32 = 1,
41+
QUANT4 = 2,
42+
QUANT8 = 3,
43+
};
44+
45+
struct GpuIndexIVFRaBitQConfig : public GpuIndexConfig {
46+
/// Number of bits used to encode each residual dimension. Supported values
47+
/// are 1 through 9.
48+
uint32_t bitsPerDim = 3;
49+
50+
/// Number of k-means iterations to use while building the coarse IVF
51+
/// quantizer.
52+
uint32_t kmeansNIterations = 20;
53+
54+
/// Maximum number of training vectors sampled from each coarse cluster.
55+
uint32_t maxTrainPointsPerCluster = 256;
56+
57+
/// Enable cuVS's fast quantization path during index construction.
58+
bool useFastQuantize = true;
59+
60+
/// Maximum number of vectors processed in a host-memory streaming build.
61+
size_t streamingBatchSize = 100000;
62+
63+
/// Force streaming construction when the input is resident on the host.
64+
bool forceStreaming = false;
65+
66+
/// Default number of IVF lists searched for each query.
67+
uint32_t nprobe = 20;
68+
69+
/// Default search implementation.
70+
IVFRaBitQSearchMode searchMode = IVFRaBitQSearchMode::QUANT4;
71+
};
72+
73+
/// Per-search IVF-RaBitQ parameters.
74+
struct SearchParametersIVFRaBitQ : SearchParameters {
75+
uint32_t nprobe = 20;
76+
IVFRaBitQSearchMode searchMode = IVFRaBitQSearchMode::QUANT4;
77+
};
78+
79+
/// GPU-native IVF-RaBitQ index backed by cuVS.
80+
///
81+
/// cuVS builds IVF-RaBitQ from a complete dataset. Consequently, the first
82+
/// call to train() or add() constructs the complete index and incremental
83+
/// additions are not supported.
84+
class GpuIndexIVFRaBitQ : public GpuIndex {
85+
public:
86+
GpuIndexIVFRaBitQ(
87+
GpuResourcesProvider* provider,
88+
int dims,
89+
idx_t nlist,
90+
faiss::MetricType metric = faiss::METRIC_L2,
91+
GpuIndexIVFRaBitQConfig config = GpuIndexIVFRaBitQConfig());
92+
93+
~GpuIndexIVFRaBitQ() override;
94+
95+
void train(idx_t n, const float* x) override;
96+
97+
/// Build the index from x. This is equivalent to train() for this index.
98+
void add(idx_t n, const float* x) override;
99+
100+
void reset() override;
101+
102+
protected:
103+
bool addImplRequiresIDs_() const override;
104+
105+
void addImpl_(idx_t n, const float* x, const idx_t* ids) override;
106+
107+
void searchImpl_(
108+
idx_t n,
109+
const float* x,
110+
int k,
111+
float* distances,
112+
idx_t* labels,
113+
const SearchParameters* search_params) const override;
114+
115+
private:
116+
const idx_t nlist_;
117+
const GpuIndexIVFRaBitQConfig ivfRabitqConfig_;
118+
std::shared_ptr<CuvsIVFRaBitQ> index_;
119+
};
120+
121+
} // namespace gpu
122+
} // namespace faiss

0 commit comments

Comments
 (0)