封神!GPU向量搜索源码全解析:从CUDA内核到Python调用,实战避坑一文拉满

cpp                  #pragma once// 防止头文件重复包含                  #include// 包含CUDA运行时头文件          /**             * 批量向量搜索（余弦相似度）             * @param d_queries: 查询向量 [batch_size x vector_dim] (GPU内存)             * @param d_db: 数据库向量 [db_size x vector_dim] (GPU内存)             * @param d_similarities: 相似度矩阵 [batch_size x db_size] (GPU内存)             * @param d_topk_indices: Top-K索引 [batch_size x k] (GPU内存)             * @param d_topk_values: Top-K相似度 [batch_size x k] (GPU内存)             * @param batch_size: 批量大小             * @param db_size: 数据库大小             * @param vector_dim: 向量维度             * @param k: Top-K数量             * @param stream: CUDA流（默认0）             * @return: CUDA错误码             */          cudaError_t cuda_vector_search_batch(          const float* d_queries,          const float* d_db,          float* d_similarities,          int* d_topk_indices,          float* d_topk_values,          int batch_size,          int db_size,          int vector_dim,          int k,          cudaStream_t stream = 0          );          // 其余3个接口（归一化、欧氏距离、点积）的声明格式类似，核心是参数和返回值的定义          // 重点说明：所有输入的d_xxx参数，都是GPU内存地址，需要用户提前分配GPU内存

