CAGRA-ACE: Add Augmentation Control

c/include/cuvs/neighbors/cagra.h

@@ -199,6 +199,14 @@ 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.
+   * 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;
 };
 
 typedef struct cuvsAceParams* cuvsAceParams_t;

c/include/cuvs/neighbors/hnsw.h

@@ -90,6 +90,14 @@ 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.
+   * 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;
 };
 
 typedef struct cuvsHnswAceParams* cuvsHnswAceParams_t;

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};
   });
 }
 

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};
   });
 }
 

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);
@@ -466,17 +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"); }
-      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;

cpp/include/cuvs/neighbors/cagra.hpp

@@ -231,6 +231,19 @@ 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.
+   * 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;
+
   ace_params() = default;
 };