Fast CAGRA Index Merge - #2352
Conversation
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
|
/ok to test 2c54392 |
aamijar
added
non-breaking
|
/ok to test |
|
/ok to test |
|
/ok to test |
|
/ok to test |
Brings the branch up to date with main (f72199e), which resolves the check-c-abi failure: the checker was reporting cuvsResourcesSetWorkspacePool, cuvsRMMAsyncMemoryResourceEnable, cuvsCagraSearchMultiPartition and cuvsSelectK as removed functions, when in fact they were added upstream after this branch's base and were simply missing here. Only cpp/src/neighbors/cagra.cuh conflicted, and only positionally: both sides append a new function immediately after the existing merge() overload. Kept both -- this branch's merge() taking merge_params, and main's multi-partition search() overload. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
|
Currently failing because the wheel builds are too big. Not sure there's much I can do about this. |
achirkin
left a comment
achirkin
left a comment
There was a problem hiding this comment.
Thanks for this extensive contribution! Since the datataset API changes are in, the main request is to rethink where we allocate the new, consolidated dataset during merge (see below).
Then, please try to refactor the code to use raft mdarrays and mdspans in place of RMM and use the workspace/large_workspace resources for all temporary allocations - this will help us to keep track of memory requirements (which seem to be plenty).
| /** Device state and tuning knobs shared by every split level: the precomputed row norms, the | ||
| * tiling and workspace capacities, and the seed feeding the deterministic leader samples. */ | ||
| struct split_context { | ||
| split_context(rmm::cuda_stream_view stream, int64_t rows) | ||
| : norms(static_cast<size_t>(rows), stream) | ||
| { | ||
| } | ||
|
|
||
| rmm::device_uvector<float> norms; | ||
| int assignment_tile_rows = ASSIGNMENT_TILE_ROWS; | ||
| size_t gemm_workspace_bytes = GEMM_WORKSPACE_BYTES; | ||
| uint64_t seed = DETERMINISTIC_SEED; | ||
| }; |
There was a problem hiding this comment.
Let's try to switch to raft primitives here: raft::device_vector<float> as the buffer and raft::resources as an argument.
Then, we need to consider: does this context fit into the bounded workspace memory or does it grow with the problem size unbounded? Depending on the answer, use raft::get_workspace_resource_ref or raft::get_large_workspace_resource_ref.
By doing this, we let the user control where the working memory for the algorithm resides, enable better accounting (via tools like raft::memory_tracking_resources), and prepare for the dry run execution (feature soon to be enabled in raft).
There was a problem hiding this comment.
I've gone through and I think this is done correctly now, let me know if I misunderstand the usage of the raft allocs. Is there a way I should be testing preparedness for dry run execution?
| if (params.attach_dataset_on_build) { | ||
| // CAGRA search assumes 16-byte row alignment (vectorized loads). Consolidate remains dense for | ||
| // the scaffold; attach through make_aligned_dataset so unaligned dims are padded. | ||
| merged_index.update_dataset(handle, make_aligned_dataset(handle, std::move(dataset), 16)); | ||
| } else { | ||
| using ds_idx_type = typename index<T, IdxT>::dataset_index_type; | ||
| merged_index.update_dataset( | ||
| handle, std::make_unique<cuvs::neighbors::empty_dataset<ds_idx_type>>(preflight.dim)); | ||
| } |
There was a problem hiding this comment.
I guess, we should move the dataset consolidation into a separate place, perhaps as a part of the dataset API. CC @HowardHuang1 to discuss what would be the place for dataset merging under the new API and who should own it?
There was a problem hiding this comment.
I can make an issue for this, but could we leave it here for now?
There was a problem hiding this comment.
Adding offline context to the thread:
The dataset merging logic now present is just copying vectors into storage that's now provided by the caller per interface changes that were introduced with the dataset API. The rebuild version of this takes a bitmap to filter them, and the fastener version doesn't.
Co-authored-by: Artem M. Chirkin <9253178+achirkin@users.noreply.github.com>
Co-authored-by: Artem M. Chirkin <9253178+achirkin@users.noreply.github.com>
1b8558a to
653e8a0
Compare
The graph merge PR (#2352) uses some of the kernels defined by CAGRA, and binary size takes a hit from them not appearing in a TU where they can be shared by both CAGRA and Fastener. Authors: - Ben Landrum (https://github.com/landrumb) Approvers: - Divye Gala (https://github.com/divyegala) URL: #2412
4 participants
This PR implements the Fastener graph merge operation. This PR supports merging for
float,half,int8, anduint8dtypes, and euclidean distances.There was originally specialization to use int8 GEMM for the integer types, but I ran into portability issues on Ada and it turns out that using the same f32 path for both is simpler and not substantially slower. Currently investigating switching the unified path to TF32 to use tensor cores.
The core logic resides in
cagra_merge_scaffold.cuh.This PR adds 7.64 MiB to
libcuvs.so, a 2.94% increase.