Skip to content
Draft
Show file tree
Hide file tree
Changes from 1 commit
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
372 changes: 367 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