Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 9 additions & 0 deletions bindings/cpp/include/svs/runtime/dynamic_vamana_index.h
Original file line number Diff line number Diff line change
Expand Up @@ -75,6 +75,15 @@ struct SVS_RUNTIME_API DynamicVamanaIndex : public VamanaIndex {

virtual size_t blocksize_bytes() const noexcept = 0;

/// @brief Storage kind currently backing the index.
///
/// When deferred compression is enabled and the threshold has not yet been crossed,
/// this returns the *initial* (uncompressed) storage kind (FP32 / FP16). After the
/// swap, this returns the same value as ``get_storage_kind()`` (the trained target
/// kind). When deferred compression is disabled, this is always equal to
/// ``get_storage_kind()``.
virtual StorageKind get_current_storage_kind() const noexcept = 0;

// Override for VamanaIndex interface
Status add(size_t, const float*) noexcept override {
return Status(
Expand Down
23 changes: 23 additions & 0 deletions bindings/cpp/include/svs/runtime/vamana_index.h
Original file line number Diff line number Diff line change
Expand Up @@ -52,6 +52,29 @@ struct SVS_RUNTIME_API VamanaIndex {

struct DynamicIndexParams {
size_t blocksize_exp = 30;

/// @brief Threshold (in number of valid vectors) at which to train and switch
/// from ``initial_storage_kind`` to the dynamic index's target storage kind.
///
/// When ``0`` (the default) deferred compression is **disabled** and the index is
/// built directly with the requested compressed storage kind (current eager
/// behavior, fully backward compatible).
///
/// When ``> 0`` and the requested target storage kind is a *trained* compressed
/// type (LVQ / LeanVec / SQ), the index is initially built using
/// ``initial_storage_kind`` (uncompressed). Once the live valid-vector count
/// reaches this threshold, statistics are trained from the accumulated data,
/// the dataset is replaced with the trained compressed form, and the existing
/// graph + ID translation are reused (no graph rebuild).
size_t deferred_compression_threshold = 0;

/// @brief Storage kind used while accumulating vectors below the delayed
/// compression threshold. Must be an *untrained* type (FP32 or FP16).
///
/// Defaults to FP32. Ignored when
/// ``deferred_compression_threshold == 0`` or the target storage kind itself is
/// already untrained.
StorageKind initial_storage_kind = StorageKind::FP32;
};

virtual Status add(size_t n, const float* x) noexcept = 0;
Expand Down
4 changes: 4 additions & 0 deletions bindings/cpp/src/dynamic_vamana_index.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -65,6 +65,10 @@ struct DynamicVamanaIndexManagerBase : public DynamicVamanaIndex {

size_t blocksize_bytes() const noexcept { return impl_->blocksize_bytes(); }

StorageKind get_current_storage_kind() const noexcept override {
return impl_->get_current_storage_kind();
}

Status
remove_selected(size_t* num_removed, const IDFilter& selector) noexcept override {
return runtime_error_wrapper([&] {
Expand Down
362 changes: 357 additions & 5 deletions bindings/cpp/src/dynamic_vamana_index_impl.h

Large diffs are not rendered by default.

72 changes: 72 additions & 0 deletions bindings/cpp/src/dynamic_vamana_index_leanvec_impl.h
Original file line number Diff line number Diff line change
Expand Up @@ -107,6 +107,21 @@ struct DynamicVamanaIndexLeanVecImpl : public DynamicVamanaIndexImpl {
lib::PowerOfTwo blocksize_bytes
) override {
assert(storage::is_leanvec_storage(this->storage_kind_));
if (this->deferred_compression_enabled() &&
data.size() <
this->dynamic_index_params_.deferred_compression_threshold) {
// Delegate the build to the base class (which builds with the
// uncompressed `initial_storage_kind`) and let our overridden
// `setup_deferred_compression_swap` install a LeanVec-aware swap closure.
DynamicVamanaIndexImpl::init_impl(data, labels, blocksize_bytes);
return;
}
// Eager path (also taken when the very first add already meets the deferred
// threshold): build the LeanVec backend directly with the configured
// training data (matrices / leanvec_dims).
if (this->deferred_compression_enabled()) {
this->current_storage_kind_ = this->storage_kind_;
}
impl_.reset(dispatch_leanvec_storage_kind(
this->storage_kind_,
[this](
Expand All @@ -133,10 +148,67 @@ struct DynamicVamanaIndexLeanVecImpl : public DynamicVamanaIndexImpl {
));
}

void setup_deferred_compression_swap(
StorageKind initial_kind, lib::PowerOfTwo blocksize_bytes
) override {
// Capture LeanVec training data by value so the closure does not depend on
// the lifetime of `*this` for those parameters.
LeanVecTrainer trainer{leanvec_dims_, leanvec_matrices_};
storage::dispatch_storage_kind<allocator_type>(
initial_kind,
[&](auto&& tag) {
using Tag = std::decay_t<decltype(tag)>;
this->install_swap_closure_with_trainer<Tag>(
blocksize_bytes, trainer
);
}
);
}

protected:
size_t leanvec_dims_;
std::optional<LeanVecMatricesType> leanvec_matrices_;

/// @brief Trainer used by the deferred-compression swap when the target is a
/// LeanVec storage kind. Reuses pre-trained matrices when supplied; otherwise
/// trains PCA matrices from the accumulated source dataset (the same path the
/// eager builder uses when `leanvec_matrices_ == std::nullopt`).
struct LeanVecTrainer {
size_t leanvec_dims;
std::optional<LeanVecMatricesType> leanvec_matrices;

// Only LeanVec target storage kinds are supported.
template <typename TargetTag>
static constexpr bool supports =
svs::leanvec::IsLeanDataset<typename TargetTag::type>;

template <typename TargetTag, typename Source, typename Pool, typename Alloc>
auto operator()(
TargetTag, const Source& source, Pool& pool, const Alloc& allocator
) const {
using TargetData = typename TargetTag::type;
if constexpr (svs::leanvec::IsLeanDataset<TargetData>) {
size_t d = leanvec_dims;
if (d == 0) {
d = (source.dimensions() + 1) / 2;
}
return TargetData::reduce(
source,
leanvec_matrices,
pool,
0,
svs::lib::MaybeStatic{d},
allocator
);
} else {
static_assert(
!sizeof(TargetData*),
"LeanVecTrainer instantiated for a non-LeanVec target type"
);
}
}
};

StorageKind check_storage_kind(StorageKind kind) {
if (!storage::is_leanvec_storage(kind)) {
throw StatusException(
Expand Down
Loading
Loading