cpp                  __global__ void cosine_similarity_batch_kernel(                  const float* __restrict__ query_vectors,// __restrict__ 禁止指针别名，提升优化效率                  const float* __restrict__ db_vectors,                  float* __restrict__ similarities,                  int batch_size,                  int db_size,                  int vector_dim                  ) {                  // 1. 计算当前线程对应的查询向量索引和数据库向量索引                  int query_idx = blockIdx.y;// 每个block对应一个查询向量（y轴方向）                  int db_idx = blockIdx.x * blockDim.x + threadIdx.x;// 每个thread对应一个数据库向量                  // 边界检查：避免线程索引超出有效范围                  if (query_idx >= batch_size || db_idx >= db_size) return;                  // 2. 共享内存缓存查询向量：每个block加载一个查询向量到共享内存                  extern __shared__ float shared_query[];// 动态共享内存，大小在Kernel调用时指定                  int tid = threadIdx.x;                  // 多线程并行加载查询向量到共享内存（避免单线程加载太慢）                  for (int i = tid; i < vector_dim; i += blockDim.x) {                  if (query_idx < batch_size) {                  shared_query[i] = query_vectors[query_idx * vector_dim + i];                  }                  }                  __syncthreads();// 等待所有线程加载完成，避免数据竞争                  // 3. 计算点积、查询向量范数、数据库向量范数                  float dot_product = 0.0f;                  float query_norm = 0.0f;                  float db_norm = 0.0f;                  for (int i = 0; i < vector_dim; i++) {                  float q_val = shared_query[i];// 从共享内存读取，速度快                  float d_val = db_vectors[db_idx * vector_dim + i];// 从全局内存读取                  dot_product += q_val * d_val;// 点积累加                  query_norm += q_val * q_val;// 查询向量范数平方累加                  db_norm += d_val * d_val;// 数据库向量范数平方累加                  }                  // 4. 计算余弦相似度，避免除以0                  float norm_product = sqrtf(query_norm) * sqrtf(db_norm);                  if (norm_product > 1e-8f) {// 阈值判断，避免除以0导致NaN                  similarities[query_idx * db_size + db_idx] = dot_product / norm_product;                  } else {                  similarities[query_idx * db_size + db_idx] = 0.0f;                  }                  }

cpp                  __global__ void topk_select_kernel(                  const float* __restrict__ similarities,                  int* __restrict__ topk_indices,                  float* __restrict__ topk_values,                  int batch_size,                  int db_size,                  int k                  ) {                  int query_idx = blockIdx.x;// 每个block对应一个查询向量                  if (query_idx >= batch_size) return;                  int tid = threadIdx.x;                  int db_idx = tid;// 每个thread对应一个数据库向量的相似度                  // 1. 读取当前线程对应的相似度值，超出范围则设为最小浮点数                  float sim_val = (db_idx < db_size) ?                  similarities[query_idx * db_size + db_idx] : -FLT_MAX;                  // 2. 共享内存存储Top-K结果（值和索引）                  extern __shared__ float s_data[];                  float* s_values = s_data;// 前k个位置存储相似度值                  int* s_indices = (int*)&s_data[k];// 后k个位置存储索引                  // 3. 初始化Top-K结果：值设为最小浮点数，索引设为-1                  if (tid < k) {                  s_values[tid] = -FLT_MAX;                  s_indices[tid] = -1;                  }                  __syncthreads();                  // 4. 归约算法筛选Top-K（位排序网络）                  for (int stride = 1; stride < db_size; stride *= 2) {                  int partner = tid ^ stride;// 找到配对线程                  if (partner < db_size) {                  float partner_val = similarities[query_idx * db_size + partner];                  // 比较当前线程和配对线程的相似度，保留较大值                  if ((tid & stride) && sim_val < partner_val) {                  sim_val = partner_val;                  db_idx = partner;                  } else if (!(tid & stride) && sim_val > partner_val) {                  // 当前值更大，保持不变                  }                  }                  __syncthreads();                  }                  // 5. 将Top-K结果写回全局内存                  if (tid < k) {                  int global_idx = query_idx * k + tid;                  topk_values[global_idx] = sim_val;                  topk_indices[global_idx] = db_idx;                  }                  }

cpp                  cudaError_t cuda_vector_search_batch(                  const float* d_queries,                  const float* d_db,                  float* d_similarities,                  int* d_topk_indices,                  float* d_topk_values,                  int batch_size,                  int db_size,                  int vector_dim,                  int k,                  cudaStream_t stream                  ) {                  const int threads = 256;// 每个block的线程数，通常设为256（GPU最佳线程数）                  const int blocks_db = (db_size + threads – 1) / threads;// 数据库向量的block数（向上取整）                  // 1. 配置相似度计算Kernel的网格和线程块                  dim3 grid(blocks_db, batch_size);// grid维度：(数据库block数, 查询向量数)                  int shared_mem_size = vector_dim * sizeof(float);// 共享内存大小（存储一个查询向量）                  // 调用余弦相似度Kernel                  cosine_similarity_batch_kernel<<>>(          d_queries, d_db, d_similarities, batch_size, db_size, vector_dim          );          // 2. 配置Top-K筛选Kernel的网格和线程块          int topk_shared_mem = k * sizeof(float) + k * sizeof(int);// 共享内存大小（存储Top-K值和索引）          topk_select_kernel<<>>(           d_similarities, d_topk_indices, d_topk_values, batch_size, db_size, k           );           // 返回CUDA错误码，便于排查问题           return cudaGetLastError();           }

cpp                  class GPUVectorSearchEngine {                  public:                  // 构造函数：初始化GPU设备、CUDA Stream、缓存内存                  GPUVectorSearchEngine(int device_id = 0) : device_id_(device_id) {                  cudaSetDevice(device_id_);// 设置使用的GPU设备ID                  cudaStreamCreate(&stream_);// 创建CUDA Stream，用于异步执行                  // 分配缓存内存（1MB），用于临时存储数据，减少内存分配开销                  cudaMalloc(&d_buffer_, 1024 * 1024 * sizeof(float));                  }                  // 析构函数：释放GPU内存和CUDA Stream，避免内存泄漏                  ~GPUVectorSearchEngine() {                  if (d_buffer_) cudaFree(d_buffer_);// 释放缓存内存                  if (stream_) cudaStreamDestroy(stream_);// 销毁CUDA Stream                  }                  // 批量向量搜索方法（核心方法）                  std::tuple         <py::array_t, py::array_t> search_batch(           py::array_tqueries_array,            py::array_tdb_array,             int k = 10             ) {             // 1. 获取numpy数组的缓冲区信息（形状、数据指针等）             py::buffer_info q_buf = queries_array.request();             py::buffer_info db_buf = db_array.request();             // 2. 解析输入参数（批量大小、向量维度、数据库大小）             int batch_size = q_buf.shape[0];             int vector_dim = q_buf.shape[1];             int db_size = db_buf.shape[0];             // 3. 分配GPU内存（存储查询向量、数据库向量、相似度矩阵、Top-K结果）             float *d_queries, *d_db, *d_similarities;             int *d_topk_indices;             float *d_topk_values;             cudaMalloc(&d_queries, batch_size * vector_dim * sizeof(float));             cudaMalloc(&d_db, db_size * vector_dim * sizeof(float));             cudaMalloc(&d_similarities, batch_size * db_size * sizeof(float));             cudaMalloc(&d_topk_indices, batch_size * k * sizeof(int));             cudaMalloc(&d_topk_values, batch_size * k * sizeof(float));             // 4. 将CPU内存的numpy数组异步传输到GPU内存             cudaMemcpyAsync(d_queries, q_buf.ptr, batch_size * vector_dim * sizeof(float),             cudaMemcpyHostToDevice, stream_);             cudaMemcpyAsync(d_db, db_buf.ptr, db_size * vector_dim * sizeof(float),             cudaMemcpyHostToDevice, stream_);             // 5. 调用C++ API，执行GPU向量搜索             cuda_vector_search_batch(d_queries, d_db, d_similarities,             d_topk_indices, d_topk_values,             batch_size, db_size, vector_dim, k, stream_);             // 6. 分配Python输出数组（Top-K索引和值）             auto indices_array = py::array_t({batch_size, k});              auto values_array = py::array_t({batch_size, k});               py::buffer_info i_buf = indices_array.request();               py::buffer_info v_buf = values_array.request();               // 7. 将GPU内存的结果异步传输回CPU内存               cudaMemcpyAsync(i_buf.ptr, d_topk_indices, batch_size * k * sizeof(int),               cudaMemcpyDeviceToHost, stream_);               cudaMemcpyAsync(v_buf.ptr, d_topk_values, batch_size * k * sizeof(float),               cudaMemcpyDeviceToHost, stream_);               // 等待所有异步操作完成               cudaStreamSynchronize(stream_);               // 8. 释放GPU内存，避免内存泄漏               cudaFree(d_queries);               cudaFree(d_db);               cudaFree(d_similarities);               cudaFree(d_topk_indices);               cudaFree(d_topk_values);               // 返回Python可直接使用的numpy数组（元组形式）               return std::make_tuple(indices_array, values_array);               }               // 其余方法（normalize、euclidean_distance）逻辑类似，都是“内存传输→GPU计算→结果回传”               private:               int device_id_;// GPU设备ID               cudaStream_t stream_;// CUDA Stream               float* d_buffer_;// 缓存内存（GPU）               };</py::array_t

cpp                  PYBIND11_MODULE(gpu_vector, m) {                  m.doc() = “GPU加速的向量搜索库”;// 模块说明文档                  // 注册GPUVectorSearchEngine类                  py::class_(m, “GPUVectorSearchEngine”)          .def(py::init())// 注册构造函数（支持传入device_id）           .def(“search_batch”, &GPUVectorSearchEngine::search_batch,           py::arg(“queries”), py::arg(“db”), py::arg(“k”) = 10)// 注册方法，设置默认参数           .def(“normalize”, &GPUVectorSearchEngine::normalize,           py::arg(“vectors”))           .def(“euclidean_distance”, &GPUVectorSearchEngine::euclidean_distance,           py::arg(“queries”), py::arg(“db”));           // 注册独立函数接口（无需创建类实例）           m.def(“search_batch”, &search_batch,           “批量向量搜索”,           py::arg(“queries”), py::arg(“db”), py::arg(“k”) = 10, py::arg(“device_id”) = 0);           // 其余函数注册逻辑类似           }

cmake                  cmake_minimum_required(VERSION 3.18)# 最低CMake版本要求                  project(gpu_vector CXX CUDA)# 项目名称，支持C++和CUDA语言                  set(CMAKE_CXX_STANDARD 17)# C++标准                  set(CMAKE_CUDA_STANDARD 17)# CUDA标准                  find_package(CUDA REQUIRED)# 查找CUDA依赖库                  find_package(pybind11 REQUIRED)# 查找pybind11依赖库                  set(CUDA_SOURCES                  vector_search.cu# CUDA源码文件                  )                  add_compile_options(-O3 -march=native)# 编译优化：O3级优化，适配本地CPU架构                  # 编译静态库（gpu_vector_lib），包含CUDA核心逻辑                  cuda_add_library(gpu_vector_lib STATIC ${CUDA_SOURCES})                  # 为静态库添加CUDA头文件路径                  target_include_directories(gpu_vector_lib PRIVATE ${CUDA_INCLUDE_DIRS})                  # 启用CUDA可分离编译，便于扩展和维护                  set_target_properties(gpu_vector_lib PROPERTIES CUDA_SEPARABLE_COMPILATION ON)                  # 编译Python扩展模块（gpu_vector），用于Python调用                  pybind11_add_module(gpu_vector pybind11_wrapper.cpp)                  # 链接静态库（依赖CUDA核心逻辑）                  target_link_libraries(gpu_vector PRIVATE gpu_vector_lib)                  # 链接CUDA库                  target_link_libraries(gpu_vector PRIVATE ${CUDA_LIBRARIES})                  # 设置Python扩展模块的输出目录（放入binary/python目录）                  set_target_properties(gpu_vector PROPERTIES LIBRARY_OUTPUT_DIRECTORY ${BINARY_DIR}/python)                  # 设置静态库的输出目录（放入binary/lib目录）                  set_target_properties(gpu_vector_lib PROPERTIES ARCHIVE_OUTPUT_DIRECTORY ${BINARY_DIR}/lib)                  # 安装配置：将扩展模块、静态库、头文件安装到系统指定目录                  install(TARGETS gpu_vector DESTINATION lib)                  install(TARGETS gpu_vector_lib DESTINATION lib)                  install(FILES vector_search.cuh DESTINATION include)

bash                  # 1. 创建build目录（用于存放编译文件）                  mkdir build && cd build                  # 2. 运行CMake，配置编译（指定Python路径，可选）                  cmake -DCMAKE_BUILD_TYPE=Release -DCMAKE_INSTALL_PREFIX=../install ..                  # 3. 编译（-j后面跟CPU核心数，加速编译，如-j8）                  make -j8                  # 4. 安装（将扩展模块、静态库、头文件安装到../install目录）                  make install

python                  import numpy as np                  import sys                  # 将扩展模块目录添加到Python路径                  sys.path.append(“../install/python”)                  import gpu_vector                  # ————————– 1. 准备测试数据 ————————–                  batch_size = 100# 查询向量批量大小                  db_size = 100000# 数据库向量数量                  vector_dim = 1024# 向量维度                  # 生成随机查询向量和数据库向量（float32类型，GPU只支持float32）                  queries = np.random.randn(batch_size, vector_dim).astype(np.float32)                  db = np.random.randn(db_size, vector_dim).astype(np.float32)                  # ————————– 2. 方式1：使用类调用（推荐，效率更高） ————————–                  # 初始化GPU向量搜索引擎（指定使用0号GPU）                  engine = gpu_vector.GPUVectorSearchEngine(device_id=0)                  # （可选）向量归一化                  queries_norm = engine.normalize(queries)                  db_norm = engine.normalize(db)                  # 批量向量搜索（Top-K=10）                  topk_indices, topk_values = engine.search_batch(queries_norm, db_norm, k=10)                  # 查看结果（topk_indices：每个查询向量的Top-10索引；topk_values：对应的相似度）                  print(“Top-K索引形状：”, topk_indices.shape)# (100, 10)                  print(“Top-K相似度形状：”, topk_values.shape)# (100, 10)                  print(“第一个查询向量的Top-3结果：”)                  print(“索引：”, topk_indices[0][:3])                  print(“相似度：”, topk_values[0][:3])                  # 计算欧氏距离                  distances = engine.euclidean_distance(queries, db)                  print(“欧氏距离矩阵形状：”, distances.shape)# (100, 100000)                  # ————————– 3. 方式2：直接调用函数（简单场景） ————————–                  topk_indices2, topk_values2 = gpu_vector.search_batch(                  queries_norm, db_norm, k=10, device_id=0                  )                  # 归一化函数调用                  vectors_norm = gpu_vector.normalize_vectors(queries, device_id=0)                  # 欧氏距离函数调用                  distances2 = gpu_vector.euclidean_distance(queries, db, device_id=0)                  # ————————– 4. 性能测试 ————————–                  import time                  # 测试向量搜索速度（10次预热，100次测试）                  engine = gpu_vector.GPUVectorSearchEngine(device_id=0)                  # 预热（避免第一次推理速度不准）                  for _ in range(10):                  engine.search_batch(queries_norm[:1], db_norm, k=10)                  # 测试100次批量搜索的平均速度                  start_time = time.time()                  for _ in range(100):                  engine.search_batch(queries_norm, db_norm, k=10)                  end_time = time.time()                  avg_time = (end_time – start_time) / 100                  qps = batch_size / avg_time# 每秒处理的查询向量数                  print(f”\n批量搜索平均时间：{avg_time*1000:.2f} ms/批”)                  print(f”QPS（每秒查询数）：{qps:.0f}”)                  # 对比CPU计算（余弦相似度）                  def cpu_cosine_similarity(queries, db):                  # CPU串行计算，仅用于对比性能                  queries_norm = queries / np.linalg.norm(queries, axis=1, keepdims=True)                  db_norm = db / np.linalg.norm(db, axis=1, keepdims=True)                  similarities = np.dot(queries_norm, db_norm.T)                  return similarities                  # 测试CPU速度（仅用1个查询向量，避免耗时过长）                  start_time_cpu = time.time()                  cpu_similarities = cpu_cosine_similarity(queries_norm[:1], db_norm)                  end_time_cpu = time.time()                  print(f”\nCPU单个查询向量计算时间：{(end_time_cpu – start_time_cpu)*1000:.2f} ms”)                  print(f”GPU比CPU快：{(end_time_cpu – start_time_cpu)/avg_time:.0f} 倍”)