From fd9bca9e35329ba8928736266e9330e8ad0614ad Mon Sep 17 00:00:00 2001
From: Julian Miller <julianmiller@nvidia.com>
Date: Tue, 22 Sep 2026 09:18:16 +0200
Subject: [PATCH 01/13] Add target alpha ACE augmentation

---
 c/include/cuvs/neighbors/cagra.h              |   5 +
 c/include/cuvs/neighbors/hnsw.h               |   5 +
 c/src/neighbors/cagra.cpp                     |   7 +-
 c/src/neighbors/hnsw.cpp                      |   4 +-
 .../src/cuvs/cuvs_ann_bench_param_parser.h    |  11 +
 cpp/include/cuvs/neighbors/cagra.hpp          |  10 +
 .../neighbors/detail/cagra/cagra_build.cuh    | 353 ++++++++++++++----
 cpp/src/neighbors/detail/hnsw.hpp             |   1 +
 fern/pages/c_api/c-api-neighbors-cagra.md     |   2 +
 fern/pages/c_api/c-api-neighbors-hnsw.md      |   2 +
 fern/pages/cpp_api/cpp-api-neighbors-cagra.md |   2 +
 .../python_api/python-api-neighbors-cagra.md  |  10 +-
 .../python_api/python-api-neighbors-hnsw.md   |  10 +-
 python/cuvs/cuvs/neighbors/cagra/cagra.pxd    |   1 +
 python/cuvs/cuvs/neighbors/cagra/cagra.pyx    |  12 +-
 python/cuvs/cuvs/neighbors/hnsw/hnsw.pxd      |   3 +-
 python/cuvs/cuvs/neighbors/hnsw/hnsw.pyx      |  11 +-
 python/cuvs/cuvs/tests/test_cagra_ace.py      |   8 +
 python/cuvs/cuvs/tests/test_hnsw_ace.py       |  10 +-
 19 files changed, 377 insertions(+), 90 deletions(-)

diff --git a/c/include/cuvs/neighbors/cagra.h b/c/include/cuvs/neighbors/cagra.h
index 602f52dbfc..0fc371e3f6 100644
--- a/c/include/cuvs/neighbors/cagra.h
+++ b/c/include/cuvs/neighbors/cagra.h
@@ -199,6 +199,11 @@ struct cuvsAceParams {
    * Useful for testing or when running alongside other memory-intensive processes.
    */
   double max_gpu_memory_gb;
+  /**
+   * Target number of augmented (spill) copies per input vector.
+   * A value of 0 disables augmentation. The default is 1.0.
+   */
+  double target_alpha;
 };
 
 typedef struct cuvsAceParams* cuvsAceParams_t;
diff --git a/c/include/cuvs/neighbors/hnsw.h b/c/include/cuvs/neighbors/hnsw.h
index 58ea5022e0..a25416018f 100644
--- a/c/include/cuvs/neighbors/hnsw.h
+++ b/c/include/cuvs/neighbors/hnsw.h
@@ -90,6 +90,11 @@ struct cuvsHnswAceParams {
    * Useful for testing or when running alongside other memory-intensive processes.
    */
   double max_gpu_memory_gb;
+  /**
+   * Target number of augmented (spill) copies per input vector.
+   * A value of 0 disables augmentation. The default is 1.0.
+   */
+  double target_alpha;
 };
 
 typedef struct cuvsHnswAceParams* cuvsHnswAceParams_t;
diff --git a/c/src/neighbors/cagra.cpp b/c/src/neighbors/cagra.cpp
index 97cc9d8bac..8e7750e0a8 100644
--- a/c/src/neighbors/cagra.cpp
+++ b/c/src/neighbors/cagra.cpp
@@ -827,6 +827,7 @@ static void _set_graph_build_params(
         ace_p.use_disk            = ace_params_c->use_disk;
         ace_p.max_host_memory_gb  = ace_params_c->max_host_memory_gb;
         ace_p.max_gpu_memory_gb   = ace_params_c->max_gpu_memory_gb;
+        ace_p.target_alpha        = ace_params_c->target_alpha;
       }
       out_params = ace_p;
       break;
@@ -1463,6 +1464,9 @@ static void _populate_cagra_index_params_from_cpp(cuvsCagraIndexParams_t c_param
     c_ace_params->ef_construction = ace_params.ef_construction;
     c_ace_params->build_dir = ace_params.build_dir.empty() ? nullptr : strdup(ace_params.build_dir.c_str());
     c_ace_params->use_disk = ace_params.use_disk;
+    c_ace_params->max_host_memory_gb = ace_params.max_host_memory_gb;
+    c_ace_params->max_gpu_memory_gb = ace_params.max_gpu_memory_gb;
+    c_ace_params->target_alpha = ace_params.target_alpha;
     c_params->graph_build_params = c_ace_params;
   }
 }
@@ -2341,7 +2345,8 @@ extern "C" cuvsError_t cuvsAceParamsCreate(cuvsAceParams_t* params)
                                 .build_dir           = build_dir,
                                 .use_disk            = ps.use_disk,
                                 .max_host_memory_gb  = ps.max_host_memory_gb,
-                                .max_gpu_memory_gb   = ps.max_gpu_memory_gb};
+                                .max_gpu_memory_gb   = ps.max_gpu_memory_gb,
+                                .target_alpha        = ps.target_alpha};
   });
 }
 
diff --git a/c/src/neighbors/hnsw.cpp b/c/src/neighbors/hnsw.cpp
index 1f19f33875..e18647eddd 100644
--- a/c/src/neighbors/hnsw.cpp
+++ b/c/src/neighbors/hnsw.cpp
@@ -47,6 +47,7 @@ void _build(cuvsResources_t res,
     ace_params.use_disk           = params->ace_params->use_disk;
     ace_params.max_host_memory_gb = params->ace_params->max_host_memory_gb;
     ace_params.max_gpu_memory_gb  = params->ace_params->max_gpu_memory_gb;
+    ace_params.target_alpha       = params->ace_params->target_alpha;
     cpp_params.graph_build_params = ace_params;
   }
 
@@ -161,7 +162,8 @@ extern "C" cuvsError_t cuvsHnswAceParamsCreate(cuvsHnswAceParams_t* params)
                                     .build_dir           = "/tmp/hnsw_ace_build",
                                     .use_disk            = false,
                                     .max_host_memory_gb  = 0,
-                                    .max_gpu_memory_gb   = 0};
+                                    .max_gpu_memory_gb   = 0,
+                                    .target_alpha        = 1.0};
   });
 }
 
diff --git a/cpp/bench/ann/src/cuvs/cuvs_ann_bench_param_parser.h b/cpp/bench/ann/src/cuvs/cuvs_ann_bench_param_parser.h
index 3b5dfd16ab..b78fb32b0c 100644
--- a/cpp/bench/ann/src/cuvs/cuvs_ann_bench_param_parser.h
+++ b/cpp/bench/ann/src/cuvs/cuvs_ann_bench_param_parser.h
@@ -406,6 +406,14 @@ void parse_build_param(const nlohmann::json& conf, cuvs::neighbors::cagra::index
       } else if constexpr (std::is_same_v<U,
                                           cuvs::neighbors::graph_build_params::nn_descent_params>) {
         parse_build_param<T, IdxT>(nn_descent_conf, arg);
+      } else if constexpr (std::is_same_v<U, cuvs::neighbors::graph_build_params::ace_params>) {
+        if (ace_conf.contains("npartitions")) { arg.npartitions = ace_conf.at("npartitions"); }
+        if (ace_conf.contains("build_dir")) { arg.build_dir = ace_conf.at("build_dir"); }
+        if (ace_conf.contains("ef_construction")) {
+          arg.ef_construction = ace_conf.at("ef_construction");
+        }
+        if (ace_conf.contains("use_disk")) { arg.use_disk = ace_conf.at("use_disk"); }
+        if (ace_conf.contains("target_alpha")) { arg.target_alpha = ace_conf.at("target_alpha"); }
       }
     },
     params.graph_build_params);
@@ -476,6 +484,9 @@ void parse_build_param(const nlohmann::json& conf,
         ace_params.ef_construction = ace_conf.at("ef_construction");
       }
       if (ace_conf.contains("use_disk")) { ace_params.use_disk = ace_conf.at("use_disk"); }
+      if (ace_conf.contains("target_alpha")) {
+        ace_params.target_alpha = ace_conf.at("target_alpha");
+      }
       cagra_params.graph_build_params = ace_params;
     }
     ::parse_build_param<T, IdxT>(conf, cagra_params);
diff --git a/cpp/include/cuvs/neighbors/cagra.hpp b/cpp/include/cuvs/neighbors/cagra.hpp
index 9969d6bf95..372ae8e0d7 100644
--- a/cpp/include/cuvs/neighbors/cagra.hpp
+++ b/cpp/include/cuvs/neighbors/cagra.hpp
@@ -231,6 +231,16 @@ struct ace_params {
    */
   double max_gpu_memory_gb = 0;
 
+  /**
+   * Target number of augmented (spill) copies per input vector.
+   *
+   * ACE samples boundary-margin scores and chooses a global threshold whose expected
+   * augmentation ratio matches this value. A value of 0 disables augmentation. Values above 1
+   * allow a vector to spill into multiple nearby partitions. The implementation bounds each
+   * vector to three spill partitions to keep memory and build work predictable.
+   */
+  double target_alpha = 1.0;
+
   ace_params() = default;
 };
 
diff --git a/cpp/src/neighbors/detail/cagra/cagra_build.cuh b/cpp/src/neighbors/detail/cagra/cagra_build.cuh
index 9a22676334..67a0b98e6a 100644
--- a/cpp/src/neighbors/detail/cagra/cagra_build.cuh
+++ b/cpp/src/neighbors/detail/cagra/cagra_build.cuh
@@ -46,10 +46,12 @@
 #include <array>
 #include <cerrno>
 #include <chrono>
+#include <cmath>
 #include <cstdio>
 #include <cstring>
 #include <exception>
 #include <filesystem>
+#include <limits>
 #include <omp.h>
 #include <optional>
 #include <type_traits>
@@ -180,8 +182,112 @@ void check_graph_degree(size_t& intermediate_degree, size_t& graph_degree, size_
   }
 }
 
