向量搜索作为大模型应用的关键技术，其性能直接影响着用户体验和系统吞吐量。最近了解到代码助手Cursor的后端实现了每秒百万次的最近邻搜索（QPS），这听起来难以置信，但通过分析其工作负载，我们发现其核心在于对数据进行分区，化整为零，从而实现了惊人的扩展性。本文将深入探讨Cursor是如何利用小索引、KDB.AI以及并行计算等技术，突破向量搜索的性能瓶颈，并展示你如何也能借鉴这些方法。
小索引的威力：化整为零的策略
传统的向量搜索理论往往关注于如何在大型向量语料库（甚至数十亿规模）中找到相似的项目。 然而，当索引缩小到几千个向量时，情况就变得简单多了。 HNSW等索引算法的延迟与数据集大小呈log(n)关系增长。但在只有几千个中等大小的向量规模下，即使是全面的搜索，在经过优化的向量操作的CPU上，也只需要几毫秒。
Cursor的巧妙之处在于，它并没有构建一个统一的巨型索引，而是为每个项目或代码仓库维护数百万个小索引。 这种做法将一个庞大的搜索问题分解成无数个小型、独立的搜索任务。 当你在自己的代码仓库中搜索时，你并不关心别人的代码仓库，反之亦然。 因此，每个查询都只会落在一个微型索引上，从而实现了数据的隔离和高效的并行处理。
案例分析： 假设一个大型代码托管平台，拥有数百万个代码仓库。 如果将所有代码向量化后放入一个统一的索引中，每次搜索都需要扫描整个索引，效率极低。 但如果为每个代码仓库维护一个独立的索引，每次搜索只需要扫描该仓库的索引，搜索范围大大缩小，速度也大大提升。 此外，不同仓库的搜索可以并行进行，进一步提升了整体吞吐量。
KDB.AI：专为分区数据设计的向量数据库
要实现百万级的向量搜索QPS，仅仅依靠小索引是不够的，还需要一个能够高效管理和查询这些索引的向量数据库。 Cursor选择了KDB.AI，一个专为分区数据设计的向量数据库。
KDB.AI的一个关键特性是它对分区数据的原生支持。 在KDB.AI中，一个表可以包含数百万个分片，每个分片都有自己的物理目录。 当你使用分区键进行过滤时，引擎会直接进入相应的目录，而不会触及其他数据。 这种机制极大地提高了查询效率。
代码示例：
以下代码展示了如何使用KDB.AI创建表，并使用user_id进行分区：
!pip install kdbai_client
import kdbai_client as kdbai
import numpy as np
import pandas as pd
from time import time
# 连接到KDB.AI
session = kdbai.Session(endpoint='http://localhost:8082')
# 获取数据库连接
database = session.database('default')
# 定义schema和index
schema = [
{'name': 'user_id', 'type': 'int32'},
{'name': 'embeddings', 'type': 'float32s'}
]
indexes = [
{'name': 'vector_index', 'column': 'embeddings', 'type': 'flat', 'params': {'dims': 64}}
]
# 创建表，并按照user_id进行分区
table = database.create_table('vectors', schema=schema, indexes=indexes, partition_column='user_id')
print("Table created successfully")
在这个例子中，我们创建了一个名为vectors的表，其中包含user_id和embeddings两个列。 user_id列用于标识用户，embeddings列用于存储向量数据。 我们还创建了一个名为vector_index的索引，用于加速向量搜索。 最重要的是，我们指定user_id作为分区列。 这意味着KDB.AI会将数据按照user_id进行分区，每个用户的数据存储在不同的分片中。
性能优势： 通过对数据进行分区，KDB.AI可以实现以下性能优势：
- 更快的查询速度： 当你使用
user_id进行过滤时，KDB.AI只会扫描包含该用户数据的分片，而不会扫描整个表。 这大大提高了查询速度。 - 更高的吞吐量： 由于每个分片都可以独立进行查询，因此KDB.AI可以并行处理多个查询，从而提高整体吞吐量。
- 更好的可扩展性： 当数据量增长时，你可以将表分成更多的分片，从而实现更好的可扩展性。
并行计算：充分利用硬件资源
除了小索引和KDB.AI之外，Cursor还利用并行计算来充分利用硬件资源，从而实现百万级的向量搜索QPS。
KDB.AI提供了两个环境变量：NUM_WRK和THREADS，用于控制工作进程的数量和每个工作进程的线程数。 在一个64核的CPU上，你可以启动4个工作进程，每个进程包含16个线程。 每个工作进程负责处理一部分分区，而每个线程可以处理该分区中的不同数据。 如果流量均匀分布在各个项目上，吞吐量几乎与核心数量成线性关系。
代码示例：
以下代码展示了如何在KDB.AI中进行向量搜索：
import time
time1 = time.time()
example_vector = np.random.randn(64).astype(np.float32)
filter_conditions = [("=", "user_id", 30)]
# 执行搜索
results = table.search(
vectors={"vector_index": [example_vector]},
n=3,
filter=filter_conditions
)
time2 = time.time()
print(time2 - time1)
# 可选，查看结果
print(results)
在这个例子中，我们使用table.search方法进行向量搜索。 我们指定了要搜索的向量、返回的最近邻数量（n=3）以及过滤条件（filter=filter_conditions）。 由于我们使用了user_id进行分区，因此KDB.AI只会扫描包含user_id=30数据的分片。
并行计算的优势： 通过使用并行计算，KDB.AI可以同时处理多个查询，从而提高整体吞吐量。 此外，KDB.AI还可以将查询分解成多个子任务，并在不同的线程上并行执行，进一步提高查询速度。
隔离作为第一原则：解耦带来的高性能
Cursor之所以能够实现百万级的QPS，关键在于其将隔离作为第一原则进行设计。 用户的项目、代码和搜索都是相互隔离的，这使得系统可以更加高效地处理每个用户的请求。
分区边界不是一种负担，而是一种机制，它将普通的服务器变成了高吞吐量的向量设备。 任何尊重这些边界的多租户产品都可以复制类似的结果，而无需使用异构GPU、定制内核或复杂的分布式算法。
总结：百万级QPS并非遥不可及
通过分析Cursor的案例，我们可以看到，实现百万级的向量搜索QPS并非遥不可及。 其核心在于：
- 小索引： 将大型索引分解成无数个小型索引，缩小搜索范围，提高搜索速度。
- KDB.AI： 使用专为分区数据设计的向量数据库，实现高效的数据管理和查询。
- 并行计算： 充分利用硬件资源，并行处理多个查询，提高整体吞吐量。
- 数据隔离：设计时就考虑到数据的隔离性，使得单个用户的请求能够高效处理，互不干扰。
如果你正在构建一个需要处理大量向量搜索请求的应用，可以借鉴Cursor的方法，使用小索引、KDB.AI和并行计算等技术，从而实现更高的性能和更好的用户体验。虽然达到百万级QPS需要进行很多优化，而本文只提及了分区这个良好的开端。