[WIP] Vpq dataset serialisation - #2480
Conversation
| // cuvsDatasetLayout_t covers standard and padded only, and every C entry point dispatches | ||
| // on that layout, so there is nothing here to hand a VPQ index to yet. | ||
| RAFT_FAIL("File holds a VPQ-compressed (CAGRA-Q) dataset, which the C API has no dataset " | ||
| "layout for; load it through the C++ API"); |
There was a problem hiding this comment.
This is not an acceptable way to handle things tha are not yet exposed to C. Folks use C because they have to (for C ABI stability guarantees, for example). We can't just throw an error and forward them to C++
| /** Host-resident dataset using its standard row layout. */ | ||
| host_standard = 4, | ||
| /** Device-resident VPQ-compressed dataset with f16 codebooks (CAGRA-Q). */ | ||
| device_vpq_f16 = 5, |
There was a problem hiding this comment.
Since we're adding new things, please rename to device_pq. VPQ is not a proper term it accidentally made its way into the public APIs (which was an oversight). Also- why the "f16"? Can we drop that?
| * error rather than a silently unusable index. For the same reason `include_dataset = false` | ||
| * produces an index that only `update_dataset` can make searchable again. | ||
| */ | ||
| void serialize(raft::resources const& handle, |
There was a problem hiding this comment.
Order of argument should be handle, params (if applicable), input, output.
There was a problem hiding this comment.
I see the other functions did not follow the proper order. Let's at least make sure the new functions are in the proper order (we can have someone update the others in a follow-up).
| * Unlike the dense overloads, `out_dataset` is required. Nothing can be searched in a VPQ index | ||
| * whose rows were dropped, so there is no use for a graph-only load, and asking for one is an | ||
| * error rather than a silently unusable index. For the same reason `include_dataset = false` | ||
| * produces an index that only `update_dataset` can make searchable again. |
There was a problem hiding this comment.
Please use proper doxygen tags to denote which arguments are inputs and outputs.
| std::unique_ptr<cuvs::neighbors::device_standard_dataset<uint8_t, int64_t>>* | ||
| out_dataset = nullptr); | ||
|
|
||
| /* vpq_f16_index overloads (CAGRA-Q). |
There was a problem hiding this comment.
vpq_f16_index doesn't mean anything to a user reading the docs. Please opt to use more descriptive summaries for comments. Please also don't call this "cagra-q" anymore. It's just cagra with quantized vectors (the vector quantization should be completely decoupled from the graph).
| void deserialize( | ||
| raft::resources const& handle, | ||
| const std::string& filename, | ||
| cuvs::neighbors::cagra::vpq_f16_index<float>* index, |
There was a problem hiding this comment.
Please make the output dataset optional. The graph does not need to be deserialized with the vectors and we should not enforce this. Also- we would like to not require the dataset be known up front wen the user is deserializing a graph that might happen to have a dataset. Goal is to have it deserialize to some dataset owned by the user, but the user shouldn't have to care about its concrete type.
In the C layer, we play to make this even more abstract.
| std::istream& is, | ||
| std::unique_ptr<cuvs::neighbors::device_vpq_dataset<half, int64_t>>* out_dataset) | ||
| { | ||
| RAFT_EXPECTS(out_dataset != nullptr, "pq::deserialize: out_dataset must not be null"); |
There was a problem hiding this comment.
We can't do this- the whole point in decoupling the dataset from the index is to satisfy use-cases where the user has already pq quantized their dataset and just wants to construct a graph. They should not be forced to store the pq vectors with the index.
2 participants
Main functionality: A VPQ dataset could be built and loaded but not saved: serialize
static_assertson any dataset view that is not dense, so compressed rows were requantized on every run and a CAGRA-Q index could not be persisted at all.This restores the VPQ writer in the dataset serializer and exposes it twice:
pq::serialize / pq::deserializefor adevice_vpq_dataset<half, int64_t>, over a stream or a filename, so a dataset can be compressed once offline and reloaded; andcagra::serialize / deserializeforvpq_f16_indexunder a newdevice_vpq_f16dataset kind, so the compressed rows travel inside the index file.There
out_datasetis required rather than optional, since those rows exist nowhere else and a graph-only load would hand back an index that cannot answer a query.A standalone
.vpqcarries the preamble the CAGRA index files already use, aNumPydtypeprefix then a version, so a future change to the encoded row layout fails cleanly instead of being misread; the index format version is unchanged, as the new kind is only an added enumerator.Also here: the C API names the VPQ kind instead of reporting such a file as invalid, and
MG_C_TESTis registered only whenBUILD_MG_ALGOSis on, where it previously failed to link.Tests:
PREPROCESSING_TESTgains*VpqSerialization*,which round-trips thepq_bits,pq_dimcombinations CAGRA-Q accepts, comparing both raw bytes and decoded vectors, and covers the refusals: an empty stream, a foreigndtypeprefix, a version from the future, a truncated payload and a null out-parameter.NEIGHBORS_ANN_CAGRA_FLOAT_UINT32_TESTgains*CagraVpqSerialize*, which saves a CAGRA-Q index, loads it back, and checks that it returns the same neighbours with the dense rows gone. It also covers saving the graph on its own, and refusing to load a compressed index as a dense one.