-// ACE: Get partition labels for partitioned approach
-// TODO(julianmi): Use all neighbors APIs.
+// A small per-vector cap bounds worst-case partition imbalance while target_alpha controls the
+// global augmentation budget. The cap is intentionally internal: the spill-cap sweep showed no
+// consistent quality benefit from exposing larger values.
+constexpr size_t ace_max_spills_per_vector = 3;
+
+template <typename IdxT>
+constexpr IdxT ace_invalid_partition()
+{
+  return std::numeric_limits<IdxT>::max();
+}
+
+inline size_t ace_select_nearest_partitions(const float* distances,
+                                            size_t n_partitions,
+                                            size_t n_keep,
+                                            std::pair<float, size_t>* nearest)
+{
+  size_t filled = 0;
+  for (size_t partition = 0; partition < n_partitions; partition++) {
+    const float distance = distances[partition];
+    if (filled == n_keep && distance >= nearest[n_keep - 1].first) { continue; }
+    size_t position = filled < n_keep ? filled++ : n_keep - 1;
+    for (; position > 0 && distance < nearest[position - 1].first; position--) {
+      nearest[position] = nearest[position - 1];
+    }
+    nearest[position] = {distance, partition};
+  }
+  return filled;
+}
+
+inline float ace_boundary_margin(float core_distance, float spill_distance, float centroid_distance)
+{
+  return centroid_distance > 0.0f
+           ? (spill_distance - core_distance) / (2.0f * std::sqrt(centroid_distance))
+           : std::numeric_limits<float>::infinity();
+}
+
+inline std::vector<float> ace_sample_margin_scores(
+  raft::resources const& res,
+  raft::device_matrix_view<const float, int64_t> sample,
+  raft::device_matrix_view<const float, int64_t> centroids,
+  raft::host_matrix_view<const float, int64_t> centroid_distances,
+  size_t max_spills)
+{
+  const size_t n_samples    = sample.extent(0);
+  const size_t dataset_dim  = sample.extent(1);
+  const size_t n_partitions = centroids.extent(0);
+  const size_t chunk_size   = std::min<size_t>(n_samples, 32 * 1024);
+  auto distances_host       = raft::make_host_matrix<float, int64_t>(chunk_size, n_partitions);
+  auto distances_device = raft::make_device_matrix<float, int64_t>(res, chunk_size, n_partitions);
+  auto stream           = raft::resource::get_cuda_stream(res).get();
+
+  std::vector<float> scores;
+  scores.reserve(n_samples * max_spills);
+  for (size_t base = 0; base < n_samples; base += chunk_size) {
+    const size_t rows = std::min(chunk_size, n_samples - base);
+    auto sample_chunk = raft::make_device_matrix_view<const float, int64_t>(
+      sample.data_handle() + base * dataset_dim, rows, dataset_dim);
+    auto distances_chunk = raft::make_device_matrix_view<float, int64_t>(
+      distances_device.data_handle(), rows, n_partitions);
+    cuvs::distance::pairwise_distance(
+      res, sample_chunk, centroids, distances_chunk, cuvs::distance::DistanceType::L2Expanded);
+    auto distances_chunk_host =
+      raft::make_host_matrix_view<float, int64_t>(distances_host.data_handle(), rows, n_partitions);
+    raft::copy(res, distances_chunk_host, raft::make_const_mdspan(distances_chunk));
+    raft::resource::sync_stream(res, stream);
+
+#pragma omp parallel
+    {
+      std::vector<float> local_scores;
+      local_scores.reserve(rows * max_spills / omp_get_num_threads() + max_spills);
+      std::vector<std::pair<float, size_t>> nearest(1 + max_spills);
+#pragma omp for
+      for (size_t row = 0; row < rows; row++) {
+        const size_t found = ace_select_nearest_partitions(
+          &distances_chunk_host(row, 0), n_partitions, 1 + max_spills, nearest.data());
+        for (size_t spill = 1; spill < found; spill++) {
+          local_scores.push_back(
+            ace_boundary_margin(nearest[0].first,
+                                nearest[spill].first,
+                                centroid_distances(nearest[0].second, nearest[spill].second)));
+        }
+      }
+#pragma omp critical
+      scores.insert(scores.end(), local_scores.begin(), local_scores.end());
+    }
+  }
+  return scores;
+}
+
+inline double ace_target_alpha_threshold(std::vector<float> scores,
+                                         size_t n_samples,
+                                         size_t max_spills,
+                                         double target_alpha)
+{
+  if (scores.empty() || target_alpha <= 0.0) { return -std::numeric_limits<double>::infinity(); }
+  if (target_alpha >= static_cast<double>(max_spills)) {
+    return std::numeric_limits<double>::infinity();
+  }
+  size_t keep = static_cast<size_t>(std::llround(target_alpha * n_samples));
+  keep        = std::clamp<size_t>(keep, 1, scores.size());
+  std::nth_element(scores.begin(), scores.begin() + keep - 1, scores.end());
+  return scores[keep - 1];
+}
+
+// ACE: Get core plus boundary-gated spill partition labels. Only core graph rows are retained;
+// spill vectors enlarge local builds but their graph rows are not merged into the final graph.
 template <typename T, typename IdxT>
 void ace_get_partition_labels(
   raft::resources const& res,
@@ -189,7 +295,7 @@ void ace_get_partition_labels(
   size_t dataset_dim,
   raft::host_matrix_view<IdxT, int64_t, raft::row_major> partition_labels,
   raft::host_matrix_view<IdxT, int64_t, raft::row_major> partition_histogram,
-  size_t min_partition_size,
+  double target_alpha,
   double sampling_rate = 0.01)
 {
   size_t dataset_size = dataset.extent(0);
@@ -202,8 +308,14 @@ void ace_get_partition_labels(
                static_cast<size_t>(dataset.extent(1)),
                dataset_dim);
   size_t n_partitions = partition_histogram.extent(0);
-  RAFT_EXPECTS(labels_dim == 2, "Labels must have 2 columns");
+  RAFT_EXPECTS(labels_dim >= 2 && labels_dim <= n_partitions,
+               "ACE labels must contain one core and at least one spill column");
   RAFT_EXPECTS(partition_histogram.extent(1) == 2, "Partition histogram must have 2 columns");
+  RAFT_EXPECTS(std::isfinite(target_alpha) && target_alpha >= 0.0 &&
+                 target_alpha <= static_cast<double>(labels_dim - 1),
+               "ACE target_alpha must be finite and in [0, %zu]",
+               labels_dim - 1);
+  RAFT_EXPECTS(n_partitions <= 8192, "ACE target_alpha currently supports at most 8192 partitions");
   cudaStream_t stream = raft::resource::get_cuda_stream(res).get();
 
   // Sampling vectors from dataset. Uses float conversion on host instead of
@@ -231,6 +343,31 @@ void ace_get_partition_labels(
   auto centroids_dev = raft::make_device_matrix<float, int64_t>(res, n_partitions, dataset_dim);
   cuvs::cluster::kmeans::fit(res, kmeans_params, sample_db_dev.view(), centroids_dev.view());
 
+  auto centroid_distances = raft::make_host_matrix<float, int64_t>(n_partitions, n_partitions);
+  auto centroid_distances_dev =
+    raft::make_device_matrix<float, int64_t>(res, n_partitions, n_partitions);
+  cuvs::distance::pairwise_distance(res,
+                                    raft::make_const_mdspan(centroids_dev.view()),
+                                    raft::make_const_mdspan(centroids_dev.view()),
+                                    centroid_distances_dev.view(),
+                                    cuvs::distance::DistanceType::L2Expanded);
+  raft::copy(
+    res, centroid_distances.view(), raft::make_const_mdspan(centroid_distances_dev.view()));
+  raft::resource::sync_stream(res, stream);
+
+  const size_t max_spills       = labels_dim - 1;
+  const double margin_threshold = ace_target_alpha_threshold(
+    ace_sample_margin_scores(res,
+                             raft::make_const_mdspan(sample_db_dev.view()),
+                             raft::make_const_mdspan(centroids_dev.view()),
+                             raft::make_const_mdspan(centroid_distances.view()),
+                             max_spills),
+    n_samples,
+    max_spills,
+    target_alpha);
+  RAFT_LOG_INFO(
+    "ACE: target_alpha %.3f -> boundary margin threshold %.6f", target_alpha, margin_threshold);
+
   // Compute distances between dataset and centroid vectors
   // Uses float conversion on host instead of batch_load_iterator to minimize GPU memory usage.
   const size_t chunk_size = 32 * 1024;
@@ -272,31 +409,34 @@ void ace_get_partition_labels(
     raft::copy(res, sub_distances, sub_distances_dev);
     raft::resource::sync_stream(res, stream);
 
-    // Find two closest partitions to each dataset vector
-#pragma omp parallel for
-    for (size_t i_sub = 0; i_sub < sub_dataset_size; i_sub++) {
-      size_t core_label      = 0;
-      size_t augmented_label = 1;
-      if (sub_distances(i_sub, 0) > sub_distances(i_sub, 1)) {
-        core_label      = 1;
-        augmented_label = 0;
-      }
-      for (size_t c = 2; c < n_partitions; c++) {
-        if (sub_distances(i_sub, c) < sub_distances(i_sub, core_label)) {
-          augmented_label = core_label;
-          core_label      = c;
-        } else if (sub_distances(i_sub, c) < sub_distances(i_sub, augmented_label)) {
-          augmented_label = c;
-        }
-      }
-      size_t i               = i_base + i_sub;
-      partition_labels(i, 0) = core_label;
-      partition_labels(i, 1) = augmented_label;
-
+    // Find the nearest core and up to three gated spill partitions per vector.
+#pragma omp parallel
+    {
+      std::vector<std::pair<float, size_t>> nearest(labels_dim);
+#pragma omp for
+      for (size_t i_sub = 0; i_sub < sub_dataset_size; i_sub++) {
+        const size_t output_row = i_base + i_sub;
+        const size_t found      = ace_select_nearest_partitions(
+          &sub_distances(i_sub, 0), n_partitions, labels_dim, nearest.data());
+        partition_labels(output_row, 0) = static_cast<IdxT>(nearest[0].second);
 #pragma omp atomic update
-      partition_histogram(core_label, 0) += 1;
+        partition_histogram(nearest[0].second, 0) += 1;
+
+        size_t output_slot = 1;
+        for (size_t rank = 1; rank < found; rank++) {
+          const float margin =
+            ace_boundary_margin(nearest[0].first,
+                                nearest[rank].first,
+                                centroid_distances(nearest[0].second, nearest[rank].second));
+          if (margin > margin_threshold) { continue; }
+          partition_labels(output_row, output_slot++) = static_cast<IdxT>(nearest[rank].second);
 #pragma omp atomic update
-      partition_histogram(augmented_label, 1) += 1;
+          partition_histogram(nearest[rank].second, 1) += 1;
+        }
+        for (; output_slot < labels_dim; output_slot++) {
+          partition_labels(output_row, output_slot) = ace_invalid_partition<IdxT>();
+        }
+      }
     }
 
     const auto completed_rows = i_base + sub_dataset_size;
@@ -350,8 +490,9 @@ void ace_check_partition_sizes(
 
   double avg_core_vectors      = static_cast<double>(total_core_vectors) / n_partitions;
   double avg_augmented_vectors = static_cast<double>(total_augmented_vectors) / n_partitions;
-  double avg_total_vectors     = 2.0 * static_cast<double>(dataset_size) / n_partitions;
-  double expected_avg_vectors  = 2.0 * static_cast<double>(dataset_size) / n_partitions;
+  double avg_total_vectors =
+    static_cast<double>(total_core_vectors + total_augmented_vectors) / n_partitions;
+  double expected_avg_vectors = avg_total_vectors;
 
   RAFT_LOG_INFO("ACE: Core vectors        - Total: %lu, Avg: %.1f, Min: %lu, Max: %lu",
                 total_core_vectors,
@@ -417,12 +558,18 @@ void ace_create_forward_and_backward_lists(
       augmented_partition_offsets(c - 1) + partition_histogram(c - 1, 1);
   }
 
+  const size_t labels_dim = partition_labels.extent(1);
+  size_t total_augmented  = 0;
+  for (size_t partition = 0; partition < n_partitions; partition++) {
+    total_augmented += partition_histogram(partition, 1);
+  }
+  RAFT_EXPECTS(static_cast<size_t>(augmented_backward_mapping.extent(0)) == total_augmented,
+               "augmented_backward_mapping must have one entry per spill copy");
+
   if (static_cast<size_t>(core_forward_mapping.extent(0)) == 0) {
     // Memory path: both backward mappings
     RAFT_EXPECTS(static_cast<size_t>(core_backward_mapping.extent(0)) == dataset_size,
                  "core_backward_mapping must be of size dataset_size");
-    RAFT_EXPECTS(static_cast<size_t>(augmented_backward_mapping.extent(0)) == dataset_size,
-                 "augmented_backward_mapping must be of size dataset_size");
 #pragma omp parallel for
     for (size_t i = 0; i < dataset_size; i++) {
       size_t core_partition_id = partition_labels(i, 0);
@@ -432,12 +579,16 @@ void ace_create_forward_and_backward_lists(
       RAFT_EXPECTS(core_id < dataset_size, "Vector ID must be smaller than dataset_size");
       core_backward_mapping(core_id) = i;
 
-      size_t augmented_partition_id = partition_labels(i, 1);
-      size_t augmented_id;
+      for (size_t spill = 1; spill < labels_dim; spill++) {
+        if (partition_labels(i, spill) == ace_invalid_partition<IdxT>()) { break; }
+        const size_t augmented_partition_id = partition_labels(i, spill);
+        size_t augmented_id;
 #pragma omp atomic capture
-      augmented_id = augmented_partition_offsets(augmented_partition_id)++;
-      RAFT_EXPECTS(augmented_id < dataset_size, "Vector ID must be smaller than dataset_size");
-      augmented_backward_mapping(augmented_id) = i;
+        augmented_id = augmented_partition_offsets(augmented_partition_id)++;
+        RAFT_EXPECTS(augmented_id < total_augmented,
+                     "Spill ID must be smaller than total augmentation");
+        augmented_backward_mapping(augmented_id) = i;
+      }
     }
   } else {
     // Disk path: all three mappings
@@ -445,8 +596,6 @@ void ace_create_forward_and_backward_lists(
                  "core_forward_mapping must be of size dataset_size");
     RAFT_EXPECTS(static_cast<size_t>(core_backward_mapping.extent(0)) == dataset_size,
                  "core_backward_mapping must be of size dataset_size");
-    RAFT_EXPECTS(static_cast<size_t>(augmented_backward_mapping.extent(0)) == dataset_size,
-                 "augmented_backward_mapping must be of size dataset_size");
     for (size_t i = 0; i < dataset_size; i++) {
       size_t core_partition_id = partition_labels(i, 0);
       size_t core_id;
@@ -455,11 +604,14 @@ void ace_create_forward_and_backward_lists(
       core_backward_mapping(core_id) = i;
       core_forward_mapping(i)        = core_id;
 
-      size_t augmented_partition_id = partition_labels(i, 1);
-      size_t augmented_id;
-      augmented_id = augmented_partition_offsets(augmented_partition_id)++;
-      RAFT_EXPECTS(augmented_id < dataset_size, "Vector ID must be smaller than dataset_size");
-      augmented_backward_mapping(augmented_id) = i;
+      for (size_t spill = 1; spill < labels_dim; spill++) {
+        if (partition_labels(i, spill) == ace_invalid_partition<IdxT>()) { break; }
+        const size_t augmented_partition_id = partition_labels(i, spill);
+        const size_t augmented_id           = augmented_partition_offsets(augmented_partition_id)++;
+        RAFT_EXPECTS(augmented_id < total_augmented,
+                     "Spill ID must be smaller than total augmentation");
+        augmented_backward_mapping(augmented_id) = i;
+      }
     }
   }
 
@@ -669,8 +821,9 @@ void ace_reorder_and_store_dataset(
   }
   RAFT_EXPECTS(total_core_vectors == dataset_size,
                "Total core vectors must be equal to dataset size");
-  RAFT_EXPECTS(total_augmented_vectors == dataset_size,
-               "Total augmented vectors must be equal to dataset size");
+  RAFT_EXPECTS(
+    total_augmented_vectors <= (static_cast<size_t>(partition_labels.extent(1)) - 1) * dataset_size,
+    "Total augmented vectors exceed the internal spill cap");
 
   // Pre-allocate file space for better performance
   const size_t vector_size   = dataset_dim * sizeof(T);
@@ -764,8 +917,7 @@ void ace_reorder_and_store_dataset(
   size_t vectors_processed     = 0;
   size_t next_progress_percent = 10;
   for (size_t i = 0; i < dataset_size; i++) {
-    size_t core_partition      = partition_labels(i, 0);
-    size_t secondary_partition = partition_labels(i, 1);
+    const size_t core_partition = partition_labels(i, 0);
 
     // Add vector to core partition buffer
     size_t core_buffer_row = core_buffer_counts(core_partition);
@@ -778,16 +930,17 @@ void ace_reorder_and_store_dataset(
       flush_core_buffer(core_partition);
     }
 
-    // Add vector to augmented partition buffer
-    size_t augmented_buffer_row = augmented_buffer_counts(secondary_partition);
-    memcpy(&augmented_buffers[secondary_partition](augmented_buffer_row, 0),
-           &dataset(i, 0),
-           dataset_dim * sizeof(T));
-    augmented_buffer_counts(secondary_partition)++;
-
-    // Flush augmented buffer if full
-    if (augmented_buffer_counts(secondary_partition) >= vectors_per_buffer) {
-      flush_augmented_buffer(secondary_partition);
+    for (size_t spill = 1; spill < static_cast<size_t>(partition_labels.extent(1)); spill++) {
+      if (partition_labels(i, spill) == ace_invalid_partition<IdxT>()) { break; }
+      const size_t spill_partition  = partition_labels(i, spill);
+      const size_t spill_buffer_row = augmented_buffer_counts(spill_partition);
+      memcpy(&augmented_buffers[spill_partition](spill_buffer_row, 0),
+             &dataset(i, 0),
+             dataset_dim * sizeof(T));
+      augmented_buffer_counts(spill_partition)++;
+      if (augmented_buffer_counts(spill_partition) >= vectors_per_buffer) {
+        flush_augmented_buffer(spill_partition);
+      }
     }
 
     vectors_processed++;
@@ -980,8 +1133,16 @@ constexpr double usable_cpu_memory_fraction = 0.8;
 // Factor to account for imbalances in the partitions (maximum allowed is 3x the average)
 constexpr double imbalance_factor = 3.0;
 
-// Current partitioning adds each vector into 2 partitions (core and augmented)
-constexpr double vector_expansion_factor = 2.0;
+constexpr double ace_label_expansion_factor = 1.0 + ace_max_spills_per_vector;
+
+inline double ace_estimated_vector_expansion_factor(double target_alpha)
+{
+  // Sampled thresholds closely track target_alpha in the validated sweeps. Keep 10% plus 0.1
+  // spill/vector as admission headroom without pessimistically charging every build for the cap.
+  const double estimated_spills =
+    std::min<double>(ace_max_spills_per_vector, std::max(0.0, target_alpha * 1.1 + 0.1));
+  return 1.0 + estimated_spills;
+}
 
 // Check if disk mode should be used for ACE based on memory constraints
 template <typename T, typename IdxT>
@@ -993,6 +1154,7 @@ bool ace_check_use_disk_mode(raft::resources const& res,
                              size_t n_partitions,
                              size_t intermediate_degree,
                              size_t graph_degree,
+                             double vector_expansion_factor,
                              std::optional<double> max_host_memory_gb,
                              std::optional<double> max_gpu_memory_gb,
                              bool guarantee_connectivity,
@@ -1048,7 +1210,7 @@ bool ace_check_use_disk_mode(raft::resources const& res,
 
   // Optimistic memory model: focus on largest arrays, assumes all partitions are of equal size
   // For memory path:
-  //   - Partition labels (core + augmented): vector_expansion_factor * dataset_size * sizeof(IdxT)
+  //   - Partition labels (core + the fixed internal spill slots)
   //   - Backward ID mapping arrays (core + augmented): vector_expansion_factor * dataset_size *
   //   sizeof(IdxT)
   //   - Avg. per-partition dataset: vector_expansion_factor * (dataset_size / n_partitions) *
@@ -1056,7 +1218,7 @@ bool ace_check_use_disk_mode(raft::resources const& res,
   //   - Avg. per-partition graph during build: vector_expansion_factor * (dataset_size /
   //   n_partitions) * (intermediate + final) * sizeof(IdxT)
   //   - Final assembled graph: dataset_size * graph_degree * sizeof(IdxT)
-  mem.partition_labels_size = vector_expansion_factor * dataset_size * sizeof(IdxT);
+  mem.partition_labels_size = ace_label_expansion_factor * dataset_size * sizeof(IdxT);
   mem.id_mapping_size       = vector_expansion_factor * dataset_size * sizeof(IdxT);
   mem.sub_dataset_size      = sub_partition_size * dataset_dim * sizeof(T);
   mem.sub_graph_size   = sub_partition_size * (intermediate_degree + graph_degree) * sizeof(IdxT);
@@ -1160,12 +1322,13 @@ void ace_validate_disk_mode_partitions(raft::resources const& res,
                                        size_t dataset_dim,
                                        size_t intermediate_degree,
                                        size_t graph_degree,
+                                       double vector_expansion_factor,
                                        bool guarantee_connectivity,
                                        ace_memory_requirements& mem)
 {
   // In disk mode, we don't need the full dataset or final graph in memory.
   // Host memory model for disk mode:
-  //   - Partition labels (core + augmented): vector_expansion_factor * dataset_size * sizeof(IdxT)
+  //   - Partition labels (core + the fixed internal spill slots)
   //   - ID mapping arrays (core + augmented): vector_expansion_factor * dataset_size * sizeof(IdxT)
   //   - Avg. per-partition dataset during processing: vector_expansion_factor * (dataset_size /
   //   n_partitions) * dataset_dim * sizeof(T)
@@ -1318,8 +1481,8 @@ auto build_from_host_matrix(raft::resources const& res,
 
 // Build CAGRA index using ACE (Augmented Core Extraction) partitioning
 // ACE enables building indexes for datasets too large to fit in GPU memory by:
-// 1. Partitioning the dataset using balanced k-means in core (non-overlapping) and augmented
-// (second-closest) partitions
+// 1. Partitioning the dataset using balanced k-means into a non-overlapping core and
+// boundary-gated spill partitions
 // 2. Building sub-indexes for each partition independently
 // 3. Concatenating sub-graphs (of core partitions) into a final unified index
 // Supports both in-memory and disk-based modes depending on available host memory.
@@ -1337,9 +1500,15 @@ auto build_ace(raft::resources const& res, const index_params& params, DatasetVi
 
   auto ace_params    = std::get<cagra::graph_build_params::ace_params>(params.graph_build_params);
   size_t npartitions = ace_params.npartitions;
-  size_t ef_construction = ace_params.ef_construction;
-  std::string build_dir  = ace_params.build_dir;
-  bool use_disk          = ace_params.use_disk;
+  size_t ef_construction    = ace_params.ef_construction;
+  std::string build_dir     = ace_params.build_dir;
+  bool use_disk             = ace_params.use_disk;
+  const double target_alpha = ace_params.target_alpha;
+  RAFT_EXPECTS(std::isfinite(target_alpha) && target_alpha >= 0.0 &&
+                 target_alpha <= static_cast<double>(ace_max_spills_per_vector),
+               "ACE target_alpha must be finite and in [0, %zu]",
+               ace_max_spills_per_vector);
+  const double vector_expansion_factor = ace_estimated_vector_expansion_factor(target_alpha);
 
   common::nvtx::range<common::nvtx::domain::cuvs> function_scope(
     "cagra::detail::build_ace<host>(%zu, %zu, %zu)",
@@ -1405,6 +1574,7 @@ auto build_ace(raft::resources const& res, const index_params& params, DatasetVi
                                                           n_partitions,
                                                           intermediate_degree,
                                                           graph_degree,
+                                                          vector_expansion_factor,
                                                           ace_params.max_host_memory_gb,
                                                           ace_params.max_gpu_memory_gb,
                                                           params.guarantee_connectivity,
@@ -1418,11 +1588,19 @@ auto build_ace(raft::resources const& res, const index_params& params, DatasetVi
                                                  dataset_dim,
                                                  intermediate_degree,
                                                  graph_degree,
+                                                 vector_expansion_factor,
                                                  params.guarantee_connectivity,
                                                  mem);
       ace_validate_partition_count(n_partitions, dataset_size, true);
     }
 
+    const size_t max_spills = std::min(ace_max_spills_per_vector, n_partitions - 1);
+    RAFT_EXPECTS(std::isfinite(target_alpha) && target_alpha >= 0.0 &&
+                   target_alpha <= static_cast<double>(max_spills),
+                 "ACE target_alpha must be finite and in [0, %zu] for %zu partitions",
+                 max_spills,
+                 n_partitions);
+
     // Preallocate space for files for better performance and fail early if not enough space.
     cuvs::util::file_descriptor reordered_fd;
     cuvs::util::file_descriptor augmented_fd;
@@ -1443,12 +1621,6 @@ auto build_ace(raft::resources const& res, const index_params& params, DatasetVi
         true);
       workspace.mark_artifact_created(ace_disk_workspace::artifact::reordered_dataset);
 
-      std::tie(augmented_fd, augmented_header_size) = cuvs::util::create_numpy_file<T>(
-        workspace.artifact_path(ace_disk_workspace::artifact::augmented_dataset),
-        {dataset_size, dataset_dim},
-        true);
-      workspace.mark_artifact_created(ace_disk_workspace::artifact::augmented_dataset);
-
       std::tie(mapping_fd, mapping_header_size) = cuvs::util::create_numpy_file<IdxT>(
         workspace.artifact_path(ace_disk_workspace::artifact::dataset_mapping),
         {dataset_size},
@@ -1461,16 +1633,14 @@ auto build_ace(raft::resources const& res, const index_params& params, DatasetVi
         true);
       workspace.mark_artifact_created(ace_disk_workspace::artifact::cagra_graph);
 
-      RAFT_LOG_DEBUG(
-        "ACE build: wrote numpy headers (reordered=%zu augmented=%zu mapping=%zu graph=%zu bytes)",
-        reordered_header_size,
-        augmented_header_size,
-        mapping_header_size,
-        graph_header_size);
+      RAFT_LOG_DEBUG("ACE build: wrote numpy headers (reordered=%zu mapping=%zu graph=%zu bytes)",
+                     reordered_header_size,
+                     mapping_header_size,
+                     graph_header_size);
     }
 
     auto partition_start     = std::chrono::high_resolution_clock::now();
-    auto partition_labels    = raft::make_host_matrix<IdxT, int64_t>(dataset_size, 2);
+    auto partition_labels    = raft::make_host_matrix<IdxT, int64_t>(dataset_size, 1 + max_spills);
     auto partition_histogram = raft::make_host_matrix<IdxT, int64_t>(n_partitions, 2);
     for (size_t c = 0; c < n_partitions; c++) {
       partition_histogram(c, 0) = 0;
@@ -1485,7 +1655,7 @@ auto build_ace(raft::resources const& res, const index_params& params, DatasetVi
                                       dataset_dim,
                                       partition_labels.view(),
                                       partition_histogram.view(),
-                                      min_partition_size);
+                                      target_alpha);
 
     ace_check_partition_sizes<IdxT>(dataset_size,
                                     n_partitions,
@@ -1501,12 +1671,31 @@ auto build_ace(raft::resources const& res, const index_params& params, DatasetVi
                   partition_elapsed,
                   min_partition_size);
 
+    size_t total_augmented_vectors = 0;
+    for (size_t partition = 0; partition < n_partitions; partition++) {
+      total_augmented_vectors += partition_histogram(partition, 1);
+    }
+    RAFT_LOG_INFO("ACE build: realized target_alpha %.3f (requested %.3f, spill_cap=%zu)",
+                  static_cast<double>(total_augmented_vectors) / dataset_size,
+                  target_alpha,
+                  max_spills);
+
+    if (use_disk_mode) {
+      std::tie(augmented_fd, augmented_header_size) = cuvs::util::create_numpy_file<T>(
+        workspace.artifact_path(ace_disk_workspace::artifact::augmented_dataset),
+        {total_augmented_vectors, dataset_dim},
+        true);
+      workspace.mark_artifact_created(ace_disk_workspace::artifact::augmented_dataset);
+      RAFT_LOG_DEBUG("ACE build: wrote augmented numpy header (%zu bytes)", augmented_header_size);
+    }
+
     // Create vector lists for each partition
     auto vectorlist_start      = std::chrono::high_resolution_clock::now();
     auto core_forward_mapping  = use_disk_mode ? raft::make_host_vector<IdxT, int64_t>(dataset_size)
                                                : raft::make_host_vector<IdxT, int64_t>(0);
     auto core_backward_mapping = raft::make_host_vector<IdxT, int64_t>(dataset_size);
-    auto augmented_backward_mapping  = raft::make_host_vector<IdxT, int64_t>(dataset_size);
+    auto augmented_backward_mapping =
+      raft::make_host_vector<IdxT, int64_t>(total_augmented_vectors);
     auto core_partition_offsets      = raft::make_host_vector<IdxT, int64_t>(n_partitions + 1);
     auto augmented_partition_offsets = raft::make_host_vector<IdxT, int64_t>(n_partitions + 1);
 
diff --git a/cpp/src/neighbors/detail/hnsw.hpp b/cpp/src/neighbors/detail/hnsw.hpp
index 19c0afd820..e4d134fae3 100644
--- a/cpp/src/neighbors/detail/hnsw.hpp
+++ b/cpp/src/neighbors/detail/hnsw.hpp
@@ -2998,6 +2998,7 @@ std::unique_ptr<index<T>> build(raft::resources const& res,
     cagra_ace_params.use_disk           = ace_params.use_disk;
     cagra_ace_params.max_host_memory_gb = ace_params.max_host_memory_gb;
     cagra_ace_params.max_gpu_memory_gb  = ace_params.max_gpu_memory_gb;
+    cagra_ace_params.target_alpha       = ace_params.target_alpha;
     cagra_params.graph_build_params     = cagra_ace_params;
   }
 
diff --git a/fern/pages/c_api/c-api-neighbors-cagra.md b/fern/pages/c_api/c-api-neighbors-cagra.md
index 13c663a7c8..d77008ab2c 100644
--- a/fern/pages/c_api/c-api-neighbors-cagra.md
+++ b/fern/pages/c_api/c-api-neighbors-cagra.md
@@ -86,6 +86,7 @@ struct cuvsAceParams {
   bool use_disk;
   double max_host_memory_gb;
   double max_gpu_memory_gb;
+  double target_alpha;
 };
 ```
 
@@ -99,6 +100,7 @@ struct cuvsAceParams {
 | `use_disk` | `bool` | Whether to use disk-based storage for ACE build.<br /><br />When true, enables disk-based operations for memory-efficient graph construction. |
 | `max_host_memory_gb` | `double` | Maximum host memory to use for ACE build in GiB.<br /><br />When set to 0 (default), uses available host memory. When set to a positive value, limits host memory usage to the specified amount. Useful for testing or when running alongside other memory-intensive processes. |
 | `max_gpu_memory_gb` | `double` | Maximum GPU memory to use for ACE build in GiB.<br /><br />When set to 0 (default), uses available GPU memory. When set to a positive value, limits GPU memory usage to the specified amount. Useful for testing or when running alongside other memory-intensive processes. |
+| `target_alpha` | `double` | Target number of augmented (spill) copies per input vector. A value of 0 disables augmentation. The default is 1.0. |
 
 <a id="cuvscagraindexparams"></a>
 ### cuvsCagraIndexParams
diff --git a/fern/pages/c_api/c-api-neighbors-hnsw.md b/fern/pages/c_api/c-api-neighbors-hnsw.md
index 4d46dbfa0e..6de8ee4a4e 100644
--- a/fern/pages/c_api/c-api-neighbors-hnsw.md
+++ b/fern/pages/c_api/c-api-neighbors-hnsw.md
@@ -47,6 +47,7 @@ struct cuvsHnswAceParams {
   bool use_disk;
   double max_host_memory_gb;
   double max_gpu_memory_gb;
+  double target_alpha;
 };
 ```
 
@@ -59,6 +60,7 @@ struct cuvsHnswAceParams {
 | `use_disk` | `bool` | Whether to use disk-based storage for ACE build. When true, enables disk-based operations for memory-efficient graph construction. |
 | `max_host_memory_gb` | `double` | Maximum host memory to use for ACE build in GiB. When set to 0 (default), uses available host memory. Useful for testing or when running alongside other memory-intensive processes. |
 | `max_gpu_memory_gb` | `double` | Maximum GPU memory to use for ACE build in GiB. When set to 0 (default), uses available GPU memory. Useful for testing or when running alongside other memory-intensive processes. |
+| `target_alpha` | `double` | Target number of augmented (spill) copies per input vector. A value of 0 disables augmentation. The default is 1.0. |
 
 <a id="cuvshnswaceparamscreate"></a>
 ### cuvsHnswAceParamsCreate
diff --git a/fern/pages/cpp_api/cpp-api-neighbors-cagra.md b/fern/pages/cpp_api/cpp-api-neighbors-cagra.md
index 2560f59ef1..f49a64f0fb 100644
--- a/fern/pages/cpp_api/cpp-api-neighbors-cagra.md
+++ b/fern/pages/cpp_api/cpp-api-neighbors-cagra.md
@@ -208,6 +208,7 @@ struct ace_params {
   bool use_disk;
   double max_host_memory_gb;
   double max_gpu_memory_gb;
+  double target_alpha;
 };
 ```
 
@@ -221,6 +222,7 @@ struct ace_params {
 | `use_disk` | `bool` | Whether to use disk-based storage for ACE build.<br /><br />When true, enables disk-based operations for memory-efficient graph construction. |
 | `max_host_memory_gb` | `double` | Maximum host memory to use for ACE build in GiB.<br /><br />When set to 0 (default), uses available host memory. When set to a positive value, limits host memory usage to the specified amount. Useful for testing or when running alongside other memory-intensive processes. |
 | `max_gpu_memory_gb` | `double` | Maximum GPU memory to use for ACE build in GiB.<br /><br />When set to 0 (default), uses available GPU memory. When set to a positive value, limits GPU memory usage to the specified amount. Useful for testing or when running alongside other memory-intensive processes. |
+| `target_alpha` | `double` | Target number of augmented (spill) copies per input vector.<br /><br />ACE samples boundary-margin scores and chooses a global threshold whose expected augmentation ratio matches this value. A value of 0 disables augmentation. Values above 1 allow a vector to spill into multiple nearby partitions. The implementation bounds each vector to three spill partitions to keep memory and build work predictable. |
 
 ## CAGRA index extend parameters
 
diff --git a/fern/pages/python_api/python-api-neighbors-cagra.md b/fern/pages/python_api/python-api-neighbors-cagra.md
index 9b1a3e3924..a6bed0d22d 100644
--- a/fern/pages/python_api/python-api-neighbors-cagra.md
+++ b/fern/pages/python_api/python-api-neighbors-cagra.md
@@ -28,11 +28,12 @@ for each partition independently.
 | `use_disk` | `bool, default = False` | Whether to use disk-based storage for ACE build. When true, enables disk-based operations for memory-efficient graph construction. |
 | `max_host_memory_gb` | `float, default = 0` | Maximum host memory to use for ACE build in GiB. When set to 0 (default), uses available host memory. Useful for testing or when running alongside other memory-intensive processes. |
 | `max_gpu_memory_gb` | `float, default = 0` | Maximum GPU memory to use for ACE build in GiB. When set to 0 (default), uses available GPU memory. Useful for testing or when running alongside other memory-intensive processes. |
+| `target_alpha` | `float, default = 1.0` | Target number of augmented copies per input vector. A value of 0 disables augmentation; values above 1 allow multi-partition spills. |
 
 **Constructor**
 
 ```python
-def __init__(self, *, npartitions=0, ef_construction=120, build_dir="/tmp/ace_build", use_disk=False, max_host_memory_gb=0, max_gpu_memory_gb=0)
+def __init__(self, *, npartitions=0, ef_construction=120, build_dir="/tmp/ace_build", use_disk=False, max_host_memory_gb=0, max_gpu_memory_gb=0, target_alpha=1.0)
 ```
 
 **Members**
@@ -45,6 +46,7 @@ def __init__(self, *, npartitions=0, ef_construction=120, build_dir="/tmp/ace_bu
 | `use_disk` | property |
 | `max_host_memory_gb` | property |
 | `max_gpu_memory_gb` | property |
+| `target_alpha` | property |
 | `get_handle` | method |
 
 ### npartitions
@@ -83,6 +85,12 @@ def max_host_memory_gb(self)
 def max_gpu_memory_gb(self)
 ```
 
+### target_alpha
+
+```python
+def target_alpha(self)
+```
+
 ### get_handle
 
 ```python
diff --git a/fern/pages/python_api/python-api-neighbors-hnsw.md b/fern/pages/python_api/python-api-neighbors-hnsw.md
index d4578b81ad..5573c07a1f 100644
--- a/fern/pages/python_api/python-api-neighbors-hnsw.md
+++ b/fern/pages/python_api/python-api-neighbors-hnsw.md
@@ -27,11 +27,12 @@ partition independently.
 | `use_disk` | `bool, default = False (optional)` | Whether to use disk-based storage for ACE build. When true, enables disk-based operations for memory-efficient graph construction. |
 | `max_host_memory_gb` | `float, default = 0 (optional)` | Maximum host memory to use for ACE build in GiB. When set to 0 (default), uses available host memory. Useful for testing or when running alongside other memory-intensive processes. |
 | `max_gpu_memory_gb` | `float, default = 0 (optional)` | Maximum GPU memory to use for ACE build in GiB. When set to 0 (default), uses available GPU memory. Useful for testing or when running alongside other memory-intensive processes. |
+| `target_alpha` | `float, default = 1.0 (optional)` | Target number of augmented copies per input vector. A value of 0 disables augmentation; values above 1 allow multi-partition spills. |
 
 **Constructor**
 
 ```python
-def __init__(self, *, npartitions=0, build_dir="/tmp/hnsw_ace_build", use_disk=False, max_host_memory_gb=0, max_gpu_memory_gb=0)
+def __init__(self, *, npartitions=0, build_dir="/tmp/hnsw_ace_build", use_disk=False, max_host_memory_gb=0, max_gpu_memory_gb=0, target_alpha=1.0)
 ```
 
 **Members**
@@ -43,6 +44,7 @@ def __init__(self, *, npartitions=0, build_dir="/tmp/hnsw_ace_build", use_disk=F
 | `use_disk` | property |
 | `max_host_memory_gb` | property |
 | `max_gpu_memory_gb` | property |
+| `target_alpha` | property |
 
 ### npartitions
 
@@ -74,6 +76,12 @@ def max_host_memory_gb(self)
 def max_gpu_memory_gb(self)
 ```
 
+### target_alpha
+
+```python
+def target_alpha(self)
+```
+
 ## IndexParams
 
 ```python
diff --git a/python/cuvs/cuvs/neighbors/cagra/cagra.pxd b/python/cuvs/cuvs/neighbors/cagra/cagra.pxd
index cfde71855e..70fb361fff 100644
--- a/python/cuvs/cuvs/neighbors/cagra/cagra.pxd
+++ b/python/cuvs/cuvs/neighbors/cagra/cagra.pxd
@@ -70,6 +70,7 @@ cdef extern from "cuvs/neighbors/cagra.h" nogil:
         bool use_disk
         double max_host_memory_gb
         double max_gpu_memory_gb
+        double target_alpha
     ctypedef cuvsAceParams* cuvsAceParams_t
 
     ctypedef struct cuvsCagraIndexParams:
diff --git a/python/cuvs/cuvs/neighbors/cagra/cagra.pyx b/python/cuvs/cuvs/neighbors/cagra/cagra.pyx
index 27880c8dd6..753492afac 100644
--- a/python/cuvs/cuvs/neighbors/cagra/cagra.pyx
+++ b/python/cuvs/cuvs/neighbors/cagra/cagra.pyx
@@ -106,6 +106,9 @@ cdef class AceParams:
         Maximum GPU memory to use for ACE build in GiB. When set to 0
         (default), uses available GPU memory. Useful for testing or
         when running alongside other memory-intensive processes.
+    target_alpha : float, default = 1.0
+        Target number of augmented copies per input vector. A value of 0
+        disables augmentation; values above 1 allow multi-partition spills.
     """
     cdef cuvsAceParams* params
     cdef bytes _build_dir_bytes  # Keep Python bytes alive for property access
@@ -124,12 +127,14 @@ cdef class AceParams:
                  build_dir="/tmp/ace_build",
                  use_disk=False,
                  max_host_memory_gb=0,
-                 max_gpu_memory_gb=0):
+                 max_gpu_memory_gb=0,
+                 target_alpha=1.0):
         self.params.npartitions = npartitions
         self.params.ef_construction = ef_construction
         self.params.use_disk = use_disk
         self.params.max_host_memory_gb = max_host_memory_gb
         self.params.max_gpu_memory_gb = max_gpu_memory_gb
+        self.params.target_alpha = target_alpha
 
         # Need to replace the default build_dir allocated by
         # cuvsAceParamsCreate
@@ -169,6 +174,10 @@ cdef class AceParams:
     def max_gpu_memory_gb(self):
         return self.params.max_gpu_memory_gb
 
+    @property
+    def target_alpha(self):
+        return self.params.target_alpha
+
     def get_handle(self):
         return <size_t>self.params
 
@@ -277,6 +286,7 @@ cdef class IndexParams:
             new_ace_params.use_disk = ace_params.params.use_disk
             new_ace_params.max_host_memory_gb = ace_params.params.max_host_memory_gb
             new_ace_params.max_gpu_memory_gb = ace_params.params.max_gpu_memory_gb
+            new_ace_params.target_alpha = ace_params.params.target_alpha
 
             # Copy the build_dir string
             if new_ace_params.build_dir != NULL:
diff --git a/python/cuvs/cuvs/neighbors/hnsw/hnsw.pxd b/python/cuvs/cuvs/neighbors/hnsw/hnsw.pxd
index 9ffb295ad3..922cc3d702 100644
--- a/python/cuvs/cuvs/neighbors/hnsw/hnsw.pxd
+++ b/python/cuvs/cuvs/neighbors/hnsw/hnsw.pxd
@@ -1,5 +1,5 @@
 #
-# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # cython: language_level=3
@@ -27,6 +27,7 @@ cdef extern from "cuvs/neighbors/hnsw.h" nogil:
         bool use_disk
         double max_host_memory_gb
         double max_gpu_memory_gb
+        double target_alpha
 
     ctypedef cuvsHnswAceParams* cuvsHnswAceParams_t
 
diff --git a/python/cuvs/cuvs/neighbors/hnsw/hnsw.pyx b/python/cuvs/cuvs/neighbors/hnsw/hnsw.pyx
index 91256a0e18..9479f04726 100644
--- a/python/cuvs/cuvs/neighbors/hnsw/hnsw.pyx
+++ b/python/cuvs/cuvs/neighbors/hnsw/hnsw.pyx
@@ -68,6 +68,9 @@ cdef class AceParams:
         Maximum GPU memory to use for ACE build in GiB. When set to 0
         (default), uses available GPU memory. Useful for testing or
         when running alongside other memory-intensive processes.
+    target_alpha : float, default = 1.0 (optional)
+        Target number of augmented copies per input vector. A value of 0
+        disables augmentation; values above 1 allow multi-partition spills.
     """
 
     cdef cuvsHnswAceParams* params
@@ -86,13 +89,15 @@ cdef class AceParams:
                  build_dir="/tmp/hnsw_ace_build",
                  use_disk=False,
                  max_host_memory_gb=0,
-                 max_gpu_memory_gb=0):
+                 max_gpu_memory_gb=0,
+                 target_alpha=1.0):
         self.params.npartitions = npartitions
         self._build_dir_bytes = build_dir.encode('utf-8')
         self.params.build_dir = self._build_dir_bytes
         self.params.use_disk = use_disk
         self.params.max_host_memory_gb = max_host_memory_gb
         self.params.max_gpu_memory_gb = max_gpu_memory_gb
+        self.params.target_alpha = target_alpha
 
     @property
     def npartitions(self):
@@ -116,6 +121,10 @@ cdef class AceParams:
     def max_gpu_memory_gb(self):
         return self.params.max_gpu_memory_gb
 
+    @property
+    def target_alpha(self):
+        return self.params.target_alpha
+
 
 cdef class IndexParams:
     """
diff --git a/python/cuvs/cuvs/tests/test_cagra_ace.py b/python/cuvs/cuvs/tests/test_cagra_ace.py
index dd5f323974..86b9832b38 100644
--- a/python/cuvs/cuvs/tests/test_cagra_ace.py
+++ b/python/cuvs/cuvs/tests/test_cagra_ace.py
@@ -22,6 +22,14 @@
 )
 
 
+def test_ace_target_alpha_params():
+    defaults = cagra.AceParams()
+    configured = cagra.AceParams(target_alpha=1.5)
+
+    assert defaults.target_alpha == 1.0
+    assert configured.target_alpha == 1.5
+
+
 def run_cagra_ace_build_search_test(
     n_rows=10000,
     n_cols=10,
diff --git a/python/cuvs/cuvs/tests/test_hnsw_ace.py b/python/cuvs/cuvs/tests/test_hnsw_ace.py
index 183d530e7c..05d5a41170 100644
--- a/python/cuvs/cuvs/tests/test_hnsw_ace.py
+++ b/python/cuvs/cuvs/tests/test_hnsw_ace.py
@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 
@@ -14,6 +14,14 @@
 from cuvs.tests.ann_utils import calc_recall, generate_data
 
 
+def test_ace_target_alpha_params():
+    defaults = hnsw.AceParams()
+    configured = hnsw.AceParams(target_alpha=1.5)
+
+    assert defaults.target_alpha == 1.0
+    assert configured.target_alpha == 1.5
+
+
 def run_hnsw_ace_build_search_test(
     n_rows=10000,
     n_cols=10,

From 3732b489d57f8bfbb188809ebd351eca4a47f93c Mon Sep 17 00:00:00 2001
From: Julian Miller <julianmiller@nvidia.com>
Date: Fri, 2 Oct 2026 10:06:35 +0200
Subject: [PATCH 02/13] Add missing Java and Rust bindings

---
 .../java-api-com-nvidia-cuvs-cuvsaceparams.md | 63 +++++++++++++++----
 .../java-api-com-nvidia-cuvs-hnswaceparams.md | 61 ++++++++++++++----
 .../java/com/nvidia/cuvs/CuVSAceParams.java   | 50 ++++++++++++++-
 .../java/com/nvidia/cuvs/HnswAceParams.java   | 34 +++++++++-
 .../nvidia/cuvs/internal/CagraIndexImpl.java  |  1 +
 .../nvidia/cuvs/internal/HnswIndexImpl.java   |  1 +
 rust/cuvs-sys/src/bindings.rs                 | 10 ++-
 7 files changed, 189 insertions(+), 31 deletions(-)

diff --git a/fern/pages/java_api/java-api-com-nvidia-cuvs-cuvsaceparams.md b/fern/pages/java_api/java-api-com-nvidia-cuvs-cuvsaceparams.md
index a3dccaefab..f02e9e686c 100644
--- a/fern/pages/java_api/java-api-com-nvidia-cuvs-cuvsaceparams.md
+++ b/fern/pages/java_api/java-api-com-nvidia-cuvs-cuvsaceparams.md
@@ -31,7 +31,7 @@ Gets the number of partitions.
 
 the number of partitions
 
-_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSAceParams.java:97`_
+_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSAceParams.java:109`_
 
 ### getEfConstruction
 
@@ -45,7 +45,7 @@ Gets the `ef_construction` parameter.
 
 the `ef_construction` parameter
 
-_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSAceParams.java:106`_
+_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSAceParams.java:118`_
 
 ### getBuildDir
 
@@ -59,7 +59,7 @@ Gets the build directory path.
 
 the build directory path
 
-_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSAceParams.java:115`_
+_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSAceParams.java:127`_
 
 ### isUseDisk
 
@@ -73,7 +73,7 @@ Gets whether disk-based mode is enabled.
 
 true if disk-based mode is enabled
 
-_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSAceParams.java:124`_
+_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSAceParams.java:136`_
 
 ### getMaxHostMemoryGb
 
@@ -87,7 +87,7 @@ Gets the maximum host memory limit in GiB.
 
 the max host memory limit (0 means use available memory)
 
-_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSAceParams.java:133`_
+_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSAceParams.java:145`_
 
 ### getMaxGpuMemoryGb
 
@@ -101,7 +101,21 @@ Gets the maximum GPU memory limit in GiB.
 
 the max GPU memory limit (0 means use available memory)
 
-_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSAceParams.java:142`_
+_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSAceParams.java:154`_
+
+### getTargetAlpha
+
+```java
+public double getTargetAlpha()
+```
+
+Gets the target number of augmented copies per input vector.
+
+**Returns**
+
+the target augmentation ratio (`0` disables augmentation)
+
+_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSAceParams.java:163`_
 
 ### withNpartitions
 
@@ -121,7 +135,7 @@ Sets the number of partitions.
 
 an instance of Builder
 
-_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSAceParams.java:194`_
+_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSAceParams.java:220`_
 
 ### withEfConstruction
 
@@ -141,7 +155,7 @@ Sets the ef_construction parameter.
 
 an instance of Builder
 
-_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSAceParams.java:205`_
+_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSAceParams.java:231`_
 
 ### withBuildDir
 
@@ -161,7 +175,7 @@ Sets the build directory path.
 
 an instance of Builder
 
-_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSAceParams.java:216`_
+_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSAceParams.java:242`_
 
 ### withUseDisk
 
@@ -181,7 +195,7 @@ Sets whether to use disk-based mode.
 
 an instance of Builder
 
-_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSAceParams.java:227`_
+_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSAceParams.java:253`_
 
 ### withMaxHostMemoryGb
 
@@ -204,7 +218,7 @@ Useful for testing or when running alongside other memory-intensive processes.
 
 an instance of Builder
 
-_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSAceParams.java:241`_
+_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSAceParams.java:267`_
 
 ### withMaxGpuMemoryGb
 
@@ -227,7 +241,30 @@ Useful for testing or when running alongside other memory-intensive processes.
 
 an instance of Builder
 
-_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSAceParams.java:255`_
+_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSAceParams.java:281`_
+
+### withTargetAlpha
+
+```java
+public Builder withTargetAlpha(double targetAlpha)
+```
+
+Sets the target number of augmented (spill) copies per input vector.
+
+A value of 0 disables augmentation. Values above 1 allow a vector to spill into multiple
+nearby partitions. The default is `1.0`.
+
+**Parameters**
+
+| Name | Description |
+| --- | --- |
+| `targetAlpha` | the target augmentation ratio |
+
+**Returns**
+
+an instance of Builder
+
+_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSAceParams.java:295`_
 
 ### build
 
@@ -241,6 +278,6 @@ Builds an instance of `CuVSAceParams`.
 
 an instance of `CuVSAceParams`
 
-_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSAceParams.java:265`_
+_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSAceParams.java:305`_
 
 _Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSAceParams.java:17`_
diff --git a/fern/pages/java_api/java-api-com-nvidia-cuvs-hnswaceparams.md b/fern/pages/java_api/java-api-com-nvidia-cuvs-hnswaceparams.md
index a7c002e06d..6fce368644 100644
--- a/fern/pages/java_api/java-api-com-nvidia-cuvs-hnswaceparams.md
+++ b/fern/pages/java_api/java-api-com-nvidia-cuvs-hnswaceparams.md
@@ -30,7 +30,7 @@ Gets the number of partitions for ACE partitioned build.
 
 the number of partitions
 
-_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/HnswAceParams.java:41`_
+_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/HnswAceParams.java:44`_
 
 ### getBuildDir
 
@@ -44,7 +44,7 @@ Gets the directory to store ACE build artifacts.
 
 the build directory path
 
-_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/HnswAceParams.java:50`_
+_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/HnswAceParams.java:53`_
 
 ### isUseDisk
 
@@ -58,7 +58,7 @@ Gets whether disk-based storage is enabled for ACE build.
 
 true if disk mode is enabled
 
-_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/HnswAceParams.java:59`_
+_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/HnswAceParams.java:62`_
 
 ### getMaxHostMemoryGb
 
@@ -72,7 +72,7 @@ Gets the maximum host memory limit in GiB.
 
 the max host memory limit (0 means use available memory)
 
-_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/HnswAceParams.java:68`_
+_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/HnswAceParams.java:71`_
 
 ### getMaxGpuMemoryGb
 
@@ -86,7 +86,21 @@ Gets the maximum GPU memory limit in GiB.
 
 the max GPU memory limit (0 means use available memory)
 
-_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/HnswAceParams.java:77`_
+_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/HnswAceParams.java:80`_
+
+### getTargetAlpha
+
+```java
+public double getTargetAlpha()
+```
+
+Gets the target number of augmented copies per input vector.
+
+**Returns**
+
+the target augmentation ratio (`0` disables augmentation)
+
+_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/HnswAceParams.java:89`_
 
 ### Builder
 
@@ -96,7 +110,7 @@ public Builder()
 
 Constructs this Builder.
 
-_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/HnswAceParams.java:110`_
+_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/HnswAceParams.java:125`_
 
 ### withNpartitions
 
@@ -129,7 +143,7 @@ constraints and a warning will be issued.
 
 an instance of Builder
 
-_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/HnswAceParams.java:131`_
+_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/HnswAceParams.java:146`_
 
 ### withBuildDir
 
@@ -150,7 +164,7 @@ Used when useDisk is true or when the graph does not fit in memory.
 
 an instance of Builder
 
-_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/HnswAceParams.java:143`_
+_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/HnswAceParams.java:158`_
 
 ### withUseDisk
 
@@ -171,7 +185,7 @@ When true, enables disk-based operations for memory-efficient graph construction
 
 an instance of Builder
 
-_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/HnswAceParams.java:155`_
+_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/HnswAceParams.java:170`_
 
 ### withMaxHostMemoryGb
 
@@ -194,7 +208,7 @@ Useful for testing or when running alongside other memory-intensive processes.
 
 an instance of Builder
 
-_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/HnswAceParams.java:169`_
+_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/HnswAceParams.java:184`_
 
 ### withMaxGpuMemoryGb
 
@@ -217,7 +231,30 @@ Useful for testing or when running alongside other memory-intensive processes.
 
 an instance of Builder
 
-_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/HnswAceParams.java:183`_
+_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/HnswAceParams.java:198`_
+
+### withTargetAlpha
+
+```java
+public Builder withTargetAlpha(double targetAlpha)
+```
+
+Sets the target number of augmented (spill) copies per input vector.
+
+A value of 0 disables augmentation. Values above 1 allow a vector to spill into multiple
+nearby partitions. The default is `1.0`.
+
+**Parameters**
+
+| Name | Description |
+| --- | --- |
+| `targetAlpha` | the target augmentation ratio |
+
+**Returns**
+
+an instance of Builder
+
+_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/HnswAceParams.java:212`_
 
 ### build
 
@@ -231,6 +268,6 @@ Builds an instance of `HnswAceParams`.
 
 an instance of `HnswAceParams`
 
-_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/HnswAceParams.java:193`_
+_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/HnswAceParams.java:222`_
 
 _Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/HnswAceParams.java:16`_
diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSAceParams.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSAceParams.java
index eddb404874..f5e93a953b 100644
--- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSAceParams.java
+++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSAceParams.java
@@ -75,19 +75,31 @@ public class CuVSAceParams {
    */
   private final double maxGpuMemoryGb;
 
+  /**
+   * Target number of augmented (spill) copies per input vector.
+   *
+   * ACE samples boundary-margin scores and chooses a global threshold whose expected augmentation
+   * ratio matches this value. A value of 0 disables augmentation. Values above 1 allow a vector to
+   * spill into multiple nearby partitions. The implementation bounds each vector to three spill
+   * partitions. The default is {@code 1.0}.
+   */
+  private final double targetAlpha;
+
   private CuVSAceParams(
       long npartitions,
       long efConstruction,
       String buildDir,
       boolean useDisk,
       double maxHostMemoryGb,
-      double maxGpuMemoryGb) {
+      double maxGpuMemoryGb,
+      double targetAlpha) {
     this.npartitions = npartitions;
     this.efConstruction = efConstruction;
     this.buildDir = buildDir;
     this.useDisk = useDisk;
     this.maxHostMemoryGb = maxHostMemoryGb;
     this.maxGpuMemoryGb = maxGpuMemoryGb;
+    this.targetAlpha = targetAlpha;
   }
 
   /**
@@ -144,6 +156,15 @@ public double getMaxGpuMemoryGb() {
     return maxGpuMemoryGb;
   }
 
+  /**
+   * Gets the target number of augmented copies per input vector.
+   *
+   * @return the target augmentation ratio ({@code 0} disables augmentation)
+   */
+  public double getTargetAlpha() {
+    return targetAlpha;
+  }
+
   @Override
   public String toString() {
     return "CuVSAceParams [npartitions="
@@ -158,6 +179,8 @@ public String toString() {
         + maxHostMemoryGb
         + ", maxGpuMemoryGb="
         + maxGpuMemoryGb
+        + ", targetAlpha="
+        + targetAlpha
         + "]";
   }
 
@@ -184,6 +207,9 @@ public static class Builder {
     /** Maximum GPU memory in GiB (0 = use available memory) */
     private double maxGpuMemoryGb = 0;
 
+    /** Target augmented copies per input vector (0 disables augmentation) */
+    private double targetAlpha = 1.0;
+
     public Builder() {}
 
     /**
@@ -258,6 +284,20 @@ public Builder withMaxGpuMemoryGb(double maxGpuMemoryGb) {
       return this;
     }
 
+    /**
+     * Sets the target number of augmented (spill) copies per input vector.
+     *
+     * A value of 0 disables augmentation. Values above 1 allow a vector to spill into multiple
+     * nearby partitions. The default is {@code 1.0}.
+     *
+     * @param targetAlpha the target augmentation ratio
+     * @return an instance of Builder
+     */
+    public Builder withTargetAlpha(double targetAlpha) {
+      this.targetAlpha = targetAlpha;
+      return this;
+    }
+
     /**
      * Builds an instance of {@link CuVSAceParams}.
      *
@@ -265,7 +305,13 @@ public Builder withMaxGpuMemoryGb(double maxGpuMemoryGb) {
      */
     public CuVSAceParams build() {
       return new CuVSAceParams(
-          npartitions, efConstruction, buildDir, useDisk, maxHostMemoryGb, maxGpuMemoryGb);
+          npartitions,
+          efConstruction,
+          buildDir,
+          useDisk,
+          maxHostMemoryGb,
+          maxGpuMemoryGb,
+          targetAlpha);
     }
   }
 }
diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/HnswAceParams.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/HnswAceParams.java
index 41370f0ac5..64c06d1180 100644
--- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/HnswAceParams.java
+++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/HnswAceParams.java
@@ -20,18 +20,21 @@ public class HnswAceParams {
   private boolean useDisk;
   private double maxHostMemoryGb;
   private double maxGpuMemoryGb;
+  private double targetAlpha;
 
   private HnswAceParams(
       long npartitions,
       String buildDir,
       boolean useDisk,
       double maxHostMemoryGb,
-      double maxGpuMemoryGb) {
+      double maxGpuMemoryGb,
+      double targetAlpha) {
     this.npartitions = npartitions;
     this.buildDir = buildDir;
     this.useDisk = useDisk;
     this.maxHostMemoryGb = maxHostMemoryGb;
     this.maxGpuMemoryGb = maxGpuMemoryGb;
+    this.targetAlpha = targetAlpha;
   }
 
   /**
@@ -79,6 +82,15 @@ public double getMaxGpuMemoryGb() {
     return maxGpuMemoryGb;
   }
 
+  /**
+   * Gets the target number of augmented copies per input vector.
+   *
+   * @return the target augmentation ratio ({@code 0} disables augmentation)
+   */
+  public double getTargetAlpha() {
+    return targetAlpha;
+  }
+
   @Override
   public String toString() {
     return "HnswAceParams [npartitions="
@@ -91,6 +103,8 @@ public String toString() {
         + maxHostMemoryGb
         + ", maxGpuMemoryGb="
         + maxGpuMemoryGb
+        + ", targetAlpha="
+        + targetAlpha
         + "]";
   }
 
@@ -104,6 +118,7 @@ public static class Builder {
     private boolean useDisk = false;
     private double maxHostMemoryGb = 0;
     private double maxGpuMemoryGb = 0;
+    private double targetAlpha = 1.0;
 
     /**
      * Constructs this Builder.
@@ -186,13 +201,28 @@ public Builder withMaxGpuMemoryGb(double maxGpuMemoryGb) {
       return this;
     }
 
+    /**
+     * Sets the target number of augmented (spill) copies per input vector.
+     *
+     * A value of 0 disables augmentation. Values above 1 allow a vector to spill into multiple
+     * nearby partitions. The default is {@code 1.0}.
+     *
+     * @param targetAlpha the target augmentation ratio
+     * @return an instance of Builder
+     */
+    public Builder withTargetAlpha(double targetAlpha) {
+      this.targetAlpha = targetAlpha;
+      return this;
+    }
+
     /**
      * Builds an instance of {@link HnswAceParams}.
      *
      * @return an instance of {@link HnswAceParams}
      */
     public HnswAceParams build() {
-      return new HnswAceParams(npartitions, buildDir, useDisk, maxHostMemoryGb, maxGpuMemoryGb);
+      return new HnswAceParams(
+          npartitions, buildDir, useDisk, maxHostMemoryGb, maxGpuMemoryGb, targetAlpha);
     }
   }
 }
diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java
index 491c5213fc..3037767b3c 100644
--- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java
+++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java
@@ -1145,6 +1145,7 @@ private static void populateNativeIndexParams(
           cuvsAceParamsMemorySegment, cuVSAceParams.getMaxHostMemoryGb());
       cuvsAceParams.max_gpu_memory_gb(
           cuvsAceParamsMemorySegment, cuVSAceParams.getMaxGpuMemoryGb());
+      cuvsAceParams.target_alpha(cuvsAceParamsMemorySegment, cuVSAceParams.getTargetAlpha());
 
       String buildDir = cuVSAceParams.getBuildDir();
       if (buildDir != null && !buildDir.isEmpty()) {
diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/HnswIndexImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/HnswIndexImpl.java
index 6f32e8deab..64e8317eb1 100644
--- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/HnswIndexImpl.java
+++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/HnswIndexImpl.java
@@ -313,6 +313,7 @@ private static CloseableHandle createHnswAceParams(Arena arena, HnswAceParams ac
     cuvsHnswAceParams.use_disk(seg, aceParams.isUseDisk());
     cuvsHnswAceParams.max_host_memory_gb(seg, aceParams.getMaxHostMemoryGb());
     cuvsHnswAceParams.max_gpu_memory_gb(seg, aceParams.getMaxGpuMemoryGb());
+    cuvsHnswAceParams.target_alpha(seg, aceParams.getTargetAlpha());
 
     String buildDir = aceParams.getBuildDir();
     if (buildDir != null) {
diff --git a/rust/cuvs-sys/src/bindings.rs b/rust/cuvs-sys/src/bindings.rs
index e723abaaea..b35cc3a56a 100644
--- a/rust/cuvs-sys/src/bindings.rs
+++ b/rust/cuvs-sys/src/bindings.rs
@@ -1099,10 +1099,11 @@ pub struct cuvsAceParams {
     pub use_disk: bool,
     pub max_host_memory_gb: f64,
     pub max_gpu_memory_gb: f64,
+    pub target_alpha: f64,
 }
 #[allow(clippy::unnecessary_operation, clippy::identity_op)]
 const _: () = {
-    ["Size of cuvsAceParams"][::std::mem::size_of::<cuvsAceParams>() - 48usize];
+    ["Size of cuvsAceParams"][::std::mem::size_of::<cuvsAceParams>() - 56usize];
     ["Alignment of cuvsAceParams"][::std::mem::align_of::<cuvsAceParams>() - 8usize];
     ["Offset of field: cuvsAceParams::npartitions"]
         [::std::mem::offset_of!(cuvsAceParams, npartitions) - 0usize];
@@ -1116,6 +1117,8 @@ const _: () = {
         [::std::mem::offset_of!(cuvsAceParams, max_host_memory_gb) - 32usize];
     ["Offset of field: cuvsAceParams::max_gpu_memory_gb"]
         [::std::mem::offset_of!(cuvsAceParams, max_gpu_memory_gb) - 40usize];
+    ["Offset of field: cuvsAceParams::target_alpha"]
+        [::std::mem::offset_of!(cuvsAceParams, target_alpha) - 48usize];
 };
 pub type cuvsAceParams_t = *mut cuvsAceParams;
 #[repr(C)]
@@ -2004,10 +2007,11 @@ pub struct cuvsHnswAceParams {
     pub use_disk: bool,
     pub max_host_memory_gb: f64,
     pub max_gpu_memory_gb: f64,
+    pub target_alpha: f64,
 }
 #[allow(clippy::unnecessary_operation, clippy::identity_op)]
 const _: () = {
-    ["Size of cuvsHnswAceParams"][::std::mem::size_of::<cuvsHnswAceParams>() - 40usize];
+    ["Size of cuvsHnswAceParams"][::std::mem::size_of::<cuvsHnswAceParams>() - 48usize];
     ["Alignment of cuvsHnswAceParams"][::std::mem::align_of::<cuvsHnswAceParams>() - 8usize];
     ["Offset of field: cuvsHnswAceParams::npartitions"]
         [::std::mem::offset_of!(cuvsHnswAceParams, npartitions) - 0usize];
@@ -2019,6 +2023,8 @@ const _: () = {
         [::std::mem::offset_of!(cuvsHnswAceParams, max_host_memory_gb) - 24usize];
     ["Offset of field: cuvsHnswAceParams::max_gpu_memory_gb"]
         [::std::mem::offset_of!(cuvsHnswAceParams, max_gpu_memory_gb) - 32usize];
+    ["Offset of field: cuvsHnswAceParams::target_alpha"]
+        [::std::mem::offset_of!(cuvsHnswAceParams, target_alpha) - 40usize];
 };
 pub type cuvsHnswAceParams_t = *mut cuvsHnswAceParams;
 unsafe extern "C" {

From 83915d72587c69c7c2d0377fd8959cdf1350bee0 Mon Sep 17 00:00:00 2001
From: Julian Miller <julianmiller@nvidia.com>
Date: Mon, 12 Oct 2026 15:41:58 +0200
Subject: [PATCH 03/13] Reject ACE spill counts that overflow index offsets

---
 .../neighbors/detail/cagra/cagra_build.cuh    | 33 ++++++++++++++-----
 .../ann_cagra/test_float_uint32_t.cu          | 15 +++++++++
 2 files changed, 40 insertions(+), 8 deletions(-)

diff --git a/cpp/src/neighbors/detail/cagra/cagra_build.cuh b/cpp/src/neighbors/detail/cagra/cagra_build.cuh
index 67a0b98e6a..b0295ebbd4 100644
--- a/cpp/src/neighbors/detail/cagra/cagra_build.cuh
+++ b/cpp/src/neighbors/detail/cagra/cagra_build.cuh
@@ -47,6 +47,8 @@
 #include <cerrno>
 #include <chrono>
 #include <cmath>
+#include <cstddef>
+#include <cstdint>
 #include <cstdio>
 #include <cstring>
 #include <exception>
@@ -182,6 +184,13 @@ void check_graph_degree(size_t& intermediate_degree, size_t& graph_degree, size_
   }
 }
 
+/** Check the realized spill count before storing prefix sums or cursors in IdxT. */
+template <typename IdxT>
+constexpr bool ace_spill_count_fits(uint64_t count)
+{
+  return count <= static_cast<uint64_t>(std::numeric_limits<IdxT>::max());
+}
+
 // A small per-vector cap bounds worst-case partition imbalance while target_alpha controls the
 // global augmentation budget. The cap is intentionally internal: the spill-cap sweep showed no
 // consistent quality benefit from exposing larger values.
@@ -550,22 +559,25 @@ void ace_create_forward_and_backward_lists(
   raft::host_vector_view<IdxT, int64_t, raft::row_major> core_partition_offsets,
   raft::host_vector_view<IdxT, int64_t, raft::row_major> augmented_partition_offsets)
 {
-  core_partition_offsets(0)      = 0;
-  augmented_partition_offsets(0) = 0;
-  for (size_t c = 1; c < n_partitions; c++) {
-    core_partition_offsets(c) = core_partition_offsets(c - 1) + partition_histogram(c - 1, 0);
-    augmented_partition_offsets(c) =
-      augmented_partition_offsets(c - 1) + partition_histogram(c - 1, 1);
-  }
-
   const size_t labels_dim = partition_labels.extent(1);
   size_t total_augmented  = 0;
   for (size_t partition = 0; partition < n_partitions; partition++) {
     total_augmented += partition_histogram(partition, 1);
   }
+  RAFT_EXPECTS(ace_spill_count_fits<IdxT>(total_augmented),
+               "ACE: realized spill count exceeds the index type's offset capacity; reduce "
+               "target_alpha or dataset size");
   RAFT_EXPECTS(static_cast<size_t>(augmented_backward_mapping.extent(0)) == total_augmented,
                "augmented_backward_mapping must have one entry per spill copy");
 
+  core_partition_offsets(0)      = 0;
+  augmented_partition_offsets(0) = 0;
+  for (size_t c = 1; c < n_partitions; c++) {
+    core_partition_offsets(c) = core_partition_offsets(c - 1) + partition_histogram(c - 1, 0);
+    augmented_partition_offsets(c) =
+      augmented_partition_offsets(c - 1) + partition_histogram(c - 1, 1);
+  }
+
   if (static_cast<size_t>(core_forward_mapping.extent(0)) == 0) {
     // Memory path: both backward mappings
     RAFT_EXPECTS(static_cast<size_t>(core_backward_mapping.extent(0)) == dataset_size,
@@ -1675,6 +1687,11 @@ auto build_ace(raft::resources const& res, const index_params& params, DatasetVi
     for (size_t partition = 0; partition < n_partitions; partition++) {
       total_augmented_vectors += partition_histogram(partition, 1);
     }
+    // target_alpha is sampled, so its product with dataset_size is not an overflow bound.
+    // Validate the realized count before allocating spill files, mappings, or IdxT offsets.
+    RAFT_EXPECTS(ace_spill_count_fits<IdxT>(total_augmented_vectors),
+                 "ACE: realized spill count exceeds the index type's offset capacity; reduce "
+                 "target_alpha or dataset size");
     RAFT_LOG_INFO("ACE build: realized target_alpha %.3f (requested %.3f, spill_cap=%zu)",
                   static_cast<double>(total_augmented_vectors) / dataset_size,
                   target_alpha,
diff --git a/cpp/tests/neighbors/ann_cagra/test_float_uint32_t.cu b/cpp/tests/neighbors/ann_cagra/test_float_uint32_t.cu
index 5dccce21da..e80789fbb3 100644
--- a/cpp/tests/neighbors/ann_cagra/test_float_uint32_t.cu
+++ b/cpp/tests/neighbors/ann_cagra/test_float_uint32_t.cu
@@ -5,6 +5,7 @@
 
 #include <gtest/gtest.h>
 
+#include "../../../src/neighbors/detail/cagra/cagra_build.cuh"
 #include "../ann_cagra.cuh"
 
 #include <utility>
@@ -154,3 +155,17 @@ TEST(AnnCagraMultiPartition, MixedGraphDegreeRejected)
 }
 
 }  // namespace cuvs::neighbors::cagra
+
+namespace cuvs::neighbors::cagra::detail {
+
+TEST(AceBudget, RejectsSpillOffsetOverflow)
+{
+  constexpr uint64_t limit = std::numeric_limits<uint32_t>::max();
+  EXPECT_TRUE(ace_spill_count_fits<uint32_t>(0));
+  EXPECT_TRUE(ace_spill_count_fits<uint32_t>(limit));
+  EXPECT_FALSE(ace_spill_count_fits<uint32_t>(limit + 1));
+  // Original IDs fit, but three spill copies per row do not.
+  EXPECT_FALSE(ace_spill_count_fits<uint32_t>(1'500'000'000ULL * 3));
+}
+
+}  // namespace cuvs::neighbors::cagra::detail

From bb828afb72dd4f77b01619856d3d6ef06309bb15 Mon Sep 17 00:00:00 2001
From: Julian Miller <julianmiller@nvidia.com>
Date: Mon, 12 Oct 2026 15:43:17 +0200
Subject: [PATCH 04/13] Subsample tied ACE margin scores to respect the target
 budget

---
 .../neighbors/detail/cagra/cagra_build.cuh    | 71 +++++++++++++------
 .../ann_cagra/test_float_uint32_t.cu          | 36 ++++++++++
 2 files changed, 87 insertions(+), 20 deletions(-)

diff --git a/cpp/src/neighbors/detail/cagra/cagra_build.cuh b/cpp/src/neighbors/detail/cagra/cagra_build.cuh
index b0295ebbd4..60984c0af8 100644
--- a/cpp/src/neighbors/detail/cagra/cagra_build.cuh
+++ b/cpp/src/neighbors/detail/cagra/cagra_build.cuh
@@ -191,6 +191,51 @@ constexpr bool ace_spill_count_fits(uint64_t count)
   return count <= static_cast<uint64_t>(std::numeric_limits<IdxT>::max());
 }
 
+/** A sampled quantile with deterministic subsampling of scores tied at the boundary. */
+struct ace_margin_gate {
+  double threshold       = -std::numeric_limits<double>::infinity();
+  double tie_probability = 0.0;
+
+  /** Select ties independently of labeling chunk size and OpenMP scheduling. */
+  bool accepts(float margin, size_t row, size_t partition) const
+  {
+    if (margin < threshold) { return true; }
+    if (margin != threshold || tie_probability <= 0.0) { return false; }
+    if (tie_probability >= 1.0) { return true; }
+    // SplitMix64 gives a stable uniform variate for each (row, spill partition).
+    uint64_t key = (static_cast<uint64_t>(row) * 0x9e3779b97f4a7c15ULL) ^ partition;
+    key          = (key ^ (key >> 30)) * 0xbf58476d1ce4e5b9ULL;
+    key          = (key ^ (key >> 27)) * 0x94d049bb133111ebULL;
+    key ^= key >> 31;
+    const double uniform = static_cast<double>(key >> 11) * 0x1.0p-53;
+    return uniform < tie_probability;
+  }
+};
+
+/** Estimate a spill gate, retaining only the required fraction of quantile ties. */
+inline ace_margin_gate ace_target_alpha_threshold(std::vector<float> scores,
+                                                  size_t n_samples,
+                                                  size_t max_spills,
+                                                  double target_alpha)
+{
+  if (scores.empty() || target_alpha <= 0.0) { return {}; }
+  if (target_alpha >= static_cast<double>(max_spills)) {
+    return {std::numeric_limits<double>::infinity(), 1.0};
+  }
+  const size_t keep =
+    std::min(static_cast<size_t>(std::llround(target_alpha * n_samples)), scores.size());
+  if (keep == 0) { return {}; }
+  std::nth_element(scores.begin(), scores.begin() + keep - 1, scores.end());
+  const float threshold = scores[keep - 1];
+  size_t below          = 0;
+  size_t equal          = 0;
+  for (const float score : scores) {
+    below += score < threshold;
+    equal += score == threshold;
+  }
+  return {threshold, static_cast<double>(keep - below) / equal};
+}
+
 // A small per-vector cap bounds worst-case partition imbalance while target_alpha controls the
 // global augmentation budget. The cap is intentionally internal: the spill-cap sweep showed no
 // consistent quality benefit from exposing larger values.
@@ -280,21 +325,6 @@ inline std::vector<float> ace_sample_margin_scores(
   return scores;
 }
 
-inline double ace_target_alpha_threshold(std::vector<float> scores,
-                                         size_t n_samples,
-                                         size_t max_spills,
-                                         double target_alpha)
-{
-  if (scores.empty() || target_alpha <= 0.0) { return -std::numeric_limits<double>::infinity(); }
-  if (target_alpha >= static_cast<double>(max_spills)) {
-    return std::numeric_limits<double>::infinity();
-  }
-  size_t keep = static_cast<size_t>(std::llround(target_alpha * n_samples));
-  keep        = std::clamp<size_t>(keep, 1, scores.size());
-  std::nth_element(scores.begin(), scores.begin() + keep - 1, scores.end());
-  return scores[keep - 1];
-}
-
 // ACE: Get core plus boundary-gated spill partition labels. Only core graph rows are retained;
 // spill vectors enlarge local builds but their graph rows are not merged into the final graph.
 template <typename T, typename IdxT>
@@ -364,8 +394,8 @@ void ace_get_partition_labels(
     res, centroid_distances.view(), raft::make_const_mdspan(centroid_distances_dev.view()));
   raft::resource::sync_stream(res, stream);
 
-  const size_t max_spills       = labels_dim - 1;
-  const double margin_threshold = ace_target_alpha_threshold(
+  const size_t max_spills = labels_dim - 1;
+  const auto margin_gate  = ace_target_alpha_threshold(
     ace_sample_margin_scores(res,
                              raft::make_const_mdspan(sample_db_dev.view()),
                              raft::make_const_mdspan(centroids_dev.view()),
@@ -374,8 +404,9 @@ void ace_get_partition_labels(
     n_samples,
     max_spills,
     target_alpha);
-  RAFT_LOG_INFO(
-    "ACE: target_alpha %.3f -> boundary margin threshold %.6f", target_alpha, margin_threshold);
+  RAFT_LOG_INFO("ACE: target_alpha %.3f -> boundary margin threshold %.6f",
+                target_alpha,
+                margin_gate.threshold);
 
   // Compute distances between dataset and centroid vectors
   // Uses float conversion on host instead of batch_load_iterator to minimize GPU memory usage.
@@ -437,7 +468,7 @@ void ace_get_partition_labels(
             ace_boundary_margin(nearest[0].first,
                                 nearest[rank].first,
                                 centroid_distances(nearest[0].second, nearest[rank].second));
-          if (margin > margin_threshold) { continue; }
+          if (!margin_gate.accepts(margin, output_row, nearest[rank].second)) { continue; }
           partition_labels(output_row, output_slot++) = static_cast<IdxT>(nearest[rank].second);
 #pragma omp atomic update
           partition_histogram(nearest[rank].second, 1) += 1;
diff --git a/cpp/tests/neighbors/ann_cagra/test_float_uint32_t.cu b/cpp/tests/neighbors/ann_cagra/test_float_uint32_t.cu
index e80789fbb3..d99a1373df 100644
--- a/cpp/tests/neighbors/ann_cagra/test_float_uint32_t.cu
+++ b/cpp/tests/neighbors/ann_cagra/test_float_uint32_t.cu
@@ -168,4 +168,40 @@ TEST(AceBudget, RejectsSpillOffsetOverflow)
   EXPECT_FALSE(ace_spill_count_fits<uint32_t>(1'500'000'000ULL * 3));
 }
 
+TEST(AceBudget, SubsamplesFiniteAndInfiniteQuantileTies)
+{
+  for (const float score : {0.0f, std::numeric_limits<float>::infinity()}) {
+    auto gate = ace_target_alpha_threshold(std::vector<float>(300, score), 100, 3, 0.25);
+    EXPECT_EQ(gate.threshold, score);
+    EXPECT_DOUBLE_EQ(gate.tie_probability, 1.0 / 12.0);
+    size_t accepted = 0;
+    for (size_t row = 0; row < 100'000; ++row) {
+      for (size_t partition = 1; partition <= 3; ++partition) {
+        accepted += gate.accepts(score, row, partition);
+      }
+    }
+    EXPECT_NEAR(static_cast<double>(accepted) / 100'000, 0.25, 0.01);
+  }
+}
+
+TEST(AceBudget, AccountsForScoresBelowThreshold)
+{
+  auto gate = ace_target_alpha_threshold({0, 1, 1, 1, 2, 2}, 2, 3, 1.0);
+  EXPECT_EQ(gate.threshold, 1);
+  EXPECT_DOUBLE_EQ(gate.tie_probability, 1.0 / 3.0);
+  EXPECT_TRUE(gate.accepts(0, 5, 1));
+  EXPECT_FALSE(gate.accepts(2, 5, 1));
+}
+
+TEST(AceBudget, ZeroAndMaximumBudgets)
+{
+  const std::vector<float> scores{0, 1, std::numeric_limits<float>::infinity()};
+  auto zero = ace_target_alpha_threshold(scores, 1, 3, 0.0);
+  auto full = ace_target_alpha_threshold(scores, 1, 3, 3.0);
+  for (const float score : scores) {
+    EXPECT_FALSE(zero.accepts(score, 0, 1));
+    EXPECT_TRUE(full.accepts(score, 0, 1));
+  }
+}
+
 }  // namespace cuvs::neighbors::cagra::detail

From ce93d3ef00d03842a0a2cd76f46d4de392dc7ca6 Mon Sep 17 00:00:00 2001
From: Julian Miller <julianmiller@nvidia.com>
Date: Mon, 12 Oct 2026 15:44:00 +0200
Subject: [PATCH 05/13] Skip spill scoring and its partition cap for core-only
 ACE builds

---
 .../neighbors/detail/cagra/cagra_build.cuh    | 69 ++++++++++++-------
 1 file changed, 43 insertions(+), 26 deletions(-)

diff --git a/cpp/src/neighbors/detail/cagra/cagra_build.cuh b/cpp/src/neighbors/detail/cagra/cagra_build.cuh
index 60984c0af8..2ef91d9e91 100644
--- a/cpp/src/neighbors/detail/cagra/cagra_build.cuh
+++ b/cpp/src/neighbors/detail/cagra/cagra_build.cuh
@@ -191,6 +191,12 @@ constexpr bool ace_spill_count_fits(uint64_t count)
   return count <= static_cast<uint64_t>(std::numeric_limits<IdxT>::max());
 }
 
+/** Only augmented labeling needs the quadratic centroid-distance matrix. */
+constexpr bool ace_partition_count_supported(size_t n_partitions, double target_alpha)
+{
+  return target_alpha == 0.0 || n_partitions <= 8192;
+}
+
 /** A sampled quantile with deterministic subsampling of scores tied at the boundary. */
 struct ace_margin_gate {
   double threshold       = -std::numeric_limits<double>::infinity();
@@ -354,7 +360,8 @@ void ace_get_partition_labels(
                  target_alpha <= static_cast<double>(labels_dim - 1),
                "ACE target_alpha must be finite and in [0, %zu]",
                labels_dim - 1);
-  RAFT_EXPECTS(n_partitions <= 8192, "ACE target_alpha currently supports at most 8192 partitions");
+  RAFT_EXPECTS(ace_partition_count_supported(n_partitions, target_alpha),
+               "ACE positive target_alpha currently supports at most 8192 partitions");
   cudaStream_t stream = raft::resource::get_cuda_stream(res).get();
 
   // Sampling vectors from dataset. Uses float conversion on host instead of
@@ -382,28 +389,33 @@ void ace_get_partition_labels(
   auto centroids_dev = raft::make_device_matrix<float, int64_t>(res, n_partitions, dataset_dim);
   cuvs::cluster::kmeans::fit(res, kmeans_params, sample_db_dev.view(), centroids_dev.view());
 
-  auto centroid_distances = raft::make_host_matrix<float, int64_t>(n_partitions, n_partitions);
-  auto centroid_distances_dev =
-    raft::make_device_matrix<float, int64_t>(res, n_partitions, n_partitions);
-  cuvs::distance::pairwise_distance(res,
-                                    raft::make_const_mdspan(centroids_dev.view()),
-                                    raft::make_const_mdspan(centroids_dev.view()),
-                                    centroid_distances_dev.view(),
-                                    cuvs::distance::DistanceType::L2Expanded);
-  raft::copy(
-    res, centroid_distances.view(), raft::make_const_mdspan(centroid_distances_dev.view()));
-  raft::resource::sync_stream(res, stream);
-
-  const size_t max_spills = labels_dim - 1;
-  const auto margin_gate  = ace_target_alpha_threshold(
-    ace_sample_margin_scores(res,
-                             raft::make_const_mdspan(sample_db_dev.view()),
-                             raft::make_const_mdspan(centroids_dev.view()),
-                             raft::make_const_mdspan(centroid_distances.view()),
-                             max_spills),
-    n_samples,
-    max_spills,
-    target_alpha);
+  // Core-only builds do not need centroid-pair distances or sampled spill scores.
+  const size_t max_spills             = labels_dim - 1;
+  const size_t centroid_distance_rows = target_alpha > 0.0 ? n_partitions : 0;
+  auto centroid_distances =
+    raft::make_host_matrix<float, int64_t>(centroid_distance_rows, centroid_distance_rows);
+  ace_margin_gate margin_gate;
+  if (target_alpha > 0.0) {
+    auto centroid_distances_dev =
+      raft::make_device_matrix<float, int64_t>(res, n_partitions, n_partitions);
+    cuvs::distance::pairwise_distance(res,
+                                      raft::make_const_mdspan(centroids_dev.view()),
+                                      raft::make_const_mdspan(centroids_dev.view()),
+                                      centroid_distances_dev.view(),
+                                      cuvs::distance::DistanceType::L2Expanded);
+    raft::copy(
+      res, centroid_distances.view(), raft::make_const_mdspan(centroid_distances_dev.view()));
+    raft::resource::sync_stream(res, stream);
+    margin_gate = ace_target_alpha_threshold(
+      ace_sample_margin_scores(res,
+                               raft::make_const_mdspan(sample_db_dev.view()),
+                               raft::make_const_mdspan(centroids_dev.view()),
+                               raft::make_const_mdspan(centroid_distances.view()),
+                               max_spills),
+      n_samples,
+      max_spills,
+      target_alpha);
+  }
   RAFT_LOG_INFO("ACE: target_alpha %.3f -> boundary margin threshold %.6f",
                 target_alpha,
                 margin_gate.threshold);
@@ -455,9 +467,11 @@ void ace_get_partition_labels(
       std::vector<std::pair<float, size_t>> nearest(labels_dim);
 #pragma omp for
       for (size_t i_sub = 0; i_sub < sub_dataset_size; i_sub++) {
-        const size_t output_row = i_base + i_sub;
-        const size_t found      = ace_select_nearest_partitions(
-          &sub_distances(i_sub, 0), n_partitions, labels_dim, nearest.data());
+        const size_t output_row         = i_base + i_sub;
+        const size_t found              = ace_select_nearest_partitions(&sub_distances(i_sub, 0),
+                                                           n_partitions,
+                                                           target_alpha > 0.0 ? labels_dim : 1,
+                                                           nearest.data());
         partition_labels(output_row, 0) = static_cast<IdxT>(nearest[0].second);
 #pragma omp atomic update
         partition_histogram(nearest[0].second, 0) += 1;
@@ -1644,6 +1658,9 @@ auto build_ace(raft::resources const& res, const index_params& params, DatasetVi
                  max_spills,
                  n_partitions);
 
+    RAFT_EXPECTS(ace_partition_count_supported(n_partitions, target_alpha),
+                 "ACE positive target_alpha currently supports at most 8192 partitions");
+
     // Preallocate space for files for better performance and fail early if not enough space.
     cuvs::util::file_descriptor reordered_fd;
     cuvs::util::file_descriptor augmented_fd;

From ffe0d2ed43389cf8c0514d7e897f9a2e26e910cd Mon Sep 17 00:00:00 2001
From: Julian Miller <julianmiller@nvidia.com>
Date: Mon, 12 Oct 2026 15:45:34 +0200
Subject: [PATCH 06/13] Budget ACE labeling GPU buffers before creating disk
 artifacts

---
 .../neighbors/detail/cagra/cagra_build.cuh    | 146 ++++++++++++++----
 .../ann_cagra/test_float_uint32_t.cu          |  52 +++++++
 python/cuvs/cuvs/tests/test_cagra_ace.py      |   2 +-
 python/cuvs/cuvs/tests/test_hnsw_ace.py       |   2 +-
 4 files changed, 173 insertions(+), 29 deletions(-)

diff --git a/cpp/src/neighbors/detail/cagra/cagra_build.cuh b/cpp/src/neighbors/detail/cagra/cagra_build.cuh
index 2ef91d9e91..64a2a6e3a9 100644
--- a/cpp/src/neighbors/detail/cagra/cagra_build.cuh
+++ b/cpp/src/neighbors/detail/cagra/cagra_build.cuh
@@ -191,10 +191,57 @@ constexpr bool ace_spill_count_fits(uint64_t count)
   return count <= static_cast<uint64_t>(std::numeric_limits<IdxT>::max());
 }
 
-/** Only augmented labeling needs the quadratic centroid-distance matrix. */
-constexpr bool ace_partition_count_supported(size_t n_partitions, double target_alpha)
+/** Match the sample size used for centroid training and margin calibration. */
+inline size_t ace_labeling_sample_count(size_t rows, size_t partitions, double rate = 0.01)
 {
-  return target_alpha == 0.0 || n_partitions <= 8192;
+  const size_t minimum = partitions > rows / 100 ? rows : 100 * partitions;
+  return std::min(rows, std::max(static_cast<size_t>(rows * rate), minimum));
+}
+
+/** Peak explicit device buffers in labeling, excluding the resource workspace. */
+inline size_t ace_labeling_gpu_bytes(
+  size_t rows, size_t dim, size_t partitions, double target_alpha, size_t batch_rows = 32768)
+{
+  const long double samples   = ace_labeling_sample_count(rows, partitions);
+  const long double centroids = partitions;
+  // Sample and centroid matrices survive through assignment. The centroid-pair device
+  // matrix is released before score sampling; assignment uses the admitted batch size.
+  const long double persistent     = (samples + centroids) * dim;
+  const long double centroid_pairs = target_alpha > 0.0 ? centroids * centroids : 0.0L;
+  const long double assignment     = static_cast<long double>(batch_rows) * (dim + centroids);
+  const long double bytes = sizeof(float) * (persistent + std::max(centroid_pairs, assignment));
+  // Saturate impossible requests instead of wrapping the admission estimate.
+  return bytes >= static_cast<long double>(std::numeric_limits<size_t>::max())
+           ? std::numeric_limits<size_t>::max()
+           : static_cast<size_t>(bytes);
+}
+
+/** Choose a labeling batch that fits both budgets; zero means the fixed buffers cannot fit. */
+inline size_t ace_labeling_batch_size(size_t rows,
+                                      size_t dim,
+                                      size_t partitions,
+                                      double target_alpha,
+                                      size_t gpu_budget,
+                                      size_t host_budget)
+{
+  if (rows == 0 || partitions == 0) { return 0; }
+  const long double samples   = ace_labeling_sample_count(rows, partitions);
+  const long double centroids = partitions;
+  const long double pairs     = target_alpha > 0.0 ? centroids * centroids : 0.0L;
+  const long double gpu_fixed = sizeof(float) * (samples + centroids) * dim;
+  // The host retains centroid-pair distances during assignment. Reserve both the final
+  // score vector and thread-local scores during calibration (at most three spills per row).
+  const long double scores =
+    target_alpha > 0.0 ? 2.0L * samples * std::min<size_t>(3, partitions - 1) : 0.0L;
+  const long double host_fixed = sizeof(float) * (samples * dim + pairs + scores);
+  const long double per_row    = sizeof(float) * (static_cast<long double>(dim) + centroids);
+  if (per_row == 0 || gpu_fixed + sizeof(float) * pairs > gpu_budget ||
+      gpu_fixed + per_row > gpu_budget || host_fixed + per_row > host_budget) {
+    return 0;
+  }
+  return static_cast<size_t>(std::min({static_cast<long double>(std::min<size_t>(rows, 32768)),
+                                       std::floor((gpu_budget - gpu_fixed) / per_row),
+                                       std::floor((host_budget - host_fixed) / per_row)}));
 }
 
 /** A sampled quantile with deterministic subsampling of scores tied at the boundary. */
@@ -283,12 +330,13 @@ inline std::vector<float> ace_sample_margin_scores(
   raft::device_matrix_view<const float, int64_t> sample,
   raft::device_matrix_view<const float, int64_t> centroids,
   raft::host_matrix_view<const float, int64_t> centroid_distances,
-  size_t max_spills)
+  size_t max_spills,
+  size_t batch_rows)
 {
   const size_t n_samples    = sample.extent(0);
   const size_t dataset_dim  = sample.extent(1);
   const size_t n_partitions = centroids.extent(0);
-  const size_t chunk_size   = std::min<size_t>(n_samples, 32 * 1024);
+  const size_t chunk_size   = std::min(n_samples, batch_rows);
   auto distances_host       = raft::make_host_matrix<float, int64_t>(chunk_size, n_partitions);
   auto distances_device = raft::make_device_matrix<float, int64_t>(res, chunk_size, n_partitions);
   auto stream           = raft::resource::get_cuda_stream(res).get();
@@ -341,8 +389,10 @@ void ace_get_partition_labels(
   raft::host_matrix_view<IdxT, int64_t, raft::row_major> partition_labels,
   raft::host_matrix_view<IdxT, int64_t, raft::row_major> partition_histogram,
   double target_alpha,
+  size_t chunk_size,
   double sampling_rate = 0.01)
 {
+  RAFT_EXPECTS(chunk_size > 0, "ACE labeling batch must contain at least one row");
   size_t dataset_size = dataset.extent(0);
   size_t labels_size  = partition_labels.extent(0);
   size_t labels_dim   = partition_labels.extent(1);
@@ -360,18 +410,13 @@ void ace_get_partition_labels(
                  target_alpha <= static_cast<double>(labels_dim - 1),
                "ACE target_alpha must be finite and in [0, %zu]",
                labels_dim - 1);
-  RAFT_EXPECTS(ace_partition_count_supported(n_partitions, target_alpha),
-               "ACE positive target_alpha currently supports at most 8192 partitions");
   cudaStream_t stream = raft::resource::get_cuda_stream(res).get();
 
   // Sampling vectors from dataset. Uses float conversion on host instead of
   // raft::matrix::sample_rows to minimize GPU memory usage.
   // TODO(julianmi): Switch to sample_rows when https://github.com/nvidia/cuvs/issues/1461 is
   // addressed.
-  size_t n_samples         = dataset_size * sampling_rate;
-  const size_t min_samples = 100 * n_partitions;
-  n_samples                = std::max(n_samples, min_samples);
-  n_samples                = std::min(n_samples, dataset_size);
+  const size_t n_samples = ace_labeling_sample_count(dataset_size, n_partitions, sampling_rate);
   RAFT_LOG_DEBUG("ACE build: partition labeling uses %lu sampled vectors", n_samples);
 
   auto sample_db = raft::make_host_matrix<float, int64_t>(n_samples, dataset_dim);
@@ -396,22 +441,25 @@ void ace_get_partition_labels(
     raft::make_host_matrix<float, int64_t>(centroid_distance_rows, centroid_distance_rows);
   ace_margin_gate margin_gate;
   if (target_alpha > 0.0) {
-    auto centroid_distances_dev =
-      raft::make_device_matrix<float, int64_t>(res, n_partitions, n_partitions);
-    cuvs::distance::pairwise_distance(res,
-                                      raft::make_const_mdspan(centroids_dev.view()),
-                                      raft::make_const_mdspan(centroids_dev.view()),
-                                      centroid_distances_dev.view(),
-                                      cuvs::distance::DistanceType::L2Expanded);
-    raft::copy(
-      res, centroid_distances.view(), raft::make_const_mdspan(centroid_distances_dev.view()));
-    raft::resource::sync_stream(res, stream);
+    {
+      auto centroid_distances_dev =
+        raft::make_device_matrix<float, int64_t>(res, n_partitions, n_partitions);
+      cuvs::distance::pairwise_distance(res,
+                                        raft::make_const_mdspan(centroids_dev.view()),
+                                        raft::make_const_mdspan(centroids_dev.view()),
+                                        centroid_distances_dev.view(),
+                                        cuvs::distance::DistanceType::L2Expanded);
+      raft::copy(
+        res, centroid_distances.view(), raft::make_const_mdspan(centroid_distances_dev.view()));
+      raft::resource::sync_stream(res, stream);
+    }
     margin_gate = ace_target_alpha_threshold(
       ace_sample_margin_scores(res,
                                raft::make_const_mdspan(sample_db_dev.view()),
                                raft::make_const_mdspan(centroids_dev.view()),
                                raft::make_const_mdspan(centroid_distances.view()),
-                               max_spills),
+                               max_spills,
+                               chunk_size),
       n_samples,
       max_spills,
       target_alpha);
@@ -422,7 +470,6 @@ void ace_get_partition_labels(
 
   // Compute distances between dataset and centroid vectors
   // Uses float conversion on host instead of batch_load_iterator to minimize GPU memory usage.
-  const size_t chunk_size = 32 * 1024;
   auto _sub_dataset       = raft::make_host_matrix<float, int64_t>(chunk_size, dataset_dim);
   auto _sub_distances     = raft::make_host_matrix<float, int64_t>(chunk_size, n_partitions);
   auto _sub_dataset_dev   = raft::make_device_matrix<float, int64_t>(res, chunk_size, dataset_dim);
@@ -1182,6 +1229,7 @@ struct ace_memory_requirements {
   size_t total_size;
   size_t available_host_memory;
   size_t available_gpu_memory;
+  size_t labeling_batch_size;
 };
 
 // Amount of host memory that can be used for the build
@@ -1201,6 +1249,38 @@ inline double ace_estimated_vector_expansion_factor(double target_alpha)
   return 1.0 + estimated_spills;
 }
 
+/** Admit labeling before creating artifacts, reserving the concurrent host labels and GPU
+ * workspace. */
+inline size_t ace_check_labeling_memory(raft::resources const& res,
+                                        size_t dataset_size,
+                                        size_t dataset_dim,
+                                        size_t n_partitions,
+                                        double target_alpha,
+                                        size_t index_size,
+                                        ace_memory_requirements const& mem)
+{
+  const size_t workspace        = raft::resource::get_workspace_total_bytes(res);
+  const long double host_budget = usable_cpu_memory_fraction * mem.available_host_memory;
+  const long double host_labels =
+    static_cast<long double>(mem.partition_labels_size) + 2.0L * n_partitions * index_size;
+  const size_t batch = ace_labeling_batch_size(
+    dataset_size,
+    dataset_dim,
+    n_partitions,
+    target_alpha,
+    workspace < mem.available_gpu_memory ? mem.available_gpu_memory - workspace : 0,
+    host_labels < host_budget ? static_cast<size_t>(host_budget - host_labels) : 0);
+  RAFT_EXPECTS(batch > 0,
+               "ACE: memory insufficient for partition labeling (%zu partitions), including "
+               "centroid-pair distances and one batch row: host %.2f GiB, GPU %.2f GiB "
+               "available. Reduce partitions or increase the memory budgets",
+               n_partitions,
+               to_gib(mem.available_host_memory),
+               to_gib(mem.available_gpu_memory));
+  RAFT_LOG_DEBUG("ACE: partition labeling batch size %zu", batch);
+  return batch;
+}
+
 // Check if disk mode should be used for ACE based on memory constraints
 template <typename T, typename IdxT>
 bool ace_check_use_disk_mode(raft::resources const& res,
@@ -1212,6 +1292,7 @@ bool ace_check_use_disk_mode(raft::resources const& res,
                              size_t intermediate_degree,
                              size_t graph_degree,
                              double vector_expansion_factor,
+                             double target_alpha,
                              std::optional<double> max_host_memory_gb,
                              std::optional<double> max_gpu_memory_gb,
                              bool guarantee_connectivity,
@@ -1319,6 +1400,14 @@ bool ace_check_use_disk_mode(raft::resources const& res,
   size_t gpu_memory_required =
     std::max(mem.sub_dataset_size, opt_dev_ws_total) + extra_gpu_workspace_size;
 
+  mem.labeling_batch_size = ace_check_labeling_memory(
+    res, dataset_size, dataset_dim, n_partitions, target_alpha, sizeof(IdxT), mem);
+  const size_t labeling_gpu_required =
+    ace_labeling_gpu_bytes(
+      dataset_size, dataset_dim, n_partitions, target_alpha, mem.labeling_batch_size) +
+    extra_gpu_workspace_size;
+  gpu_memory_required = std::max(gpu_memory_required, labeling_gpu_required);
+
   bool gpu_memory_limited = mem.available_gpu_memory < gpu_memory_required;
 
   RAFT_LOG_INFO("ACE: Estimated GPU memory required: %.2f GiB, available: %.2f GiB",
@@ -1380,6 +1469,7 @@ void ace_validate_disk_mode_partitions(raft::resources const& res,
                                        size_t intermediate_degree,
                                        size_t graph_degree,
                                        double vector_expansion_factor,
+                                       double target_alpha,
                                        bool guarantee_connectivity,
                                        ace_memory_requirements& mem)
 {
@@ -1520,6 +1610,8 @@ void ace_validate_disk_mode_partitions(raft::resources const& res,
       to_gib(new_opt_host_ws),
       to_gib(new_opt_dev_ws + extra_gpu_workspace_size));
   }
+  mem.labeling_batch_size = ace_check_labeling_memory(
+    res, dataset_size, dataset_dim, n_partitions, target_alpha, sizeof(IdxT), mem);
 }
 
 template <typename T, typename IdxT, typename DatasetViewT>
@@ -1632,6 +1724,7 @@ auto build_ace(raft::resources const& res, const index_params& params, DatasetVi
                                                           intermediate_degree,
                                                           graph_degree,
                                                           vector_expansion_factor,
+                                                          target_alpha,
                                                           ace_params.max_host_memory_gb,
                                                           ace_params.max_gpu_memory_gb,
                                                           params.guarantee_connectivity,
@@ -1646,6 +1739,7 @@ auto build_ace(raft::resources const& res, const index_params& params, DatasetVi
                                                  intermediate_degree,
                                                  graph_degree,
                                                  vector_expansion_factor,
+                                                 target_alpha,
                                                  params.guarantee_connectivity,
                                                  mem);
       ace_validate_partition_count(n_partitions, dataset_size, true);
@@ -1658,9 +1752,6 @@ auto build_ace(raft::resources const& res, const index_params& params, DatasetVi
                  max_spills,
                  n_partitions);
 
-    RAFT_EXPECTS(ace_partition_count_supported(n_partitions, target_alpha),
-                 "ACE positive target_alpha currently supports at most 8192 partitions");
-
     // Preallocate space for files for better performance and fail early if not enough space.
     cuvs::util::file_descriptor reordered_fd;
     cuvs::util::file_descriptor augmented_fd;
@@ -1715,7 +1806,8 @@ auto build_ace(raft::resources const& res, const index_params& params, DatasetVi
                                       dataset_dim,
                                       partition_labels.view(),
                                       partition_histogram.view(),
-                                      target_alpha);
+                                      target_alpha,
+                                      mem.labeling_batch_size);
 
     ace_check_partition_sizes<IdxT>(dataset_size,
                                     n_partitions,
diff --git a/cpp/tests/neighbors/ann_cagra/test_float_uint32_t.cu b/cpp/tests/neighbors/ann_cagra/test_float_uint32_t.cu
index d99a1373df..8d07034e54 100644
--- a/cpp/tests/neighbors/ann_cagra/test_float_uint32_t.cu
+++ b/cpp/tests/neighbors/ann_cagra/test_float_uint32_t.cu
@@ -204,4 +204,56 @@ TEST(AceBudget, ZeroAndMaximumBudgets)
   }
 }
 
+TEST(AceBudget, LabelingMemoryIncludesConcurrentBuffers)
+{
+  constexpr size_t rows       = 1'000'000;
+  constexpr size_t partitions = 8192;
+  constexpr size_t dim        = 128;
+  constexpr size_t expected =
+    sizeof(float) * ((819'200ULL + partitions) * dim + 32768ULL * (dim + partitions));
+  EXPECT_EQ(ace_labeling_gpu_bytes(rows, dim, partitions, 1.0), expected);
+  EXPECT_GT(expected, 1ULL << 30);
+  EXPECT_GT(ace_labeling_gpu_bytes(rows, dim, partitions * 2, 0.0),
+            ace_labeling_gpu_bytes(rows, dim, partitions, 0.0));
+  EXPECT_EQ(ace_labeling_sample_count(1000, 20), 1000);
+}
+
+TEST(AceBudget, LabelingBatchRespectsHostAndGpuBudgetsAbove8192Partitions)
+{
+  constexpr size_t rows = 2'000'000, partitions = 16384, dim = 128;
+  constexpr size_t samples    = 1'638'400;
+  constexpr size_t gpu_fixed  = sizeof(float) * (samples + partitions) * dim;
+  constexpr size_t pairs      = sizeof(float) * partitions * partitions;
+  constexpr size_t host_fixed = sizeof(float) * (samples * dim + 6 * samples) + pairs;
+  constexpr size_t per_row    = sizeof(float) * (dim + partitions);
+  constexpr size_t ample      = 8ULL << 30;
+  EXPECT_EQ(ace_labeling_batch_size(rows, dim, partitions, 1.0, ample, ample), 32768);
+  // This GPU budget fits the pair matrix but only 20000 assignment rows.
+  EXPECT_EQ(ace_labeling_batch_size(rows, dim, partitions, 1.0, gpu_fixed + 20000 * per_row, ample),
+            20000);
+  EXPECT_EQ(ace_labeling_batch_size(rows, dim, partitions, 1.0, ample, host_fixed + 7 * per_row),
+            7);
+  EXPECT_EQ(ace_labeling_batch_size(rows, dim, partitions, 1.0, ample, host_fixed + per_row - 1),
+            0);
+  EXPECT_EQ(ace_labeling_batch_size(rows, dim, partitions, 1.0, gpu_fixed + pairs - 1, ample), 0);
+  EXPECT_EQ(ace_labeling_gpu_bytes(rows, dim, partitions, 1.0, 20000), gpu_fixed + 20000 * per_row);
+  // Core-only builds need no pair matrix or score vectors.
+  EXPECT_GT(ace_labeling_batch_size(rows, dim, partitions, 0.0, gpu_fixed + pairs - 1, ample), 0);
+  EXPECT_EQ(ace_labeling_batch_size(10, dim, 2, 1.0, ample, ample), 10);
+  EXPECT_EQ(ace_labeling_batch_size(rows, dim, partitions, 1.0, 0, ample), 0);
+}
+
+TEST(AceBudget, LabelingMemorySaturatesInsteadOfWrapping)
+{
+  EXPECT_EQ(ace_labeling_gpu_bytes(1'000'000, std::numeric_limits<size_t>::max(), 8192, 1.0),
+            std::numeric_limits<size_t>::max());
+  EXPECT_EQ(ace_labeling_batch_size(1'000'000,
+                                    std::numeric_limits<size_t>::max(),
+                                    16384,
+                                    1.0,
+                                    std::numeric_limits<size_t>::max(),
+                                    std::numeric_limits<size_t>::max()),
+            0);
+}
+
 }  // namespace cuvs::neighbors::cagra::detail
diff --git a/python/cuvs/cuvs/tests/test_cagra_ace.py b/python/cuvs/cuvs/tests/test_cagra_ace.py
index 86b9832b38..61683cf52b 100644
--- a/python/cuvs/cuvs/tests/test_cagra_ace.py
+++ b/python/cuvs/cuvs/tests/test_cagra_ace.py
@@ -295,7 +295,7 @@ def test_cagra_ace_tiny_memory_limit_triggers_disk_mode():
             ef_construction=100,
             build_dir=temp_dir,
             use_disk=False,  # Not explicitly requesting disk mode
-            max_host_memory_gb=0.001,  # Tiny limit to force disk mode
+            max_host_memory_gb=0.004,  # Fit labeling buffers but force disk mode
             max_gpu_memory_gb=0.0,  # No GPU memory limit
         )
 
diff --git a/python/cuvs/cuvs/tests/test_hnsw_ace.py b/python/cuvs/cuvs/tests/test_hnsw_ace.py
index 05d5a41170..badf15e28d 100644
--- a/python/cuvs/cuvs/tests/test_hnsw_ace.py
+++ b/python/cuvs/cuvs/tests/test_hnsw_ace.py
@@ -247,7 +247,7 @@ def test_hnsw_ace_tiny_memory_limit_triggers_disk_mode():
             npartitions=2,
             build_dir=temp_dir,
             use_disk=False,  # Not explicitly requesting disk mode
-            max_host_memory_gb=0.001,  # Tiny limit to force disk mode
+            max_host_memory_gb=0.004,  # Fit labeling buffers but force disk mode
             max_gpu_memory_gb=0.0,  # No GPU memory limit
         )
 

From 697bb51a5c59bb43fd74e51d62223b8c9129821e Mon Sep 17 00:00:00 2001
From: Julian Miller <julianmiller@nvidia.com>
Date: Mon, 12 Oct 2026 15:47:48 +0200
Subject: [PATCH 07/13] Document ACE target alpha limits across language APIs

---
 c/include/cuvs/neighbors/cagra.h                           | 3 +++
 c/include/cuvs/neighbors/hnsw.h                            | 3 +++
 cpp/include/cuvs/neighbors/cagra.hpp                       | 3 +++
 fern/pages/c_api/c-api-neighbors-cagra.md                  | 2 +-
 fern/pages/c_api/c-api-neighbors-hnsw.md                   | 2 +-
 fern/pages/cpp_api/cpp-api-neighbors-cagra.md              | 2 +-
 .../java_api/java-api-com-nvidia-cuvs-cuvsaceparams.md     | 7 +++++--
 .../java_api/java-api-com-nvidia-cuvs-hnswaceparams.md     | 7 +++++--
 fern/pages/python_api/python-api-neighbors-cagra.md        | 2 +-
 fern/pages/python_api/python-api-neighbors-hnsw.md         | 2 +-
 .../src/main/java/com/nvidia/cuvs/CuVSAceParams.java       | 3 +++
 .../src/main/java/com/nvidia/cuvs/HnswAceParams.java       | 3 +++
 python/cuvs/cuvs/neighbors/cagra/cagra.pyx                 | 3 +++
 python/cuvs/cuvs/neighbors/hnsw/hnsw.pyx                   | 3 +++
 14 files changed, 36 insertions(+), 9 deletions(-)

diff --git a/c/include/cuvs/neighbors/cagra.h b/c/include/cuvs/neighbors/cagra.h
index 0fc371e3f6..46b2ab1f84 100644
--- a/c/include/cuvs/neighbors/cagra.h
+++ b/c/include/cuvs/neighbors/cagra.h
@@ -202,6 +202,9 @@ struct cuvsAceParams {
   /**
    * Target number of augmented (spill) copies per input vector.
    * A value of 0 disables augmentation. The default is 1.0.
+   * Must be finite and in [0, min(3, N - 1)], where N is the effective partition
+   * count after build-time adjustment.
+   * This is an expected budget; sampling can change the realized spill count.
    */
   double target_alpha;
 };
diff --git a/c/include/cuvs/neighbors/hnsw.h b/c/include/cuvs/neighbors/hnsw.h
index a25416018f..1cec7f1048 100644
--- a/c/include/cuvs/neighbors/hnsw.h
+++ b/c/include/cuvs/neighbors/hnsw.h
@@ -93,6 +93,9 @@ struct cuvsHnswAceParams {
   /**
    * Target number of augmented (spill) copies per input vector.
    * A value of 0 disables augmentation. The default is 1.0.
+   * Must be finite and in [0, min(3, N - 1)], where N is the effective partition
+   * count after build-time adjustment.
+   * This is an expected budget; sampling can change the realized spill count.
    */
   double target_alpha;
 };
diff --git a/cpp/include/cuvs/neighbors/cagra.hpp b/cpp/include/cuvs/neighbors/cagra.hpp
index bb902b384f..964b429945 100644
--- a/cpp/include/cuvs/neighbors/cagra.hpp
+++ b/cpp/include/cuvs/neighbors/cagra.hpp
@@ -238,6 +238,9 @@ struct ace_params {
    * augmentation ratio matches this value. A value of 0 disables augmentation. Values above 1
    * allow a vector to spill into multiple nearby partitions. The implementation bounds each
    * vector to three spill partitions to keep memory and build work predictable.
+   * Must be finite and in [0, min(3, N - 1)], where N is the effective partition count
+   * after build-time adjustment. Sampling can change
+   * the realized spill count; the target is not a strict allocation bound.
    */
   double target_alpha = 1.0;
 
diff --git a/fern/pages/c_api/c-api-neighbors-cagra.md b/fern/pages/c_api/c-api-neighbors-cagra.md
index d77008ab2c..b05d5a6b52 100644
--- a/fern/pages/c_api/c-api-neighbors-cagra.md
+++ b/fern/pages/c_api/c-api-neighbors-cagra.md
@@ -100,7 +100,7 @@ struct cuvsAceParams {
 | `use_disk` | `bool` | Whether to use disk-based storage for ACE build.<br /><br />When true, enables disk-based operations for memory-efficient graph construction. |
 | `max_host_memory_gb` | `double` | Maximum host memory to use for ACE build in GiB.<br /><br />When set to 0 (default), uses available host memory. When set to a positive value, limits host memory usage to the specified amount. Useful for testing or when running alongside other memory-intensive processes. |
 | `max_gpu_memory_gb` | `double` | Maximum GPU memory to use for ACE build in GiB.<br /><br />When set to 0 (default), uses available GPU memory. When set to a positive value, limits GPU memory usage to the specified amount. Useful for testing or when running alongside other memory-intensive processes. |
-| `target_alpha` | `double` | Target number of augmented (spill) copies per input vector. A value of 0 disables augmentation. The default is 1.0. |
+| `target_alpha` | `double` | Target number of augmented (spill) copies per input vector. A value of 0 disables augmentation. The default is 1.0. Must be finite and in [0, min(3, N - 1)], where N is the effective partition count after build-time adjustment. This is an expected budget; sampling can change the realized spill count. |
 
 <a id="cuvscagraindexparams"></a>
 ### cuvsCagraIndexParams
diff --git a/fern/pages/c_api/c-api-neighbors-hnsw.md b/fern/pages/c_api/c-api-neighbors-hnsw.md
index 6de8ee4a4e..449fe570b4 100644
--- a/fern/pages/c_api/c-api-neighbors-hnsw.md
+++ b/fern/pages/c_api/c-api-neighbors-hnsw.md
@@ -60,7 +60,7 @@ struct cuvsHnswAceParams {
 | `use_disk` | `bool` | Whether to use disk-based storage for ACE build. When true, enables disk-based operations for memory-efficient graph construction. |
 | `max_host_memory_gb` | `double` | Maximum host memory to use for ACE build in GiB. When set to 0 (default), uses available host memory. Useful for testing or when running alongside other memory-intensive processes. |
 | `max_gpu_memory_gb` | `double` | Maximum GPU memory to use for ACE build in GiB. When set to 0 (default), uses available GPU memory. Useful for testing or when running alongside other memory-intensive processes. |
-| `target_alpha` | `double` | Target number of augmented (spill) copies per input vector. A value of 0 disables augmentation. The default is 1.0. |
+| `target_alpha` | `double` | Target number of augmented (spill) copies per input vector. A value of 0 disables augmentation. The default is 1.0. Must be finite and in [0, min(3, N - 1)], where N is the effective partition count after build-time adjustment. This is an expected budget; sampling can change the realized spill count. |
 
 <a id="cuvshnswaceparamscreate"></a>
 ### cuvsHnswAceParamsCreate
diff --git a/fern/pages/cpp_api/cpp-api-neighbors-cagra.md b/fern/pages/cpp_api/cpp-api-neighbors-cagra.md
index f49a64f0fb..f3181f1b68 100644
--- a/fern/pages/cpp_api/cpp-api-neighbors-cagra.md
+++ b/fern/pages/cpp_api/cpp-api-neighbors-cagra.md
@@ -222,7 +222,7 @@ struct ace_params {
 | `use_disk` | `bool` | Whether to use disk-based storage for ACE build.<br /><br />When true, enables disk-based operations for memory-efficient graph construction. |
 | `max_host_memory_gb` | `double` | Maximum host memory to use for ACE build in GiB.<br /><br />When set to 0 (default), uses available host memory. When set to a positive value, limits host memory usage to the specified amount. Useful for testing or when running alongside other memory-intensive processes. |
 | `max_gpu_memory_gb` | `double` | Maximum GPU memory to use for ACE build in GiB.<br /><br />When set to 0 (default), uses available GPU memory. When set to a positive value, limits GPU memory usage to the specified amount. Useful for testing or when running alongside other memory-intensive processes. |
-| `target_alpha` | `double` | Target number of augmented (spill) copies per input vector.<br /><br />ACE samples boundary-margin scores and chooses a global threshold whose expected augmentation ratio matches this value. A value of 0 disables augmentation. Values above 1 allow a vector to spill into multiple nearby partitions. The implementation bounds each vector to three spill partitions to keep memory and build work predictable. |
+| `target_alpha` | `double` | Target number of augmented (spill) copies per input vector.<br /><br />ACE samples boundary-margin scores and chooses a global threshold whose expected augmentation ratio matches this value. A value of 0 disables augmentation. Values above 1 allow a vector to spill into multiple nearby partitions. The implementation bounds each vector to three spill partitions to keep memory and build work predictable. Must be finite and in [0, min(3, N - 1)], where N is the effective partition count after build-time adjustment. Sampling can change the realized spill count; the target is not a strict allocation bound. |
 
 ## CAGRA index extend parameters
 
diff --git a/fern/pages/java_api/java-api-com-nvidia-cuvs-cuvsaceparams.md b/fern/pages/java_api/java-api-com-nvidia-cuvs-cuvsaceparams.md
index f02e9e686c..d15f4ca069 100644
--- a/fern/pages/java_api/java-api-com-nvidia-cuvs-cuvsaceparams.md
+++ b/fern/pages/java_api/java-api-com-nvidia-cuvs-cuvsaceparams.md
@@ -253,6 +253,9 @@ Sets the target number of augmented (spill) copies per input vector.
 
 A value of 0 disables augmentation. Values above 1 allow a vector to spill into multiple
 nearby partitions. The default is `1.0`.
+Must be finite and in `[0, min(3, N - 1)]`, where `N` is the effective
+partition count after build-time adjustment.
+Sampling can change the realized spill count.
 
 **Parameters**
 
@@ -264,7 +267,7 @@ nearby partitions. The default is `1.0`.
 
 an instance of Builder
 
-_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSAceParams.java:295`_
+_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSAceParams.java:298`_
 
 ### build
 
@@ -278,6 +281,6 @@ Builds an instance of `CuVSAceParams`.
 
 an instance of `CuVSAceParams`
 
-_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSAceParams.java:305`_
+_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSAceParams.java:308`_
 
 _Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSAceParams.java:17`_
diff --git a/fern/pages/java_api/java-api-com-nvidia-cuvs-hnswaceparams.md b/fern/pages/java_api/java-api-com-nvidia-cuvs-hnswaceparams.md
index 6fce368644..85d8b42d14 100644
--- a/fern/pages/java_api/java-api-com-nvidia-cuvs-hnswaceparams.md
+++ b/fern/pages/java_api/java-api-com-nvidia-cuvs-hnswaceparams.md
@@ -243,6 +243,9 @@ Sets the target number of augmented (spill) copies per input vector.
 
 A value of 0 disables augmentation. Values above 1 allow a vector to spill into multiple
 nearby partitions. The default is `1.0`.
+Must be finite and in `[0, min(3, N - 1)]`, where `N` is the effective
+partition count after build-time adjustment.
+Sampling can change the realized spill count.
 
 **Parameters**
 
@@ -254,7 +257,7 @@ nearby partitions. The default is `1.0`.
 
 an instance of Builder
 
-_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/HnswAceParams.java:212`_
+_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/HnswAceParams.java:215`_
 
 ### build
 
@@ -268,6 +271,6 @@ Builds an instance of `HnswAceParams`.
 
 an instance of `HnswAceParams`
 
-_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/HnswAceParams.java:222`_
+_Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/HnswAceParams.java:225`_
 
 _Source: `java/cuvs-java/src/main/java/com/nvidia/cuvs/HnswAceParams.java:16`_
diff --git a/fern/pages/python_api/python-api-neighbors-cagra.md b/fern/pages/python_api/python-api-neighbors-cagra.md
index a6bed0d22d..c762386a35 100644
--- a/fern/pages/python_api/python-api-neighbors-cagra.md
+++ b/fern/pages/python_api/python-api-neighbors-cagra.md
@@ -28,7 +28,7 @@ for each partition independently.
 | `use_disk` | `bool, default = False` | Whether to use disk-based storage for ACE build. When true, enables disk-based operations for memory-efficient graph construction. |
 | `max_host_memory_gb` | `float, default = 0` | Maximum host memory to use for ACE build in GiB. When set to 0 (default), uses available host memory. Useful for testing or when running alongside other memory-intensive processes. |
 | `max_gpu_memory_gb` | `float, default = 0` | Maximum GPU memory to use for ACE build in GiB. When set to 0 (default), uses available GPU memory. Useful for testing or when running alongside other memory-intensive processes. |
-| `target_alpha` | `float, default = 1.0` | Target number of augmented copies per input vector. A value of 0 disables augmentation; values above 1 allow multi-partition spills. |
+| `target_alpha` | `float, default = 1.0` | Target number of augmented copies per input vector. A value of 0 disables augmentation; values above 1 allow multi-partition spills. Must be finite and in [0, min(3, N - 1)], where N is the effective partition count after build-time adjustment. Sampling can change the realized spill count. |
 
 **Constructor**
 
diff --git a/fern/pages/python_api/python-api-neighbors-hnsw.md b/fern/pages/python_api/python-api-neighbors-hnsw.md
index 5573c07a1f..2429656bfe 100644
--- a/fern/pages/python_api/python-api-neighbors-hnsw.md
+++ b/fern/pages/python_api/python-api-neighbors-hnsw.md
@@ -27,7 +27,7 @@ partition independently.
 | `use_disk` | `bool, default = False (optional)` | Whether to use disk-based storage for ACE build. When true, enables disk-based operations for memory-efficient graph construction. |
 | `max_host_memory_gb` | `float, default = 0 (optional)` | Maximum host memory to use for ACE build in GiB. When set to 0 (default), uses available host memory. Useful for testing or when running alongside other memory-intensive processes. |
 | `max_gpu_memory_gb` | `float, default = 0 (optional)` | Maximum GPU memory to use for ACE build in GiB. When set to 0 (default), uses available GPU memory. Useful for testing or when running alongside other memory-intensive processes. |
-| `target_alpha` | `float, default = 1.0 (optional)` | Target number of augmented copies per input vector. A value of 0 disables augmentation; values above 1 allow multi-partition spills. |
+| `target_alpha` | `float, default = 1.0 (optional)` | Target number of augmented copies per input vector. A value of 0 disables augmentation; values above 1 allow multi-partition spills. Must be finite and in [0, min(3, N - 1)], where N is the effective partition count after build-time adjustment. Sampling can change the realized spill count. |
 
 **Constructor**
 
diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSAceParams.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSAceParams.java
index f5e93a953b..89cd724e19 100644
--- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSAceParams.java
+++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CuVSAceParams.java
@@ -289,6 +289,9 @@ public Builder withMaxGpuMemoryGb(double maxGpuMemoryGb) {
      *
      * A value of 0 disables augmentation. Values above 1 allow a vector to spill into multiple
      * nearby partitions. The default is {@code 1.0}.
+     * Must be finite and in {@code [0, min(3, N - 1)]}, where {@code N} is the effective
+     * partition count after build-time adjustment.
+     * Sampling can change the realized spill count.
      *
      * @param targetAlpha the target augmentation ratio
      * @return an instance of Builder
diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/HnswAceParams.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/HnswAceParams.java
index 64c06d1180..9545ad769f 100644
--- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/HnswAceParams.java
+++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/HnswAceParams.java
@@ -206,6 +206,9 @@ public Builder withMaxGpuMemoryGb(double maxGpuMemoryGb) {
      *
      * A value of 0 disables augmentation. Values above 1 allow a vector to spill into multiple
      * nearby partitions. The default is {@code 1.0}.
+     * Must be finite and in {@code [0, min(3, N - 1)]}, where {@code N} is the effective
+     * partition count after build-time adjustment.
+     * Sampling can change the realized spill count.
      *
      * @param targetAlpha the target augmentation ratio
      * @return an instance of Builder
diff --git a/python/cuvs/cuvs/neighbors/cagra/cagra.pyx b/python/cuvs/cuvs/neighbors/cagra/cagra.pyx
index 753492afac..2412c01043 100644
--- a/python/cuvs/cuvs/neighbors/cagra/cagra.pyx
+++ b/python/cuvs/cuvs/neighbors/cagra/cagra.pyx
@@ -109,6 +109,9 @@ cdef class AceParams:
     target_alpha : float, default = 1.0
         Target number of augmented copies per input vector. A value of 0
         disables augmentation; values above 1 allow multi-partition spills.
+        Must be finite and in [0, min(3, N - 1)], where N is the effective
+        partition count after build-time adjustment. Sampling can change the
+        realized spill count.
     """
     cdef cuvsAceParams* params
     cdef bytes _build_dir_bytes  # Keep Python bytes alive for property access
diff --git a/python/cuvs/cuvs/neighbors/hnsw/hnsw.pyx b/python/cuvs/cuvs/neighbors/hnsw/hnsw.pyx
index 9479f04726..a9c8cab172 100644
--- a/python/cuvs/cuvs/neighbors/hnsw/hnsw.pyx
+++ b/python/cuvs/cuvs/neighbors/hnsw/hnsw.pyx
@@ -71,6 +71,9 @@ cdef class AceParams:
     target_alpha : float, default = 1.0 (optional)
         Target number of augmented copies per input vector. A value of 0
         disables augmentation; values above 1 allow multi-partition spills.
+        Must be finite and in [0, min(3, N - 1)], where N is the effective
+        partition count after build-time adjustment. Sampling can change the
+        realized spill count.
     """
 
     cdef cuvsHnswAceParams* params

From 0b3ecfc9e758515625f988f271ac5dfdb071a48f Mon Sep 17 00:00:00 2001
From: Julian Miller <julianmiller@nvidia.com>
Date: Mon, 12 Oct 2026 15:49:03 +0200
Subject: [PATCH 08/13] Exercise ACE alpha validation and zero and multi-spill
 builds

---
 python/cuvs/cuvs/tests/test_cagra_ace.py | 56 +++++++++++++++++++++++-
 python/cuvs/cuvs/tests/test_hnsw_ace.py  | 52 +++++++++++++++++++++-
 2 files changed, 105 insertions(+), 3 deletions(-)

diff --git a/python/cuvs/cuvs/tests/test_cagra_ace.py b/python/cuvs/cuvs/tests/test_cagra_ace.py
index 61683cf52b..bfb5d911e7 100644
--- a/python/cuvs/cuvs/tests/test_cagra_ace.py
+++ b/python/cuvs/cuvs/tests/test_cagra_ace.py
@@ -23,6 +23,7 @@
 
 
 def test_ace_target_alpha_params():
+    """Preserve default and explicit alpha values in the Python parameter object."""
     defaults = cagra.AceParams()
     configured = cagra.AceParams(target_alpha=1.5)
 
@@ -43,7 +44,10 @@ def run_cagra_ace_build_search_test(
     ef_construction=100,
     use_disk=False,
     hierarchy="gpu",
+    target_alpha=1.0,
+    expected_recall=0.7,
 ):
+    """Build and search ACE indexes, reloading HNSW output for disk builds."""
     dataset = generate_data((n_rows, n_cols), dtype)
     queries = generate_data((n_queries, n_cols), dtype)
     if metric == "inner_product":
@@ -63,6 +67,7 @@ def run_cagra_ace_build_search_test(
             ef_construction=ef_construction,
             build_dir=temp_dir,
             use_disk=use_disk,
+            target_alpha=target_alpha,
         )
 
         # Build parameters
@@ -112,7 +117,8 @@ def run_cagra_ace_build_search_test(
             skl_idx = nn_skl.kneighbors(queries, return_distance=False)
 
             recall = calc_recall(out_idx_host, skl_idx)
-            assert recall > 0.7
+            assert recall >= expected_recall
+            assert np.all((out_idx_host >= 0) & (out_idx_host < n_rows))
 
             # test that we can get the cagra graph from the index
             graph = index.graph
@@ -165,7 +171,9 @@ def run_cagra_ace_build_search_test(
             )
 
             recall = calc_recall(out_idx, skl_idx)
-            assert recall > 0.7
+            assert recall >= expected_recall
+            assert np.all((out_idx >= 0) & (out_idx < n_rows))
+            assert np.all(np.isfinite(out_dist))
 
 
 def _build_ace_with_disk_workspace(dataset, build_dir):
@@ -324,3 +332,47 @@ def test_cagra_ace_tiny_memory_limit_triggers_disk_mode():
         assert os.path.exists(reordered_file), (
             "Reordered dataset file should exist when disk mode is triggered"
         )
+
+
+@pytest.mark.parametrize(
+    ("target_alpha", "npartitions"),
+    [(-1.0, 4), (float("nan"), 4), (float("inf"), 4), (3.5, 4), (1.5, 2)],
+)
+def test_ace_rejects_invalid_target_alpha(target_alpha, npartitions, tmp_path):
+    """Validate alpha through the build, including the effective partition limit."""
+    dataset = np.zeros((4096, 16), dtype=np.float32)
+    build_dir = tmp_path / "ace"
+    ace_params = cagra.AceParams(
+        npartitions=npartitions,
+        target_alpha=target_alpha,
+        use_disk=True,
+        build_dir=str(build_dir),
+    )
+    params = cagra.IndexParams(build_algo="ace", ace_params=ace_params)
+    with pytest.raises(CuvsException, match="target_alpha"):
+        cagra.build(params, dataset)
+    assert not build_dir.exists()
+
+
+@pytest.mark.parametrize(
+    ("target_alpha", "use_disk", "npartitions"),
+    [(0.0, False, 2), (1.5, True, 3)],
+)
+def test_ace_target_alpha_build_search(target_alpha, use_disk, npartitions):
+    """Exercise zero/multiple spills and reload disk indexes before searching."""
+    run_cagra_ace_build_search_test(
+        # Keep 1000 rows per partition so ACE preserves the requested count.
+        n_rows=1000 * npartitions,
+        n_cols=8,
+        n_queries=16,
+        k=4,
+        npartitions=npartitions,
+        ef_construction=64,
+        hierarchy="none",
+        target_alpha=target_alpha,
+        use_disk=use_disk,
+        # Core-only graphs can be disconnected; do not impose a quality target.
+        expected_recall=0.0 if target_alpha == 0.0 else 0.7,
+        graph_degree=16,
+        intermediate_graph_degree=32,
+    )
diff --git a/python/cuvs/cuvs/tests/test_hnsw_ace.py b/python/cuvs/cuvs/tests/test_hnsw_ace.py
index badf15e28d..f5ceeeb419 100644
--- a/python/cuvs/cuvs/tests/test_hnsw_ace.py
+++ b/python/cuvs/cuvs/tests/test_hnsw_ace.py
@@ -10,11 +10,13 @@
 from sklearn.neighbors import NearestNeighbors
 from sklearn.preprocessing import normalize
 
+from cuvs.common.exceptions import CuvsException
 from cuvs.neighbors import hnsw
 from cuvs.tests.ann_utils import calc_recall, generate_data
 
 
 def test_ace_target_alpha_params():
+    """Preserve default and explicit alpha values in the Python parameter object."""
     defaults = hnsw.AceParams()
     configured = hnsw.AceParams(target_alpha=1.5)
 
@@ -33,7 +35,9 @@ def run_hnsw_ace_build_search_test(
     ef_construction=100,
     use_disk=False,
     hierarchy="gpu",
+    target_alpha=1.0,
     expected_recall=0.9,
+    M=32,
 ):
     """
     Test HNSW index build using ACE algorithm.
@@ -60,12 +64,13 @@ def run_hnsw_ace_build_search_test(
             npartitions=npartitions,
             build_dir=temp_dir,
             use_disk=use_disk,
+            target_alpha=target_alpha,
         )
 
         # Build parameters with ACE configuration
         index_params = hnsw.IndexParams(
             hierarchy=hierarchy,
-            M=32,
+            M=M,
             ef_construction=ef_construction,
             metric=metric,
             ace_params=ace_params,
@@ -119,6 +124,8 @@ def run_hnsw_ace_build_search_test(
         nn_skl.fit(dataset)
         skl_idx = nn_skl.kneighbors(queries, return_distance=False)
 
+        assert np.all((out_idx >= 0) & (out_idx < n_rows))
+        assert np.all(np.isfinite(out_dist))
         recall = calc_recall(out_idx, skl_idx)
         assert recall >= expected_recall, (
             f"Recall {recall:.3f} is below expected {expected_recall}"
@@ -275,3 +282,46 @@ def test_hnsw_ace_tiny_memory_limit_triggers_disk_mode():
         assert os.path.exists(reordered_file), (
             "Reordered dataset file should exist when disk mode is triggered"
         )
+
+
+@pytest.mark.parametrize(
+    ("target_alpha", "npartitions"),
+    [(-1.0, 4), (float("nan"), 4), (float("inf"), 4), (3.5, 4), (1.5, 2)],
+)
+def test_ace_rejects_invalid_target_alpha(target_alpha, npartitions, tmp_path):
+    """Validate alpha through the build, including the effective partition limit."""
+    dataset = np.zeros((4096, 16), dtype=np.float32)
+    build_dir = tmp_path / "ace"
+    ace_params = hnsw.AceParams(
+        npartitions=npartitions,
+        target_alpha=target_alpha,
+        use_disk=True,
+        build_dir=str(build_dir),
+    )
+    params = hnsw.IndexParams(hierarchy="none", ace_params=ace_params)
+    with pytest.raises(CuvsException, match="target_alpha"):
+        hnsw.build(params, dataset)
+    assert not build_dir.exists()
+
+
+@pytest.mark.parametrize(
+    ("target_alpha", "use_disk", "npartitions"),
+    [(1.5, False, 3), (0.0, True, 2)],
+)
+def test_ace_target_alpha_build_search(target_alpha, use_disk, npartitions):
+    """Exercise zero/multiple spills and reload disk indexes before searching."""
+    run_hnsw_ace_build_search_test(
+        # Keep 1000 rows per partition so ACE preserves the requested count.
+        n_rows=1000 * npartitions,
+        n_cols=8,
+        n_queries=16,
+        k=4,
+        npartitions=npartitions,
+        ef_construction=64,
+        hierarchy="none",
+        target_alpha=target_alpha,
+        use_disk=use_disk,
+        # Core-only graphs can be disconnected; do not impose a quality target.
+        expected_recall=0.0 if target_alpha == 0.0 else 0.7,
+        M=8,
+    )

From 44d969b2602c20dfe3b957443a7168ce20a5570a Mon Sep 17 00:00:00 2001
From: Julian Miller <julianmiller@nvidia.com>
Date: Mon, 12 Oct 2026 15:51:41 +0200
Subject: [PATCH 09/13] Document ACE budget helpers and Python alpha accessors

---
 .../neighbors/detail/cagra/cagra_build.cuh    | 57 ++++++++++++-------
 .../python_api/python-api-neighbors-cagra.md  |  2 +
 .../python_api/python-api-neighbors-hnsw.md   |  2 +
 python/cuvs/cuvs/neighbors/cagra/cagra.pyx    |  2 +
 python/cuvs/cuvs/neighbors/hnsw/hnsw.pyx      |  2 +
 5 files changed, 46 insertions(+), 19 deletions(-)

diff --git a/cpp/src/neighbors/detail/cagra/cagra_build.cuh b/cpp/src/neighbors/detail/cagra/cagra_build.cuh
index 64a2a6e3a9..35e4a2851a 100644
--- a/cpp/src/neighbors/detail/cagra/cagra_build.cuh
+++ b/cpp/src/neighbors/detail/cagra/cagra_build.cuh
@@ -294,12 +294,14 @@ inline ace_margin_gate ace_target_alpha_threshold(std::vector<float> scores,
 // consistent quality benefit from exposing larger values.
 constexpr size_t ace_max_spills_per_vector = 3;
 
+/** Mark an unused spill slot without aliasing a valid partition ID. */
 template <typename IdxT>
 constexpr IdxT ace_invalid_partition()
 {
   return std::numeric_limits<IdxT>::max();
 }
 
+/** Select up to n_keep nearest partitions, ordered by squared distance. */
 inline size_t ace_select_nearest_partitions(const float* distances,
                                             size_t n_partitions,
                                             size_t n_keep,
@@ -318,6 +320,7 @@ inline size_t ace_select_nearest_partitions(const float* distances,
   return filled;
 }
 
+/** Distance to the core/spill bisector; coincident centroids yield an infinite margin. */
 inline float ace_boundary_margin(float core_distance, float spill_distance, float centroid_distance)
 {
   return centroid_distance > 0.0f
@@ -325,6 +328,7 @@ inline float ace_boundary_margin(float core_distance, float spill_distance, floa
            : std::numeric_limits<float>::infinity();
 }
 
+/** Collect candidate spill margins from the centroid-training sample for budget calibration. */
 inline std::vector<float> ace_sample_margin_scores(
   raft::resources const& res,
   raft::device_matrix_view<const float, int64_t> sample,
@@ -379,8 +383,10 @@ inline std::vector<float> ace_sample_margin_scores(
   return scores;
 }
 
-// ACE: Get core plus boundary-gated spill partition labels. Only core graph rows are retained;
-// spill vectors enlarge local builds but their graph rows are not merged into the final graph.
+/**
+ * ACE: Get core plus boundary-gated spill partition labels. Only core graph rows are retained;
+ * spill vectors enlarge local builds but their graph rows are not merged into the final graph.
+ */
 template <typename T, typename IdxT>
 void ace_get_partition_labels(
   raft::resources const& res,
@@ -552,7 +558,9 @@ void ace_get_partition_labels(
   }
 }
 
-// ACE: Check partition sizes for stable KNN graph construction
+/**
+ * ACE: Check partition sizes for stable KNN graph construction
+ */
 template <typename IdxT>
 void ace_check_partition_sizes(
   size_t dataset_size,
@@ -637,8 +645,10 @@ void ace_check_partition_sizes(
   }
 }
 
-// ACE: Create forward/backward mappings between original and reordered vector IDs
-// The in-memory path can be parallelized but the disk path requires ordering.
+/**
+ * ACE: Create forward/backward mappings between original and reordered vector IDs
+ * The in-memory path can be parallelized but the disk path requires ordering.
+ */
 template <typename IdxT>
 void ace_create_forward_and_backward_lists(
   size_t dataset_size,
@@ -876,9 +886,11 @@ void ace_adjust_sub_graph_ids_disk(
   raft::resource::sync_stream(res);
 }
 
-// ACE: Reorder dataset based on partition assignments and store to disk
-// Writes two files: reordered_dataset.npy (core partitions) and augmented_dataset.npy (secondary
-// partitions). Uses buffered writes optimized for NVMe storage.
+/**
+ * ACE: Reorder dataset based on partition assignments and store to disk
+ * Writes two files: reordered_dataset.npy (core partitions) and augmented_dataset.npy (secondary
+ * partitions). Uses buffered writes optimized for NVMe storage.
+ */
 template <typename T, typename IdxT>
 void ace_reorder_and_store_dataset(
   raft::resources const& res,
@@ -1240,6 +1252,7 @@ constexpr double imbalance_factor = 3.0;
 
 constexpr double ace_label_expansion_factor = 1.0 + ace_max_spills_per_vector;
 
+/** Estimate core-plus-spill storage with headroom; this is not an overflow bound. */
 inline double ace_estimated_vector_expansion_factor(double target_alpha)
 {
   // Sampled thresholds closely track target_alpha in the validated sweeps. Keep 10% plus 0.1
@@ -1281,7 +1294,9 @@ inline size_t ace_check_labeling_memory(raft::resources const& res,
   return batch;
 }
 
-// Check if disk mode should be used for ACE based on memory constraints
+/**
+ * Check if disk mode should be used for ACE based on memory constraints
+ */
 template <typename T, typename IdxT>
 bool ace_check_use_disk_mode(raft::resources const& res,
                              bool use_disk,
@@ -1460,7 +1475,9 @@ inline void ace_validate_partition_count(size_t n_partitions,
                  : "ACE: number of partitions cannot exceed dataset size");
 }
 
-// Validate and adjust partitions for disk mode memory requirements
+/**
+ * Validate and adjust partitions for disk mode memory requirements
+ */
 template <typename T, typename IdxT>
 void ace_validate_disk_mode_partitions(raft::resources const& res,
                                        size_t& n_partitions,
@@ -1628,15 +1645,17 @@ auto build_from_host_matrix(raft::resources const& res,
                             DatasetViewT const& dataset)
   -> cuvs::neighbors::cagra::index<T, IdxT, DatasetViewT>;
 
-// Build CAGRA index using ACE (Augmented Core Extraction) partitioning
-// ACE enables building indexes for datasets too large to fit in GPU memory by:
-// 1. Partitioning the dataset using balanced k-means into a non-overlapping core and
-// boundary-gated spill partitions
-// 2. Building sub-indexes for each partition independently
-// 3. Concatenating sub-graphs (of core partitions) into a final unified index
-// Supports both in-memory and disk-based modes depending on available host memory.
-// In disk mode, the graph is stored in build_dir and dataset is reordered on disk.
-// The returned index is not usable for search. Use the created files for search instead.
+/**
+ * Build CAGRA index using ACE (Augmented Core Extraction) partitioning
+ * ACE enables building indexes for datasets too large to fit in GPU memory by:
+ * 1. Partitioning the dataset using balanced k-means into a non-overlapping core and
+ * boundary-gated spill partitions
+ * 2. Building sub-indexes for each partition independently
+ * 3. Concatenating sub-graphs (of core partitions) into a final unified index
+ * Supports both in-memory and disk-based modes depending on available host memory.
+ * In disk mode, the graph is stored in build_dir and dataset is reordered on disk.
+ * The returned index is not usable for search. Use the created files for search instead.
+ */
 template <typename T, typename IdxT, typename DatasetViewT>
   requires cuvs::neighbors::is_host_dataset_view_v<DatasetViewT>
 auto build_ace(raft::resources const& res, const index_params& params, DatasetViewT const& dataset)
diff --git a/fern/pages/python_api/python-api-neighbors-cagra.md b/fern/pages/python_api/python-api-neighbors-cagra.md
index c762386a35..c6fc4e8fbd 100644
--- a/fern/pages/python_api/python-api-neighbors-cagra.md
+++ b/fern/pages/python_api/python-api-neighbors-cagra.md
@@ -91,6 +91,8 @@ def max_gpu_memory_gb(self)
 def target_alpha(self)
 ```
 
+Return the expected number of spill copies per input vector.
+
 ### get_handle
 
 ```python
diff --git a/fern/pages/python_api/python-api-neighbors-hnsw.md b/fern/pages/python_api/python-api-neighbors-hnsw.md
index 2429656bfe..497e24e43a 100644
--- a/fern/pages/python_api/python-api-neighbors-hnsw.md
+++ b/fern/pages/python_api/python-api-neighbors-hnsw.md
@@ -82,6 +82,8 @@ def max_gpu_memory_gb(self)
 def target_alpha(self)
 ```
 
+Return the expected number of spill copies per input vector.
+
 ## IndexParams
 
 ```python
diff --git a/python/cuvs/cuvs/neighbors/cagra/cagra.pyx b/python/cuvs/cuvs/neighbors/cagra/cagra.pyx
index 2412c01043..a0613fccaa 100644
--- a/python/cuvs/cuvs/neighbors/cagra/cagra.pyx
+++ b/python/cuvs/cuvs/neighbors/cagra/cagra.pyx
@@ -132,6 +132,7 @@ cdef class AceParams:
                  max_host_memory_gb=0,
                  max_gpu_memory_gb=0,
                  target_alpha=1.0):
+        """Initialize ACE parameters, including the expected spill-copy budget."""
         self.params.npartitions = npartitions
         self.params.ef_construction = ef_construction
         self.params.use_disk = use_disk
@@ -179,6 +180,7 @@ cdef class AceParams:
 
     @property
     def target_alpha(self):
+        """Return the expected number of spill copies per input vector."""
         return self.params.target_alpha
 
     def get_handle(self):
diff --git a/python/cuvs/cuvs/neighbors/hnsw/hnsw.pyx b/python/cuvs/cuvs/neighbors/hnsw/hnsw.pyx
index a9c8cab172..663b424997 100644
--- a/python/cuvs/cuvs/neighbors/hnsw/hnsw.pyx
+++ b/python/cuvs/cuvs/neighbors/hnsw/hnsw.pyx
@@ -94,6 +94,7 @@ cdef class AceParams:
                  max_host_memory_gb=0,
                  max_gpu_memory_gb=0,
                  target_alpha=1.0):
+        """Initialize ACE parameters, including the expected spill-copy budget."""
         self.params.npartitions = npartitions
         self._build_dir_bytes = build_dir.encode('utf-8')
         self.params.build_dir = self._build_dir_bytes
@@ -126,6 +127,7 @@ cdef class AceParams:
 
     @property
     def target_alpha(self):
+        """Return the expected number of spill copies per input vector."""
         return self.params.target_alpha
 
 

From a4a6caaad2d71ba39d3cfc97d2dc1ae7584d74fc Mon Sep 17 00:00:00 2001
From: Julian Miller <julianmiller@nvidia.com>
Date: Tue, 13 Oct 2026 11:01:47 +0200
Subject: [PATCH 10/13] Fix implementation include paths for CAGRA helper tests

---
 cpp/tests/CMakeLists.txt | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt
index eee4150eed..a9b9318c1d 100644
--- a/cpp/tests/CMakeLists.txt
+++ b/cpp/tests/CMakeLists.txt
@@ -241,6 +241,10 @@ ConfigureTest(
   GPUS 1
   PERCENT 100
 )
+# ACE helper tests include cagra_build.cuh and its generated distance declarations.
+target_include_directories(
+  NEIGHBORS_ANN_CAGRA_FLOAT_UINT32_TEST PRIVATE "${CUVS_SOURCE_DIR}/src" "${CUVS_BINARY_DIR}/src"
+)
 
 ConfigureTest(
   NAME NEIGHBORS_ANN_CAGRA_FILTER_UDF_TEST

From e8554ab82df0006ff595ebe47a1eed6bd5c56608 Mon Sep 17 00:00:00 2001
From: Julian Miller <julianmiller@nvidia.com>
Date: Tue, 13 Oct 2026 12:52:27 +0200
Subject: [PATCH 11/13] Add private dependencies required by CAGRA helper tests

---
 cpp/tests/CMakeLists.txt | 1 +
 1 file changed, 1 insertion(+)

diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt
index a9b9318c1d..6b29ec29cc 100644
--- a/cpp/tests/CMakeLists.txt
+++ b/cpp/tests/CMakeLists.txt
@@ -238,6 +238,7 @@ ConfigureTest(
 ConfigureTest(
   NAME NEIGHBORS_ANN_CAGRA_FLOAT_UINT32_TEST
   PATH neighbors/ann_cagra/test_float_uint32_t.cu neighbors/ann_cagra/test_iterative_cagra_q.cu
+  ADDITIONAL_DEP cuco::cuco kvikio::kvikio
   GPUS 1
   PERCENT 100
 )

From 3bfb56f9540c98a8c3e2cc92b1b80ad732723864 Mon Sep 17 00:00:00 2001
From: Julian Miller <julianmiller@nvidia.com>
Date: Tue, 13 Oct 2026 15:44:44 +0200
Subject: [PATCH 12/13] Use feasible budgets in ACE host-memory fallback tests

---
 cpp/tests/neighbors/ann_hnsw_ace.cuh | 15 ++++++++-------
 1 file changed, 8 insertions(+), 7 deletions(-)

diff --git a/cpp/tests/neighbors/ann_hnsw_ace.cuh b/cpp/tests/neighbors/ann_hnsw_ace.cuh
index 90a3e46caf..9837c8cf4c 100644
--- a/cpp/tests/neighbors/ann_hnsw_ace.cuh
+++ b/cpp/tests/neighbors/ann_hnsw_ace.cuh
@@ -393,7 +393,7 @@ class AnnHnswAceTest : public ::testing::TestWithParam<AnnHnswAceInputs> {
 
   void testHnswAceMemoryLimitFallback()
   {
-    // This test verifies that setting tiny memory limits forces disk mode automatically
+    // This test verifies that a constrained host memory budget forces disk mode automatically
     // Create temporary directory for ACE build
     std::string temp_dir = std::string("/tmp/cuvs_hnsw_ace_memlimit_test_") +
                            std::to_string(std::time(nullptr)) + "_" +
@@ -412,16 +412,17 @@ class AnnHnswAceTest : public ::testing::TestWithParam<AnnHnswAceInputs> {
       hnsw_params.hierarchy = hnsw::HnswHierarchy::GPU;
       hnsw_params.M         = 32;
 
-      // Configure ACE parameters with tiny memory limits to force disk mode
+      // Configure ACE parameters with a host memory limit to force disk mode
       auto ace_params            = graph_build_params::ace_params();
       ace_params.npartitions     = ps.npartitions;
       ace_params.ef_construction = ps.ef_construction;
       ace_params.build_dir       = temp_dir;
       ace_params.use_disk        = false;  // Not explicitly requesting disk mode
 
-      // Selected host memory limit should enforce disk mode
-      ace_params.max_host_memory_gb  = 0.001;
-      ace_params.max_gpu_memory_gb   = 3.0;
+      // Leave room for labeling while keeping the host budget below the in-memory build
+      // estimate. Do not cap GPU memory below the device-dependent RAFT workspace budget.
+      ace_params.max_host_memory_gb  = ps.max_host_memory_gb;
+      ace_params.max_gpu_memory_gb   = ps.max_gpu_memory_gb;
       hnsw_params.graph_build_params = ace_params;
 
       // Build HNSW index using ACE - should automatically fall back to disk mode
@@ -998,8 +999,8 @@ inline std::vector<AnnHnswAceInputs> generate_hnsw_ace_memory_fallback_inputs()
      false,  // use_disk (not explicitly set, should be triggered by memory limit)
      cuvs::distance::DistanceType::L2Expanded,
      0.0,    // min_recall (not checked in fallback test)
-     0.001,  // max_host_memory_gb (tiny limit to force disk mode)
-     0.001}  // max_gpu_memory_gb (tiny limit to force disk mode)
+     0.004,  // max_host_memory_gb (fits labeling buffers but forces disk mode)
+     0.0}    // max_gpu_memory_gb (use device memory; this tests host-memory fallback)
   };
 }
 

From 0308314e714c2498fc950d2f48fbaae68d7c74ef Mon Sep 17 00:00:00 2001
From: Julian Miller <julianmiller@nvidia.com>
Date: Fri, 16 Oct 2026 13:28:33 +0200
Subject: [PATCH 13/13] Simplify ACE score storage and reuse budget and
 parameter checks

---
 .../src/cuvs/cuvs_ann_bench_param_parser.h    | 17 +----
 cpp/src/neighbors/detail/cagra/cagra_ace.cuh  | 70 +++++--------------
 .../ann_cagra/test_float_uint32_t.cu          | 22 +-----
 3 files changed, 22 insertions(+), 87 deletions(-)

diff --git a/cpp/bench/ann/src/cuvs/cuvs_ann_bench_param_parser.h b/cpp/bench/ann/src/cuvs/cuvs_ann_bench_param_parser.h
index b78fb32b0c..90b73d37a6 100644
--- a/cpp/bench/ann/src/cuvs/cuvs_ann_bench_param_parser.h
+++ b/cpp/bench/ann/src/cuvs/cuvs_ann_bench_param_parser.h
@@ -474,20 +474,9 @@ void parse_build_param(const nlohmann::json& conf,
     } else if (conf.value("graph_build_algo", "") == "ACE") {
       cagra_params.graph_build_params = cuvs::neighbors::cagra::graph_build_params::ace_params{};
     }
-    // Parse ACE parameters if provided
-    nlohmann::json ace_conf = collect_conf_with_prefix(conf, "ace_");
-    if (!ace_conf.empty()) {
-      auto ace_params = cuvs::neighbors::cagra::graph_build_params::ace_params();
-      if (ace_conf.contains("npartitions")) { ace_params.npartitions = ace_conf.at("npartitions"); }
-      if (ace_conf.contains("build_dir")) { ace_params.build_dir = ace_conf.at("build_dir"); }
-      if (ace_conf.contains("ef_construction")) {
-        ace_params.ef_construction = ace_conf.at("ef_construction");
-      }
-      if (ace_conf.contains("use_disk")) { ace_params.use_disk = ace_conf.at("use_disk"); }
-      if (ace_conf.contains("target_alpha")) {
-        ace_params.target_alpha = ace_conf.at("target_alpha");
-      }
-      cagra_params.graph_build_params = ace_params;
+    // Preserve ACE precedence among prefixed configs; the shared parser applies the values.
+    if (!collect_conf_with_prefix(conf, "ace_").empty()) {
+      cagra_params.graph_build_params = cuvs::neighbors::cagra::graph_build_params::ace_params{};
     }
     ::parse_build_param<T, IdxT>(conf, cagra_params);
     return cagra_params;
diff --git a/cpp/src/neighbors/detail/cagra/cagra_ace.cuh b/cpp/src/neighbors/detail/cagra/cagra_ace.cuh
index 03725adbc5..205edb87c3 100644
--- a/cpp/src/neighbors/detail/cagra/cagra_ace.cuh
+++ b/cpp/src/neighbors/detail/cagra/cagra_ace.cuh
@@ -177,24 +177,6 @@ inline size_t ace_labeling_sample_count(size_t rows, size_t partitions, double r
   return std::min(rows, std::max(static_cast<size_t>(rows * rate), minimum));
 }
 
-/** Peak explicit device buffers in labeling, excluding the resource workspace. */
-inline size_t ace_labeling_gpu_bytes(
-  size_t rows, size_t dim, size_t partitions, double target_alpha, size_t batch_rows = 32768)
-{
-  const long double samples   = ace_labeling_sample_count(rows, partitions);
-  const long double centroids = partitions;
-  // Sample and centroid matrices survive through assignment. The centroid-pair device
-  // matrix is released before score sampling; assignment uses the admitted batch size.
-  const long double persistent     = (samples + centroids) * dim;
-  const long double centroid_pairs = target_alpha > 0.0 ? centroids * centroids : 0.0L;
-  const long double assignment     = static_cast<long double>(batch_rows) * (dim + centroids);
-  const long double bytes = sizeof(float) * (persistent + std::max(centroid_pairs, assignment));
-  // Saturate impossible requests instead of wrapping the admission estimate.
-  return bytes >= static_cast<long double>(std::numeric_limits<size_t>::max())
-           ? std::numeric_limits<size_t>::max()
-           : static_cast<size_t>(bytes);
-}
-
 /** Choose a labeling batch that fits both budgets; zero means the fixed buffers cannot fit. */
 inline size_t ace_labeling_batch_size(size_t rows,
                                       size_t dim,
@@ -208,10 +190,10 @@ inline size_t ace_labeling_batch_size(size_t rows,
   const long double centroids = partitions;
   const long double pairs     = target_alpha > 0.0 ? centroids * centroids : 0.0L;
   const long double gpu_fixed = sizeof(float) * (samples + centroids) * dim;
-  // The host retains centroid-pair distances during assignment. Reserve both the final
-  // score vector and thread-local scores during calibration (at most three spills per row).
+  // The host retains centroid-pair distances during assignment and one score per candidate
+  // spill during calibration (at most three spills per row).
   const long double scores =
-    target_alpha > 0.0 ? 2.0L * samples * std::min<size_t>(3, partitions - 1) : 0.0L;
+    target_alpha > 0.0 ? samples * std::min<size_t>(3, partitions - 1) : 0.0L;
   const long double host_fixed = sizeof(float) * (samples * dim + pairs + scores);
   const long double per_row    = sizeof(float) * (static_cast<long double>(dim) + centroids);
   if (per_row == 0 || gpu_fixed + sizeof(float) * pairs > gpu_budget ||
@@ -324,8 +306,8 @@ inline std::vector<float> ace_sample_margin_scores(
   auto distances_device = raft::make_device_matrix<float, int64_t>(res, chunk_size, n_partitions);
   auto stream           = raft::resource::get_cuda_stream(res).get();
 
-  std::vector<float> scores;
-  scores.reserve(n_samples * max_spills);
+  // Every sample has max_spills candidates; disjoint slots avoid per-thread buffers and locking.
+  std::vector<float> scores(n_samples * max_spills);
   for (size_t base = 0; base < n_samples; base += chunk_size) {
     const size_t rows = std::min(chunk_size, n_samples - base);
     auto sample_chunk = raft::make_device_matrix_view<const float, int64_t>(
@@ -341,22 +323,18 @@ inline std::vector<float> ace_sample_margin_scores(
 
 #pragma omp parallel
     {
-      std::vector<float> local_scores;
-      local_scores.reserve(rows * max_spills / omp_get_num_threads() + max_spills);
       std::vector<std::pair<float, size_t>> nearest(1 + max_spills);
 #pragma omp for
       for (size_t row = 0; row < rows; row++) {
         const size_t found = ace_select_nearest_partitions(
           &distances_chunk_host(row, 0), n_partitions, 1 + max_spills, nearest.data());
         for (size_t spill = 1; spill < found; spill++) {
-          local_scores.push_back(
+          scores[(base + row) * max_spills + spill - 1] =
             ace_boundary_margin(nearest[0].first,
                                 nearest[spill].first,
-                                centroid_distances(nearest[0].second, nearest[spill].second)));
+                                centroid_distances(nearest[0].second, nearest[spill].second));
         }
       }
-#pragma omp critical
-      scores.insert(scores.end(), local_scores.begin(), local_scores.end());
     }
   }
   return scores;
@@ -537,9 +515,7 @@ void ace_get_partition_labels(
   }
 }
 
-/**
- * ACE: Check partition sizes for stable KNN graph construction
- */
+// Check partition sizes for stable KNN graph construction
 template <typename IdxT>
 void ace_check_partition_sizes(
   size_t dataset_size,
@@ -577,9 +553,8 @@ void ace_check_partition_sizes(
 
   double avg_core_vectors      = static_cast<double>(total_core_vectors) / n_partitions;
   double avg_augmented_vectors = static_cast<double>(total_augmented_vectors) / n_partitions;
-  double avg_total_vectors =
+  double expected_avg_vectors =
     static_cast<double>(total_core_vectors + total_augmented_vectors) / n_partitions;
-  double expected_avg_vectors = avg_total_vectors;
 
   RAFT_LOG_INFO("ACE: Core vectors        - Total: %lu, Avg: %.1f, Min: %lu, Max: %lu",
                 total_core_vectors,
@@ -623,10 +598,8 @@ void ace_check_partition_sizes(
   }
 }
 
-/**
- * ACE: Create forward/backward mappings between original and reordered vector IDs
- * The in-memory path can be parallelized but the disk path requires ordering.
- */
+// Create forward/backward mappings between original and reordered vector IDs
+// The in-memory path can be parallelized but the disk path requires ordering.
 template <typename IdxT>
 void ace_create_forward_and_backward_lists(
   size_t dataset_size,
@@ -864,11 +837,9 @@ void ace_adjust_sub_graph_ids_disk(
   raft::resource::sync_stream(res);
 }
 
-/**
- * ACE: Reorder dataset based on partition assignments and store to disk
- * Writes two files: reordered_dataset.npy (core partitions) and augmented_dataset.npy (secondary
- * partitions). Uses buffered writes optimized for NVMe storage.
- */
+// Reorder dataset based on partition assignments and store to disk
+// Writes two files: reordered_dataset.npy (core partitions) and augmented_dataset.npy (secondary
+// partitions). Uses buffered writes optimized for NVMe storage.
 template <typename T, typename IdxT>
 void ace_reorder_and_store_dataset(
   raft::host_matrix_view<const T, int64_t, row_major> dataset,
@@ -1308,9 +1279,7 @@ inline size_t ace_check_labeling_memory(raft::resources const& res,
   return batch;
 }
 
-/**
- * Check if disk mode should be used for ACE based on memory constraints
- */
+// Check if disk mode should be used for ACE based on memory constraints
 template <typename T, typename IdxT>
 bool ace_check_use_disk_mode(raft::resources const& res,
                              bool use_disk,
@@ -1400,11 +1369,6 @@ bool ace_check_use_disk_mode(raft::resources const& res,
 
   mem.labeling_batch_size = ace_check_labeling_memory(
     res, dataset_size, dataset_dim, n_partitions, target_alpha, sizeof(IdxT), mem);
-  const size_t labeling_gpu_required =
-    ace_labeling_gpu_bytes(
-      dataset_size, dataset_dim, n_partitions, target_alpha, mem.labeling_batch_size) +
-    extra_gpu_workspace_size;
-  gpu_memory_required = std::max(gpu_memory_required, labeling_gpu_required);
 
   bool gpu_memory_limited = mem.available_gpu_memory < gpu_memory_required;
 
@@ -1458,9 +1422,7 @@ inline void ace_validate_partition_count(size_t n_partitions,
                  : "ACE: number of partitions cannot exceed dataset size");
 }
 
-/**
- * Validate and adjust partitions for disk mode memory requirements
- */
+// Validate and adjust partitions for disk mode memory requirements
 template <typename T, typename IdxT>
 void ace_validate_disk_mode_partitions(raft::resources const& res,
                                        size_t& n_partitions,
diff --git a/cpp/tests/neighbors/ann_cagra/test_float_uint32_t.cu b/cpp/tests/neighbors/ann_cagra/test_float_uint32_t.cu
index 43f8540f4a..0861dcf903 100644
--- a/cpp/tests/neighbors/ann_cagra/test_float_uint32_t.cu
+++ b/cpp/tests/neighbors/ann_cagra/test_float_uint32_t.cu
@@ -204,29 +204,16 @@ TEST(AceBudget, ZeroAndMaximumBudgets)
   }
 }
 
-TEST(AceBudget, LabelingMemoryIncludesConcurrentBuffers)
-{
-  constexpr size_t rows       = 1'000'000;
-  constexpr size_t partitions = 8192;
-  constexpr size_t dim        = 128;
-  constexpr size_t expected =
-    sizeof(float) * ((819'200ULL + partitions) * dim + 32768ULL * (dim + partitions));
-  EXPECT_EQ(ace_labeling_gpu_bytes(rows, dim, partitions, 1.0), expected);
-  EXPECT_GT(expected, 1ULL << 30);
-  EXPECT_GT(ace_labeling_gpu_bytes(rows, dim, partitions * 2, 0.0),
-            ace_labeling_gpu_bytes(rows, dim, partitions, 0.0));
-  EXPECT_EQ(ace_labeling_sample_count(1000, 20), 1000);
-}
-
 TEST(AceBudget, LabelingBatchRespectsHostAndGpuBudgetsAbove8192Partitions)
 {
   constexpr size_t rows = 2'000'000, partitions = 16384, dim = 128;
   constexpr size_t samples    = 1'638'400;
   constexpr size_t gpu_fixed  = sizeof(float) * (samples + partitions) * dim;
   constexpr size_t pairs      = sizeof(float) * partitions * partitions;
-  constexpr size_t host_fixed = sizeof(float) * (samples * dim + 6 * samples) + pairs;
+  constexpr size_t host_fixed = sizeof(float) * (samples * dim + 3 * samples) + pairs;
   constexpr size_t per_row    = sizeof(float) * (dim + partitions);
   constexpr size_t ample      = 8ULL << 30;
+  EXPECT_EQ(ace_labeling_sample_count(1000, 20), 1000);
   EXPECT_EQ(ace_labeling_batch_size(rows, dim, partitions, 1.0, ample, ample), 32768);
   // This GPU budget fits the pair matrix but only 20000 assignment rows.
   EXPECT_EQ(ace_labeling_batch_size(rows, dim, partitions, 1.0, gpu_fixed + 20000 * per_row, ample),
@@ -236,17 +223,14 @@ TEST(AceBudget, LabelingBatchRespectsHostAndGpuBudgetsAbove8192Partitions)
   EXPECT_EQ(ace_labeling_batch_size(rows, dim, partitions, 1.0, ample, host_fixed + per_row - 1),
             0);
   EXPECT_EQ(ace_labeling_batch_size(rows, dim, partitions, 1.0, gpu_fixed + pairs - 1, ample), 0);
-  EXPECT_EQ(ace_labeling_gpu_bytes(rows, dim, partitions, 1.0, 20000), gpu_fixed + 20000 * per_row);
   // Core-only builds need no pair matrix or score vectors.
   EXPECT_GT(ace_labeling_batch_size(rows, dim, partitions, 0.0, gpu_fixed + pairs - 1, ample), 0);
   EXPECT_EQ(ace_labeling_batch_size(10, dim, 2, 1.0, ample, ample), 10);
   EXPECT_EQ(ace_labeling_batch_size(rows, dim, partitions, 1.0, 0, ample), 0);
 }
 
-TEST(AceBudget, LabelingMemorySaturatesInsteadOfWrapping)
+TEST(AceBudget, LabelingMemoryRejectsImpossibleDimensions)
 {
-  EXPECT_EQ(ace_labeling_gpu_bytes(1'000'000, std::numeric_limits<size_t>::max(), 8192, 1.0),
-            std::numeric_limits<size_t>::max());
   EXPECT_EQ(ace_labeling_batch_size(1'000'000,
                                     std::numeric_limits<size_t>::max(),
                                     16384,