Support Standalone RaBitQ in cuVS - #2377
Open
Stardust-SJF wants to merge 4 commits into
Open
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.
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
1 participant
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.