[WIP] Improve CAGRA-Q performance and add support for PQ_LEN=8

cpp/include/cuvs/neighbors/cagra.hpp

@@ -204,6 +204,8 @@ enum class search_algo {
 
 enum class hash_mode { HASH = 0, SMALL = 1, AUTO = 100 };
 
+enum class internal_dtype { F16 = 0, E5M2 = 1, AUTO = 100 };
+
 struct search_params : cuvs::neighbors::search_params {
   /** Maximum number of queries to search at the same time (batch size). Auto select when 0.*/
   size_t max_queries = 0;
@@ -277,6 +279,10 @@ struct search_params : cuvs::neighbors::search_params {
    * negative, in which case the filtering rate is automatically calculated.
    */
   float filtering_rate = -1.0;
+
+  /** Data type of the query vector and codebook table on shared memory. Currently, only VPQ
+   * supports FP8. **/
+  internal_dtype smem_dtype = internal_dtype::AUTO;
 };
 
 /**

cpp/src/neighbors/detail/cagra/cagra_search.cuh

@@ -153,6 +153,10 @@ void search_main(raft::resources const& res,
   // Dispatch search parameters based on the dataset kind.
   if (auto* strided_dset = dynamic_cast<const strided_dataset<T, ds_idx_type>*>(&index.data());
       strided_dset != nullptr) {
+    if (params.smem_dtype != cuvs::neighbors::cagra::internal_dtype::AUTO ||
+        params.smem_dtype != cuvs::neighbors::cagra::internal_dtype::F16) {
+      RAFT_LOG_WARN("In this search mode, only AUTO or F16 are supported as the smem_dtype.");
+    }
     // Search using a plain (strided) row-major dataset
     RAFT_EXPECTS(index.metric() != cuvs::distance::DistanceType::CosineExpanded ||
                    index.dataset_norms().has_value(),

cpp/src/neighbors/detail/cagra/compute_distance_00_generate.py

@@ -30,11 +30,14 @@
 """
 
 mxdim_team = [(128, 8), (256, 16), (512, 32)]
+vpq_2_4_mxdim_team = [(64, 4), (128, 8), (256, 16), (512, 32)]
+vpq_8_mxdim_team = [(128, 4), (256, 8), (512, 16), (1024, 32)]
+vrq_mxdim_team = [(64, 4), (128, 8), (256, 16), (512, 32)]
 # mxdim_team = [(64, 8), (128, 16), (256, 32)]
 # mxdim_team = [(32, 8), (64, 16), (128, 32)]
 
 pq_bits = [8]
-pq_lens = [2, 4]
+pq_lens = [2, 4, 8]
 
 # rblock = [(256, 4), (512, 2), (1024, 1)]
 # rcandidates = [32]
@@ -76,26 +79,33 @@
                 f.write(template.format(includes=includes, content=content))
                 cmake_list.append(f"  src/neighbors/detail/cagra/{path}")
 
-        # CAGRA-Q
-        for code_book_t in code_book_types:
-            for pq_len in pq_lens:
+    for pq_len in pq_lens:
+        vpq_mxdim_team = (
+            vpq_8_mxdim_team if pq_len == 8 else vpq_2_4_mxdim_team
+        )
+        for mxdim, team in vpq_mxdim_team:
+            # CAGRA-Q
+            for code_book_t in code_book_types:
                 for pq_bit in pq_bits:
                     for metric in ["L2Expanded"]:
-                        path = f"compute_distance_vpq_{metric}_{type_path}_dim{mxdim}_t{team}_{pq_bit}pq_{pq_len}subd_{code_book_t}.cu"
-                        includes = '#include "compute_distance_vpq-impl.cuh"'
-                        params = f"{metric_prefix}{metric}, {team}, {mxdim}, {pq_bit}, {pq_len}, {code_book_t}, {data_t}, {idx_t}, {distance_t}"
-                        spec = f"vpq_descriptor_spec<{params}>"
-                        content = f"""template struct {spec};"""
-                        specs.append(spec)
-                        with open(path, "w") as f:
-                            f.write(
-                                template.format(
-                                    includes=includes, content=content
-                                )
-                            )
-                            cmake_list.append(
-                                f"  src/neighbors/detail/cagra/{path}"
+                        for enable_fp8 in ["true", "false"]:
+                            path = f"compute_distance_vpq_{metric}_{type_path}_dim{mxdim}_t{team}_{pq_bit}pq_{pq_len}subd_{code_book_t}_fp8{enable_fp8}.cu"
+                            includes = (
+                                '#include "compute_distance_vpq-impl.cuh"'
                             )
+                            params = f"{metric_prefix}{metric}, {team}, {mxdim}, {pq_bit}, {pq_len}, {code_book_t}, {data_t}, {idx_t}, {distance_t}, {enable_fp8}"
+                            spec = f"vpq_descriptor_spec<{params}>"
+                            content = f"""template struct {spec};"""
+                            specs.append(spec)
+                            with open(path, "w") as f:
+                                f.write(
+                                    template.format(
+                                        includes=includes, content=content
+                                    )
+                                )
+                                cmake_list.append(
+                                    f"  src/neighbors/detail/cagra/{path}"
+                                )
 
 # CAGRA (Binary Hamming distance)
 for mxdim, team in mxdim_team: