Repository navigation
Support Standalone RaBitQ in cuVS - #2377
Open
Stardust-SJF wants to merge 4 commits into
Open
Stardust-SJF wants to merge 4 commits into
Stardust-SJF wants to merge 4 commits into
Conversation
Imports the standalone RaBitQ GPU quantizer as a first-class preprocessing
component, alongside the existing scalar, binary and pq quantizers. It is
independent of cuvs::neighbors::ivf_rabitq and shares no code with it.
Public API (cuvs::preprocessing::quantize::rabitq, header-installed with the
rest of include/cuvs):
rotator_kind { matmul, fht_kac }, delta_kind, params
rotator - kind + padded_dim + the rotation state (raft containers)
quantizer - params + const_scaling_factor + dim/padded_dim
pipeline - composes the two; neither owns the other
init_rotator / init_quantizer / init_pipeline
rotate, transform, transform_residuals
The rotator and the quantizer are independent peer types, mirroring how
IVFGPU already holds RotatorGPU and DataQuantizerGPU as siblings; `pipeline`
is the pq-like convenience that composes them. The factories take a `dim`
rather than a dataset because RaBitQ needs no training: they generate the
random rotation and derive the scaling constant from (dim, ex_bits, seed).
State is plain caller-owned data in raft containers, following
binary::quantizer, so the rotator travels with the codes and no serialize API
is required - matching every other preprocessing quantizer, none of which has
I/O (persistence lives at the index layer, cf. detail/dataset_serialize.hpp).
Import notes:
- The imported sources land in detail/rabitq/ essentially verbatim, wrapped in
...::rabitq::detail. Include paths, the enclosing namespace and the error
macros are the only edits; kernels and numerics are untouched.
- CUDA_CHECK/CUBLAS_CHECK, which printed and called exit(), are now
RAFT_CUDA_TRY/RAFT_CUBLAS_TRY so failures throw instead of killing the
process.
- class RotatorGPU is dissolved: it owned raw cudaMalloc buffers and its own
cublasHandle_t. Its state now lives in the public rotator, and detail exposes
free launchers over caller-owned buffers, taking the cuBLAS handle and the
stream from raft::resources.
- The pipeline's scratch buffer comes from the workspace resource instead of
cudaMalloc, and its glue kernels run on the resource stream.
- The fht_kac segment size is derived from the padded dimension, matching the
standalone.
Known follow-up: quantize.cu and rescale_search.cu still launch on the default
stream. Since cuVS builds with CUDA_API_PER_THREAD_DEFAULT_STREAM that does not
implicitly sync with the resource stream, so pipeline.cu bridges them with two
TODO(rabitq)-marked host syncs. Threading a stream through those two files
removes both.
Tested on an RTX PRO 6000 (sm_120a): PREPROCESSING_TEST passes 368/368,
including 141 new RaBitQ cases covering both rotator kinds, uint8/uint16 codes,
padded/pow2/non-pow2 dims, rotation norm preservation, non-degenerate codes,
finite factors and run-to-run determinism.
The imported quantizer entry points launched on the per-thread default stream,
which under CUDA_API_PER_THREAD_DEFAULT_STREAM does not implicitly synchronize
with the resource stream. That was bridged by four host-side syncs, so every
transform() paid two full host<->device round trips and serialized the caller's
stream. Thread the stream through instead and delete the bridge.
- quantize_{fused,full}_on_residuals and the launch_quantize_* helpers now take
cudaStream_t as their first parameter (first, so the existing trailing default
arguments still work, and without a default of its own so no caller can fall
back to the default stream by accident). Both kernel launches use it.
- get_const_scaling_factor takes raft::resources const& first: it needs the
stream, the workspace memory resource and the device properties.
- Its three cudaMalloc/cudaFree pairs become rmm::device_uvector on the
workspace resource, so the scratch comes from the caller's memory resource and
is freed on the stream rather than at a device-wide synchronization point.
- cudaGetDeviceProperties(&prop, 0) becomes
raft::resource::get_device_properties(res). Reading device 0's
sharedMemPerBlock to pick block_size was wrong whenever the current device was
not device 0, and it was the one CUDA call in the file with no RAFT_CUDA_TRY;
the resource accessor is also cached, so this drops a driver query.
- Both cub::DeviceReduce::Sum calls take the stream and are wrapped in
RAFT_CUDA_TRY.
- sync_before_quantize / sync_after_quantize are removed from both pipeline.cu
and rabitq.cu along with all four call sites.
The transform() path now issues no host synchronization at all. The two
remaining syncs are init-only and necessary: the async read-back of the scaling
factor in get_const_scaling_factor (it returns a host float) and the state
uploads in init_state_matmul / init_state_fht_kac (they copy from host vectors
that die at scope exit).
No kernel body, no numerics and no rotator behaviour changed. PREPROCESSING_TEST
still passes 368/368 on an RTX PRO 6000 (sm_120a), unchanged from before.
Aligns the public factories with the verb cuVS already uses for functions that return a new object, and with the split ivf_rabitq itself observes: - make_* returns a new object. ivf_rabitq's internals use it for exactly that (make_compute_inner_products_with_lut_launcher and friends), and cuVS's public headers use it ~106 times across make_const_mdspan, make_scalar_view, make_list_extents, make_aligned_dataset, make_strided_dataset and others. - init_* populates existing state and returns void, e.g. void IVFGPU::init_clusters(cluster_sizes). Every other init_* in the tree is a detail-level seeder (init_centroids, init_plus_plus, init_random_graph, init_key), never public API. Before this change init_rotator / init_quantizer / init_pipeline were the only init_* functions in the whole public API surface, and they return objects, so make_* is both the conventional and the more accurate verb. The detail layer keeps init_state_matmul / init_state_fht_kac: they fill a caller-provided device buffer and return void, which is the init_* semantic. PREPROCESSING_TEST still passes 368/368; no behaviour change.
Brings rabitq to parity with scalar/binary/pq, which all expose a C API under
c/. 20 functions, each returning cuvsError_t with its body wrapped in
cuvs::core::translate_exceptions:
- cuvsRabitqQuantizerParams + ParamsCreate/ParamsDestroy, defaults mirroring the
C++ member initializers.
- Opaque cuvsRabitqRotator / cuvsRabitqQuantizer handles + Create/Destroy.
- cuvsRabitqRotatorMake / QuantizerMake / PipelineMake.
- cuvsRabitqRotate, cuvsRabitqTransform[Full], cuvsRabitqTransformResiduals[Full].
- Getters: rotator kind/padded_dim/state, quantizer dim/padded_dim/
const_scaling_factor.
Notes on the shape:
- No opaque pipeline handle. The C++ pipeline only composes the two peers, so
cuvsRabitqPipelineMake takes two out-parameters instead of adding a third
handle type to unwrap.
- The handles carry only { uintptr_t addr }. pq and binary also carry a
DLDataType because their C++ class is quantizer<T> and the field records which
instantiation to cast back to; rabitq's rotator and quantizer are
non-templated, so the field would be permanently dead.
- The rotator getters are load-bearing rather than conveniences: there is no
serialize API by design, so GetKind/GetPaddedDim/GetState are the only way a C
caller can keep the rotator alongside the codes it produced. GetState follows
cuvsProductQuantizerGetPqCodebook's ownership semantics (a non-owning view of
rotator-owned device memory) but its rank and dtype depend on the kind: 2-D
float32 for matmul, 1-D uint8 for fht_kac.
- Device memory only. pq/scalar/binary dispatch between host and device
overloads; every rabitq entry point takes device views, so host tensors are
rejected rather than routed.
- transform and transform_residuals are each overloaded on their output type
alone, which C cannot express, hence the explicit ...Full spellings.
- rabitq.h had to be added to c/include/cuvs/core/all.h: header_check.cmake
globs c/include and fails an existing test for any header missing from it.
Tested on an RTX PRO 6000 (sm_120a): RABITQ_C_TEST passes 2/2, covering both
rotator kinds end to end (PipelineMake -> Transform -> GetState -> Destroy), and
PREPROCESSING_TEST still passes 368/368.
cjnolet
added
improvement
This was referenced Oct 10, 2026
tfeher
requested changes
Oct 19, 2026
tfeher
left a comment
tfeher
left a comment
Contributor
There was a problem hiding this comment.
Thank you Jifan for the PR! We are excited to have a standalone RaBitQ quantizer in cuVS. The code is well structured and cleanly implemented. In general, we should aim for more code reuse. Here is my first batch of comments.
| * @brief RaBitQ quantizer parameters. | ||
| * | ||
| * These parameters fully describe the quantizer: RaBitQ is data-oblivious, so | ||
| * none of them are learned from a dataset (see make_quantizer / make_rotator). |
Contributor
There was a problem hiding this comment.
This means there is no clustering involved. We also don't subtract a global mean value from the vectors, right? Do we have assumption about data distribution (e.g. zero mean), or it is not required?
Comment on lines
+168
to
+172
| struct pipeline { | ||
| /** Rotation applied before quantization. */ | ||
| rabitq::rotator rotator; | ||
| /** Quantizer applied to the rotated residuals. */ | ||
| rabitq::quantizer quantizer; |
Contributor
There was a problem hiding this comment.
Here pipeline is the high level object, the entry point that the user should use for RaBitQ quantizing data. In itself it is logical. But it does not match the naming conventions used in the other quantizers we have. E.g. PQ quantizer is actually a pipeline of multiple steps, still the high level object is called quantizer. I recommend to rename pipeline as "quantizer" and find a different name for the existing quantizer. Tagging @achirkin for suggestions.
Comment on lines
+283
to
+285
| * A temporary `n_rows x padded_dim` float buffer holding the rotated residuals | ||
| * is allocated from the workspace resource, so large datasets should be | ||
| * transformed in batches (a pool workspace resource is recommended). |
Contributor
There was a problem hiding this comment.
Most of cuVS algorithms would do this batching under the hood automatically. I think it would make sense to implement it here too.
Contributor
There was a problem hiding this comment.
The algorithms currently only accept device input data. Large host dataset would need manual batching and copying. We should consider accepting host input data as well.
| // Pipeline: pad(dim -> rotator.padded_dim) -> rotate -> [subtract rotated | ||
| // centroid] -> fused quantization on the rotated residuals. | ||
| // | ||
| // Design note: the centroid is rotated with a separate 1-row rotate() call |
Contributor
There was a problem hiding this comment.
Won't this make kernel launch overhead a problem?
| // conventions. Numerically stable enough for the few-thousand dimensions used | ||
| // by RaBitQ. Output is row-major. | ||
| // --------------------------------------------------------------------------- | ||
| std::vector<float> random_orthogonal_matrix(int64_t D, uint64_t seed) |
Contributor
There was a problem hiding this comment.
I believe ivf_pq::hepers::make_rotation_matrix could do the same task. We should use that instead. (Ideally we should move the helper to preprocessing, and use it both from ivf_pq and ivf_rabitq).
| auto stream = raft::resource::get_cuda_stream(res); | ||
| constexpr int kBlock = 256; | ||
| const size_t total = N * D; | ||
| const int grid = static_cast<int>((total + kBlock - 1) / kBlock); |
Contributor
There was a problem hiding this comment.
Suggested change
| const int grid = static_cast<int>((total + kBlock - 1) / kBlock); | |
| const int grid = raft::ceildiv<int>(total, kBlock); |
| pad_and_rotate(res, d_centroid, 1, dim, rotator, d_rotated_centroid); | ||
| constexpr int kBlock = 256; | ||
| const size_t total = N * D; | ||
| subtract_row_kernel<<<static_cast<int>((total + kBlock - 1) / kBlock), kBlock, 0, stream>>>( |
Contributor
There was a problem hiding this comment.
This can be expressed with raft::linewise_op:
auto rows = raft::make_device_matrix_view<float, int64_t>(d_residuals, N, D);
raft::matrix::linewise_op<raft::Apply::ALONG_ROWS>(
res, raft::make_const_mdspan(rows), rows, raft::sub_op{},
raft::make_device_vector_view<const float, int64_t>(d_rotated_centroid, dim));|
|
||
| namespace { | ||
|
|
||
| __global__ void pad_rows_kernel(const float* __restrict__ d_src, |
Contributor
There was a problem hiding this comment.
This can be expressed with raft::linalg::map:
void pad_rows(raft::resources const& res,
const float* d_src, float* d_dst,
size_t N, size_t dim, size_t D) {
raft::linalg::map_offset(
res,
raft::make_device_vector_view<float, size_t>(d_dst, N * D),
[d_src, dim, D] __device__(size_t idx) {
const size_t row = idx / D;
const size_t col = idx % D;
return col < dim ? d_src[row * dim + col] : 0.0f;
});
}
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This PR introduces the standalone RaBitQ quantization components into cuVS, which provide a convenient way to quantize vectors into RaBitQ codes.
Note that it depends on #2362. Only the commits on top of
cuvs_ivf_rabitq_syn_fix_fusedare intended to be reviewed in this PR.