动态批处理#

动态批处理允许将小的搜索请求分组为批次，以增加设备占用率和吞吐量，同时将延迟保持在限制范围内。

#include

命名空间 cuvs::neighbors::dynamic_batching

索引构建参数#

struct index_params : public cuvs::neighbors::index_params#
#include <dynamic_batching.hpp>

公共成员

int64_t k#

在构建时，搜索的邻居数量是固定的。

int64_t max_batch_size = 100#

提交给上游索引的批次的最大大小。

size_t n_queues = 3#

独立请求队列的数量。

每个队列都与一个唯一的CUDA流和IO设备缓冲区相关联。如果并发请求的数量很高，使用多个队列可以在其他队列忙碌时填充数据并准备批次。此外，队列是并发提交的；这可以通过隐藏内核启动延迟来更好地利用GPU，从而有助于提高吞吐量。

bool conservative_dispatch = false#

默认情况下（conservative_dispatch = false），第一个将查询提交到批处理的CPU线程会尽快（在批处理未满之前）调度上游搜索函数。在这种情况下，它在调用上游搜索时不知道最终的批处理大小，因此每次都以最大批处理大小运行上游搜索，即使批处理中只有一个有效查询。这减少了延迟，但以浪费GPU资源为代价。

替代行为（conservative_dispatch = true）更为保守：调度线程启动收集输入查询的内核，但等待直到批次填满或等待时间超过。只有在那时，它才会获取实际的批次大小并启动上游搜索。因此，以暴露上游搜索延迟为代价，减少了GPU资源的浪费。

经验法则: 对于较大的 max_batch_size，设置 conservative_dispatch = true，否则保持禁用状态。

索引搜索参数#

struct search_params : public cuvs::neighbors::search_params#
#include <dynamic_batching.hpp>

公共成员

double dispatch_timeout_ms = 1.0#

请求可以在队列中停留多长时间（毫秒）。请注意，这仅影响调度时间，并不反映完整的请求延迟；后者取决于上游搜索参数和批量大小。

索引#

template<typename T, typename IdxT>
struct index : public cuvs::neighbors::index#
#include <dynamic_batching.hpp>

轻量级动态批处理索引包装器。

一个轻量级的动态批处理索引管理一个单一的索引和一个单一的搜索参数集。这个结构应该通过复制语义在多个用户之间共享：对底层实现的访问是通过共享指针管理的，参与者之间的并发搜索是线程安全的。

使用示例 

using namespace cuvs::neighbors;
// When creating a dynamic batching index, k parameter has to be passed explicitly.
// The first empty braces default-initialize the parent `neighbors::index_params` (unused).
dynamic_batching::index_params dynb_index_params{{}, k};
// Construct the index by wrapping the upstream index and search parameters.
dynamic_batching::index<float, uint32_t> index{
    res, dynb_index_params, upstream_index, upstream_search_params
};
// Use default search parameters
dynamic_batching::search_params search_params;
// Search K nearest neighbours
auto neighbors = raft::make_device_matrix<uint32_t>(res, n_queries, k);
auto distances = raft::make_device_matrix<float>(res, n_queries, k);
dynamic_batching::search(
    res, search_params, index, queries, neighbors.view(), distances.view()
);

优先队列

动态批处理索引对优先处理个别请求的支持有限。批处理器中只有一个队列池，并且没有功能可以优先处理一个批次而不是另一个批次。每个请求中传递的search_params::dispatch_timeout_ms参数在内部聚合，并且批次在任何超时时间之前都会被分发。在这种逻辑下，高优先级请求永远不能比之前提交的低优先级请求更早处理。

然而，动态批处理索引是轻量级的，不包含任何全局或静态状态。这意味着很容易组合多个批处理器。例如，您可以为每个优先级类构建一个批处理索引：

using namespace cuvs::neighbors;
// Large batch size (128), couple queues (2),
//   enabled conservative dispatch - all for better throughput
dynamic_batching::index_params low_priority_params{{}, k, 128, 2, true};
// Small batch size (16), more queues (4),
//   disabled conservative dispatch - to minimize latency with reasonable throughput
dynamic_batching::index_params high_priority_params{{}, k, 16, 4, false};
// Construct the indexes by wrapping the upstream index and search parameters.
dynamic_batching::index<float, uint32_t> low_priority_index{
    res, low_priority_params, upstream_index, upstream_search_params
};
dynamic_batching::index<float, uint32_t> high_priority_index{
    res, high_priority_params, upstream_index, upstream_search_params
};
// Define a combined search function with priority selection
double high_priority_threshold_ms = 0.1;
auto search_function =
   [low_priority_index, high_priority_index, high_priority_threshold_ms](
     raft::resources const &res,
     dynamic_batching::search_params search_params,
     raft::device_matrix_view<const float, int64_t> queries,
     raft::device_matrix_view<uint32_t, int64_t> neighbors,
     raft::device_matrix_view<float, int64_t> distances) {
   dynamic_batching::search(
       res,
       search_params,
       search_params.dispatch_timeout_ms < high_priority_threshold_ms
         ? high_priority_index : low_priority_index,
       queries,
       neighbors,
       distances
   );
};

Template Parameters:

  • T – 数据类型

  • IdxT – 索引类型

公共函数

template<typename Upstream>
index(const raft::resources &res, const cuvs::neighbors::dynamic_batching::index_params &params, const Upstream &upstream_index, const typename Upstream::search_params_type &upstream_params, const cuvs::neighbors::filtering::base_filter *sample_filter = nullptr)#

通过包装上游索引来构建动态批处理索引。

Template Parameters:

上游 – 上游索引类型

Parameters:

  • res – [in] raft资源

  • params – [in] 动态批处理参数

  • upstream_index – [in] 执行搜索的原始索引（在动态批处理索引的生命周期内，引用必须保持有效）

  • upstream_params – [in] 批处理中所有查询的原始索引搜索参数（这些参数在动态批处理索引的生命周期内按值捕获）

  • sample_filter – [in] 过滤函数，如果有的话，必须对批次中的所有请求相同（指针在动态批处理索引的生命周期内必须保持有效）