为什么 Milvus 中 GPU 没有比 CPU 加速向量搜索？

Question

我正在尝试使用 GPU 来加速 Milvus 中向量查询和搜索向量维度为 1536 的 100 万个数据点。但是，与 CPU 相比，我没有看到 GPU 有任何显着的性能提升。这是我的设置和我尝试过的：

环境：

• CPU：2 个 Intel(R) Xeon(R) Gold 5320 CPU @ 2.20GHz（104 线程）

• GPU：2 个 NVIDIA A100 80GB PCIe

• NVIDIA 驱动程序版本：555.42.02

• CUDA版本：12.5

步骤：

在 Docker 中运行 Milvus 后启动了 CPU 服务。 在使用 Docker Compose 运行具有 GPU 支持的 Milvus 后启动了 GPU 服务。 使用 hello_milvus.py 示例的修改版本插入数据：

数据量改为100万条，维度改为1536 对于 CPU：使用 IVF_FLAT 索引，nlist=2048 对于 GPU：使用 nlist=2048 的 GPU_IVF_FLAT 索引

使用以下Python代码搜索数据：

from pymilvus import MilvusClient, DataType
import time
import numpy as np
import string
import random

milvus_uri = ""
collection_name = ""
client = MilvusClient(uri=milvus_uri)
client.load_collection(collection_name)

search_params = {
   "metric_type": "L2",
   "params": {"nprobe": 32},
}
vectors_to_search = [np.random.rand(1536).tolist()]
start_time = time.time()
result = client.search(
       collection_name=collection_name,
       data=vectors_to_search, 
       filter=filter_expr,
       anns_field="embeddings", 
       search_params=search_params, 
       limit=10, 
       output_fields=["random"],
       consistency_level="Eventually"
   )
end_time = time.time()
print(f"time cost {end_time-start_time}")

GPU IVF 平面（n探针：32）：

并发	QPS
1	681
5	594
10	546

CPU IVF 平坦（nprobe：32）：

并发	QPS
1	680
5	609
10	580

我的问题

为什么GPU没有加速效果？请帮我看看上面的操作是否有问题。真诚希望得到帮助，谢谢

Answer 1

对于CPU索引来说，一次搜索请求的时间成本包括：

• 请求的网络传输
• ANN 计算
• 网络传输结果

对于GPU索引，一次搜索请求的时间成本包括以下内容：

• 请求的网络传输
• 目标向量从CPU内存复制到GPU内存
• ANN 计算
• 结果从 GPU 内存复制到 CPU 内存
• 网络传输结果

我们可以看到在CPU内存和GPU内存之间复制数据会产生额外的时间成本。 GPU搜索的优点是大NQ搜索，因为GPU具有很强的并行计算能力。 对于小数据集和小NQ搜索，CPU索引和GPU索引没有太大区别。

一个大的NQ搜索是这样的：

NQ = 10000
target_vectors = []
for i in range(NQ):
  target_vectors.append(random_vector)

results = collection.search(
            data=target_vectors,
            anns_field="xxx",
            param=search_params,
            limit=topk,
            consistency_level="Eventually",
        )

您可以尝试增加“并发数”，或者为每个请求提高 NQ 值。

NQ = 100
vectors_to_search = [np.random.rand(1536).tolist() for _ in range (NQ)]

我认为，为了获得更高的QPS，最好在循环外生成随机向量。您可以在线程启动之前预先创建随机向量列表。然后从循环内的列表中选择随机向量。