GPU-Resident Top-K for Agentic RAG: I Built a CUDA Kernel So My Retrieval Step Would Stop Bouncing Off the GPU

用于智能体 RAG 的 GPU 常驻 Top-K：我构建了一个 CUDA 内核，让我的检索步骤不再在 GPU 上反复跳转 | Towards Data Science

智能体 AI

用于智能体 RAG 的 GPU 常驻 Top-K：我构建了一个 CUDA 内核，让我的检索步骤不再在 GPU 上反复跳转

通过替换 Python 往返的开销，采用自定义的 GPU 内存架构，实现多跳 RAG 的确定性微秒级尾部延迟。

Anubhab Banerjee

2026 年 6 月 19 日

31 分钟阅读

分享

通过将检索循环保留在 GPU 上，我们消除了 PCIe 传输的开销，并完全利用了设备内存带宽。

这是一次高度实证的、共 343 行的 CUDA Top-K 检索之旅。该内核、CPU 神谕和基准测试套件证明了标准智能体 RAG 往返——在 PCIe 总线上传输查询——是您流水线中的隐形杀手。通过将相似性搜索保留在设备内存中，该架构即使在 7 年前的 GTX 1080 上也能实现比优化的 CPU 基线快 8.6 倍的速度。

这是“生产级智能体推理”系列文章的第三部分。每篇文章都会从智能体 LLM 流水线中移除一种冗余的工作。第一部分消除了冗余的预填充。第二部分消除了冗余的等待——多个微智能体如何通过时间切片共享一个 GPU。第三部分（本文）通过自定义的 CUDA Top-K 内核，将 RAG 检索保留在 GPU 上。第四部分则在智能体交接时持久化智能体状态，使得下一个智能体永远不会遇到冷启动问题。

关键要点

问题：在智能体 RAG 中，每个需要上下文的工具调用都会触发一次相似性搜索。默认的流水线会将查询嵌入从 GPU 发送到 Python，让 CPU 对 N 个语料库行进行评分并选择最佳的 K 个，然后将答案返回。这种往返就是隐形的税。计算没有问题；开销是账单。我们都知道，无论你想去哪里，旅行从不便宜（双关语，意指旅行开销）。

简单的解决方法：将语料库一次性上传到 VRAM，然后将相似性评分、Top-K 选择和合并步骤保留在设备上。只有很小的每个查询嵌入（D 个浮点数）和 K 个结果在 PCIe 上传输。

证据：在第一部分和第二部分使用的同款 7 年前的 GTX 1080 上，GPU 常驻路径的检索步骤比 CPU 暴力基线快 8.57 倍。在 K=8 时，它在所有 15 个扫描配置（N ∈ {10k, 50k, 100k, 500k, 1M}, D ∈ {384, 768, 1024}）上都取得了胜利，速度提升从 2.43 倍到 8.57 倍。在 K=32 时，它在 15 个配置中的 13 个上取得胜利，最高达到 7.76 倍。在 K=100 时，V1 选择器故意保持简单，CPU 在 15 个配置中的 14 个上取得胜利。最后一句话是诚实的部分（即使我撒了谎，你也很容易发现）。

TL;DR：默认的智能体 RAG 将 GPU 视为一个服务设备，将检索视为一个 Python 的问题。每次工具调用都会将查询嵌入 D→H 发送出去，让 CPU 计算 N 个点积，对候选结果进行排序，选择前 K 个，然后将索引和分数 H→D 发送回来。对于一个在每次推理步骤中调用向量存储十次的智能体，这个往返过程是主要的成本 —— 不是模型，不是嵌入，而是传输。CUDA-TopK-Retrieval 保持语料库驻留在设备上，在 GPU 上完全运行评分 + 每块部分 Top-K + 多路合并，并暴露一个小型的 C++ 协调器 API（upload_corpus_rowmajor 一次，search_resident 每个查询一次）。每个查询的主机接触字节减少到一个 D 长度的嵌入上行和 2K 个结果下行。在 GTX 1080 上，经过 45 个配置的扫描，GPU 驻留路径在所有 15 个 K=8 的配置上都优于 CPU 往返基准（2.43 倍到 8.57 倍，峰值在 N=1M，D=1024）以及在 15 个 K=32 配置中的 13 个（两个失败的配置都在最小的 N=10k，D=384 和 D=768，此时往返本身已经很便宜；大 N K=32 的胜利倍数上升到 7.76 倍）。在 K=100 时，V1 内核故意保持简单 —— 每块单线程冒泡排序与串行合并 —— CPU 在 15 个配置中的 14 个上获胜；这个上限是本文诚实的要点，也是第 4 部分的清晰铺垫。

GitHub 仓库：https://github.com/AnubhabBanerjee/cuda-topk-retrieval

（在我们开始之前，快速坦白一下：我来自 5G/6G RAN 工程背景。基站的波束选择看起来与 RAG Top-K 非常相似 —— UE 通过接收功率对候选波束的码本进行评分，并通过空中接口回传最佳的几个。下面有一个专门的部分 —— 第 8 部分 —— 但这也是这个内核之所以以这种形式存在的原因。）

架构思维模型 —— 阅读时请保持这个模型在脑海中。

agent.embed(query) → cudaMemcpy H→D (D floats) → row_dot_scores_kernel → partial_topk_block_kernel (P blocks) → merge_partial_topk_kernel → cudaMemcpy D→H (K indices + K scores)

以下所有内容只是对该流程中某一部分的评论。

CUDA TopK 检索概述

1. 一个坦白：你智能体中的每次 RAG 步骤都是一次微小的 PCIe 之旅

在本系列的第 2 部分中，我们成功地隔离了 LLM 智能体的推理循环，使令牌生成在设备上保持高速运行。我们设计了一个避免停滞的系统。但当我们给该智能体一个搜索外部知识库的工具 —— 这是任何多跳检索增强生成（RAG）管道的核心 —— 我们就默默地破坏了所有之前获得的性能，并遇到了瓶颈。如果你曾经通过 Python 检索器将“智能体”管道连接到向量存储，那么每次工具调用实际上发生了什么（稍微夸张一点地描述）：

你：“智能体，帮我找到与‘如何在第 80C 节下申领扣除？’最相关的五个片段。” 智能体：“好的。在 GPU 上嵌入查询。✅” 智能体：“现在将查询嵌入发送回主机。”（cudaMemcpy D→H，约 1,024 个浮点数） Python 检索器：“收到。NumPy 循环。点积 N 次。argpartition。Top-5。”（CPU 对语料库的半百万行逐行计算得分，而一个 9 TFLOP 的 GPU 只能旁观） Python 检索器：“完成。这是索引和得分。” 智能体：“不错。现在将它们发送回 GPU。”（cudaMemcpy H→D，10 个数字） 智能体：“准备好了。问题是什么来着？”

代理拥有一块性能非常不错的 GPU。语料库正存储在 4 GB 的显存中。查询嵌入已经位于 GPU 上 —— 我们就是在那生成的。然后，在每次检索跳转时，我们都要将查询发送回主机，使用 NumPy / FAISS-on-CPU / 手写的循环进行暴力相似性计算，再将结果发送回来。

你的 GPU 的利用率：在检索步骤中大部分时间都处于空闲状态。你的 PCIe 总线：被迫承担了它没有申请的工作量。你的代理的工具调用延迟：被一些既不是模型也不是嵌入的东西主导。这就是个笑话。

这也是每一个能扩展到“内存中只放十个片段”之外的代理式 RAG 演示背后的肮脏秘密。每次检索跳转都会从 GPU 上弹回主机，再弹回 GPU，语料库越大，这种开销就越严重。在一百万行 1024 维嵌入向量的情况下，仅仅来回一次 —— 还不包括评分，仅仅是来回一次 —— 就已经消耗了检索步骤预算的大部分。

CUDA-TopK-Retrieval 就是在你决定来回一次是可选的，并且你宁愿写 343 行 CUDA 代码，也不愿让代理每次找邻居时都去主机内存度假时发生的事情。

现在想象一下这项工作背后的真实负载。它不是“一个问题对应五个片段”。而是多个专业化的微代理 —— 每个代理都运行自己的 RAG 跳转，每个代理都需要对同一语料库进行 Top-K 检索，每个代理目前在每次工具调用时都要支付自己的 PCIe 费用。本系列文章的第一部分消除了预填充的来回。第二部分让 GPU 可以在多个代理之间共享。第三部分说：既然它们现在可以公平地共享显卡，那就不要再让每个代理都回主机去查找邻居了。

2. 为什么存在 Top-K 检索？（一分钟速览）

如果你已经了解过相关内容，可以跳过这一部分。对于其他人，尤其是刚接触这个领域的读者，这里有一个简短、非专业的解释。

现代代理不会将整个知识库塞进提示中。它会进行检索。对于每一个需要有依据上下文的推理步骤，它会将查询嵌入到一个固定维度的向量中（D 个浮点数 —— 通常为 384、768 或 1024），然后将该向量与语料库中每个预嵌入片段的行（N 行，每行也是 D 个浮点数）进行相似性评分，并返回相似性最高的 K 行语料。这就是全部内容。这就是 Top-K 向量搜索。检索增强生成（RAG）只是礼貌地表示“Top-K 加上一个提示模板”。

相似性有两种常见形式。点积是便宜的一种：每个维度只需要一次融合乘加操作，总工作量为 N×D。余弦相似性是点积除以两个 L2 范数的乘积，如果你在数据摄入时对语料库进行一次预归一化，余弦相似性就变成了免费的点积。大多数生产级的向量存储都采用这种预归一化技巧，并称之为“余弦”，同时在查询时执行原始的点积计算。CUDA-TopK-Retrieval 内核支持这两种方式 —— 只是在余弦模式下，它会乘以一个预先计算好的每行范数指针。

主流工具（FAISS、hnswlib、cuVS 的 Python 端、你最喜欢的 SaaS 向量数据库）都执行这种评分 + Top-K 操作。其中大多数都做得很好。问题在于它们在哪里执行这些操作。几乎所有的代理框架都会从 Python 调用检索器，而一旦 Python 处于热点路径上，检索步骤就不再是 GPU 操作 —— 它变成了一个一端连接 GPU 的 PCIe 操作。

解决方法不是“更好的算法”。而是“更短的公路旅行”。

3. “将语料库保留在 GPU 上”的灵光一现（以及为什么听起来简单，实际上却困难）

这个想法很简单：

• 在数据摄入时，将语料库一次性上传到 VRAM。

• 对于每一个到来的查询，使用 cudaMemcpy 将一个很小的 D 维浮点嵌入向量复制到设备上。

• 启动一个评分内核，其中每个 CUDA 线程对应语料库中的一行，计算点积。

• 启动一个部分 Top-K 内核，每个块扫描一个互不重叠的行范围，输出自己的局部候选列表。

• 最后，启动一个合并内核，遍历每个块的头部，以最佳优先顺序输出全局 Top-K。

你只将 2K 个数字通过 cudaMemcpy 返回到主机：K 个索引和 K 个分数。

这是一种“将内存检索视为硬件原语，而非软件 API 调用”的范式。唯一的原因是，要实现这一点需要超过 30 行 PyTorch 脚本，是因为有三个繁琐的边界情况会立即破坏这种简单的方法。

#### 问题 A：在 GPU 上执行 Top-K 在结构上很笨拙

对向量进行评分是简单的部分。这仅仅是矩阵乘法，而你的 GPU 本质上就是为了做这个而生的——这是硬件的最爱。然而，选择部分却是浪漫的终结。让 GPU 执行一个完整的 O(N log N) 排序，只是为了获取 Top-K 的结果，这在计算上是令人反感的；这就像为了找到一张收据而将整个回收箱按字母顺序排列。你可以尝试使用 O(N) 的 argpartition，但这需要树遍历，这会将 GPU 内存合并变成数百万次未对齐的读取。锦标赛选择虽然快速，但前提是你要花整个周末来调试边界情况。而一旦你屈服，使用 thrust 或 cub 排序原语，恭喜你：你刚刚将轻量级、独立的 C++ 管道感染了巨大的构建依赖。

架构特意选择了无聊的答案。它依赖于每个块的微小 O(K²) 冒泡排序，对一个互不重叠的行范围进行扫描，由每个块的单一线程驱动，并以串行多路合并结束。在纸面上，这听起来很糟糕。但在实践中，它工作得非常完美，原因正是内核注释中坦率提到的：

// 单线程块扫描，为它的行分区生成一个局部 Top-K 列表。
// 这不是最快的全局选择方法，但它易于理解，并且与 CPU 的排序规则完全一致。
__device__ void bubble_downward(float* const s, int* const ids, const int n) {
  // 微小的 O(K^2) 排序是可接受的，因为 K 有上限（kMaxSupportedK），并且这个排序在每个块的单一线程上运行。
  for (int i = 0; i < n - 1; ++i) {
    for (int j = 0; j < n - 1 - i; ++j) {
      if (device_is_better(s[j + 1], ids[j + 1], s[j], ids[j])) {
        const float ts = s[j];
        s[j] = s[j + 1];
        s[j + 1] = ts;
        const int ti = ids[j];
        ids[j] = ids[j + 1];
        ids[j + 1] = ti;
      }
    }
  }
}

这是设计合同：V1 选择可审计性而非巧妙性。整个内核足够小，审阅者可以在咖啡时间通读一遍，与 CPU 的“神谕”对齐，并说服自己 GPU 的输出是位精确的。当某天有人想要一个 2 倍的内核时，这个将被替换为一个 warp 专用的锦标赛选择器——文章的第 9 节明确承诺了这一点。对于 K ≤ 32，单线程的冒泡排序确实足够好。对于 K=100，它会急剧下降。这个下降点被记录和基准测试了。（再次参见第 9 节。）

#### 问题 B：GPU 和 CPU 必须在位对位上达成一致的平局规则

嗯，就像世界杯足球小组赛一样，平局发生的频率很高。两个语料库行的得分在 fp32 精度下是相同的。那么，哪一行会获胜呢？

如果 CPU 的“预言者”和 GPU 内核在平局的处理上意见不一致，那么你永远无法信任一个基准测试，因为每一个“不匹配”的警报都变得模棱两可：是 GPU 得分计算错误，还是 GPU 只是用不同的方式处理平局？你将会在凌晨 3 点的 Slack 通道中耗费整整一周的时间。

解决方法是用一句话定义比较器，并在两个地方实现它——一次在主机上，一次在设备上——并确保这两个实现完全相同。

在主机端，它看起来像这样：

// 在浮点数相等语义下，(score, index) 对的字典序“更好”关系。
// 我们使用严格弱排序用于 std::partial_sort：得分越高越好；在完全平局时，索引较小的更好。
bool is_better_score_pair(const float32_t score_lhs, const index_t idx_lhs, const float32_t score_rhs,
                          const index_t idx_rhs) {
  // 主键：相似度得分（对于检索，得分越高越好）。
  if (score_lhs != score_rhs) {
    return score_lhs > score_rhs;
  }
  // 确定性平局处理：优先选择较小的语料库行 ID，以匹配稳定的数据库主键。
  return idx_lhs < idx_rhs;
}

在设备端，它看起来像这样：

// 设备端的比较器副本，用于避免 __device__ 代码路径的跨 TU 链接问题。
__device__ bool device_is_better(const float score_lhs, const int idx_lhs, const float score_rhs, const int idx_rhs) {
  // 与 topk::is_better_score_pair 的排序语义相同，以实现位相同平局处理。
  if (score_lhs != score_rhs) {
    return score_lhs > score_rhs;
  }
  return idx_lhs < idx_rhs;
}

这就是整个平局处理策略，用五行代码重复了两次。得分高的胜出；在完全平局时，索引较小的语料库行胜出。CPU 的“预言者”在 std::partial_sort 中使用它，GPU 在冒泡排序和多路归并中使用它，基准测试框架将不会开始计时，直到 GPU 的输出与 CPU 的输出完全一致——索引顺序相同，得分在较小的 fp32 容差范围内。

这个单一的比较器是文章能够引用加速比的根本原因。没有它，“GPU 比 CPU 快 8 倍”就只是“GPU 比 CPU 快 8 倍地以不同的方式出错”。

#### 问题 C：VRAM 非常宝贵，而最糟糕的 malloc 地点是热点路径

每次查询都分配 GPU 内存，就像每次去杂货店都需要签一份新车租赁合同一样。这是将 1 毫秒的搜索变成 50 毫秒交通堵塞的最简单方式。

相反，GpuTopkEngine::initialize 会提前购买好这辆车。它在引擎启动时运行所有的 cudaMalloc 调用，为最糟糕的配置调整缓冲区的大小。一旦引擎开始处理查询，热点路径就完全摆脱了内存管理。它只是快速的内核启动和微小的数据复制。没有碎片化，没有与分配器的协商，cudaMalloc 永远不会再出现在你的性能分析中。

4. 四阶段流水线（真正酷的部分）

步骤 0：引擎初始化 —— 为 (max_n, max_d, max_k) 分配八个 cudaMallocs   (GpuTopkEngine::initialize)
步骤 1：一次性将语料库上传到 VRAM                                      (upload_corpus_rowmajor)
步骤 2：每个查询 —— 将嵌入向量从主机上传到设备                                     (search_resident, 第一行)
步骤 3：在设备上对 N 行进行评分                                             (row_dot_scores_kernel)
步骤 4：每个块的部分 Top-K                                            (partial_topk_block_kernel)
步骤 5：合并为全局 Top-K                                          (merge_partial_topk_kernel)
步骤 6：将 K 个索引 + K 个评分从设备上传到主机                                       (search_resident, 最后几行)

让我们通过实际代码逐个了解这些步骤。代码片段甚至比源文件还要短。

#### 步骤 1 —— 一次性上传语料库

这是使文章其余部分成为可能的枯燥步骤。语料库在每次摄入时只上传一次，并在引擎的整个生命周期中保持在 VRAM 中：

cudaError_t GpuTopkEngine::upload_corpus_rowmajor(const float32_t* const host_corpus_rowmajor, const index_t N,
                                                  const index_t D) {
  if (N > max_n_ || D > max_d_) {
    return cudaErrorInvalidValue;
  }
  const std::size_t corpus_bytes = sizeof(float) * static_cast<std::size_t>(N) * static_cast<std::size_t>(D);
  const cudaError_t st = cudaMemcpy(d_corpus_, host_corpus_rowmajor, corpus_bytes, cudaMemcpyHostToDevice);
  if (st != cudaSuccess) {
    return st;
  }
  resident_n_ = N;
  resident_d_ = D;
  return cudaSuccess;
}

这就是整个摄入 API。在 1024 维的情况下，一百万向量正好是 4 GB 的 float32 数据，完美地适合于老款 GTX 1080 的 8 GB VRAM。当语料库达到一千万向量时会发生什么？这将变成一个分布式系统问题，而不是内核问题。如果数据超过 VRAM，你需要一个分片策略，我们在第 9 节中会进行介绍。但目前，我们在这里是为了解决计算瓶颈，而不是发明一个新的数据库。

#### 步骤 2 —— 在设备上对 N 行进行评分

每个语料库行对应一个 CUDA 线程。每个块有 256 个线程。每个线程在 D 维度上累加点积，并将一个 float 写入密集的 scores[N] 缓冲区：

// 按行优先方式计算点积，可选余弦归一化；在v1版本中，为了清晰度，牺牲了沿D方向的合并读取。
// 微架构说明：每个线程处理一行很简单；后续可以将D跨线程块进行分块，以提高算术强度。
__global__ void row_dot_scores_kernel(const float* const corpus, const float* const query, const float* const row_l2,
                                      const float query_l2, const int N, const int D, const int cosine_enabled,
                                      float* const scores) {
  // 将每个CUDA线程映射到一个语料库行，以使归约逻辑易于与CPU参考进行核对。
  const int row = static_cast<int>(blockIdx.x) * static_cast<int>(blockDim.x) + static_cast<int>(threadIdx.x);
  if (row >= N) {
    return;
  }
  float acc = 0.0F;
  const int base = row * D;
  for (int col = 0; col < D; ++col) {
    acc += corpus[static_cast<std::size_t>(base + col)] * query[static_cast<std::size_t>(col)];
  }
  if (cosine_enabled != 0) {
    const float denom = query_l2 * row_l2_fetch(row_l2, row);
    scores[static_cast<std::size_t>(row)] =
        denom > 0.0F ? (acc / denom) : -std::numeric_limits<float>::infinity();
  } else {
    scores[static_cast<std::size_t>(row)] = acc;
  }
}

每个线程处理一行是最简单的映射方式。代码注释坦率地承认了这一点：后续可以将D跨线程块进行分块，以提高算术强度。对于V1版本，这为审计者提供了与CPU循环一一对应的映射关系，让他们可以安心入睡。

#### 第三步 — 每个线程块构建自己的局部Top-K

现在是最尴尬的部分。从N中选出Top K在概念上是“排序并切片”，但完整的排序会浪费大部分工作。我们将行范围分配给P个线程块（最多128个），每个线程块处理其不重叠的切片，并使用第3节中的微型冒泡排序，然后写出自己的局部Top-K列表：

const int P = std::min(kMaxPartialBlocks, std::max(1, (static_cast<int>(N) + 4095) / 4096));
  partial_topk_block_kernel<<<P, 1>>>(d_scores_, static_cast<int>(N), static_cast<int>(K), P, d_partial_scores_,
                                      d_partial_indices_);

每个线程块一个线程。是的，这在纸面上是浪费的。这也是为什么人类可以在二十分钟内审计这个内核——策略 if (threadIdx.x != 0 || blockIdx.x >= P) return; 将整个线程块内的推理简化为“线程块的车道0拥有行[start, end)。”每个线程块的s[]和ids[]数组都存储在寄存器/本地内存中，大小由编译时的kMaxSupportedK = 256上限决定。

#### 第四步 — 将局部结果合并为全局Top-K

最后，一个线程在一个线程块上遍历P个游标，这些游标指向每个线程块的列表。每个列表已经是最佳优先的。选择最佳头部；输出；推进该游标；重复K次：

for (int out = 0; out < K; ++out) {
    int best_p = -1;
    float best_s = -std::numeric_limits<float>::infinity();
    int best_i = std::numeric_limits<int>::max();
    for (int p = 0; p < P; ++p) {
      if (heads[p] >= K) {
        continue;
      }
      const float s = partial_scores[static_cast<std::size_t>(p * K + heads[p])];
      const int idx = partial_indices[static_cast<std::size_t>(p * K + heads[p])];
      if (best_p < 0 || device_is_better(s, idx, best_s, best_i)) {
        best_p = p;
        best_s = s;
        best_i = idx;
      }
    }
    out_scores[static_cast<std::size_t>(out)] = best_s;
    out_indices[static_cast<std::size_t>(out)] = best_i;
    heads[best_p] += 1;
  }

合并操作效率极高：最多执行 P * K 次读取和恰好 K 次写入，均由单个线程完成。为了避免浮点数混乱，device_is_better 比较器强制执行严格确定性——如果两个头部在分数上并列，则具有较低语料行索引的头部获胜，完美地模仿了 CPU 的预言机。最后，两次微小的 cudaMemcpy 调用将 K 个胜出的索引和分数送回主机。代理程序接收这些数据，RAG 循环再次启动。

这就是整个热点路径：一次 H -> D 嵌入传输，三次内核启动，以及两次微小的 D -> H 结果复制。没有 Python 主机循环，没有框架开销，也没有任何 PCIe 假期。

5. 收据（即数字）

现在，让我们将其与基线进行比较，看看是否值得付出这些努力。

在基准测试人员到来之前，先快速说明一下方法：以下所有比较均在与第 1 部分和第 2 部分相同的 GPU 上运行（NVIDIA GeForce GTX 1080，Pascal sm_61，8 GB），驱动程序版本 535.309.01，CUDA 12.2，主机 CPU 为 Intel Core i7-8700K，编译器标志为 -O3 -march=native --expt-relaxed-constexpr。进行了三次试验，一次预热，种子为 1，固定 RNG（使用 std::mt19937 和 std::normal_distribution），高斯嵌入，L2 归一化在点积模式下。完整的扫描范围是 N ∈ {10k, 50k, 100k, 500k, 1M} × D ∈ {384, 768, 1024} × K ∈ {8, 32, 100} → 45 种配置，所有测量均通过 cudaEventElapsedTime 进行，使用 cudaDeviceSynchronize 括起每个间隔。代码在 src/host/bench_main.cpp 中；原始数据在 examples/example-run-results/benchmark_run_results.csv 中。

计时的两条路径：

• GPU-驻留（处理）。语料库已经在设备上。每个计时迭代：cudaMemcpy 查询 H→D（D 个浮点数）+ 评分内核 + 每块部分 Top-K + 合并 + cudaMemcpy K 个评分 D→H + cudaMemcpy K 个索引 D→H。端到端。

• CPU 往返（基线）。模拟默认代理流程：cudaMemcpy 查询 D→H + CPU 暴力评分 + std::partial_sort 使用相同的比较器 + cudaMemcpy 索引 H→D + cudaMemcpy 评分 H→D。端到端。

两条路径都在同一进程中运行，使用相同的查询字节和相同的比较器。唯一的区别是工作执行的位置。如果你曾经担任过“PCIe 没问题，我们在隔离中基准测试内核”的职位，这就是当你停止假装往返是免费时所付出的代价。

###### 标题（GTX 1080，三次试验，平均毫秒，比率由每试验平均值计算）：

配置（N × D，K）

基线平均（毫秒）

GPU 平均（毫秒）

加速比

10,000 × 1024，K=8

9.56

1.35

7.10×

100,000 × 768，K=8

70.66

25.70

2.75×

500,000 × 1024，K=8

483.90

69.79

6.93×

1,000,000 × 1024，K=8

977.80

114.12

8.57×

1,000,000 × 1024，K=32

973.89

125.46

7.76×

10,000 × 384，K=100

3.37

155.25

GPU 慢 46 倍

1,000,000 × 384，K=100

329.49

682.38

GPU 慢 2.07 倍

是的，你看到的数字是正确的。

前五行展示了文章的核心观点：在 K=8 时，GPU 常驻路径在所有 15 种配置中都胜出，提速比从 N=50k、D=384 时的礼貌的 2.43 倍，到 N=1M、D=1024 时的响亮的 8.57 倍不等。在 K=32 时，它在 15 种配置中的 13 种中胜出，仅在 N 最小（10k）的两种配置中失败，分别是 D=384 和 D=768，此时往返时间本身仅花费约 3–7 毫秒，GPU 的三次内核启动几乎没有空间进行摊销。当你达到实际的智能体语料库规模（N ≥ 50k）时，K=32 也轻松胜出，最高提速比达到 7.76 倍。这些显著的加速并不是“神奇内核”的加速，而是“我们没有理由再将语料库送回主机内存”的加速。GPU 一直都会在这场竞赛中胜出；它唯一失败的原因是我们一直让它做不必要的往返。

最后两行展示了这篇文章有资格自称诚实的地方。在 K=100 时，每个块的单通道冒泡排序变为 O(K²) = O(10,000) 次顺序比较，而串行合并需要遍历 P × K 个头部位置。CPU 的 std::partial_sort 是基于堆的，由编译器向量化，其复杂度为 O(N log K)，对 K=100 更友好。因此，GPU 在 15 个 K=100 配置中有 14 个失败，有时提速比为 2 倍，有时高达 46 倍。（唯一一个 GPU 仍然胜出的 K=100 配置是 N=1M、D=1024，提速比为 1.44 倍，因为此时有足够的评分工作来主导选择上限。15 行中的一行并不能算作例外；它只是一个好奇点。）这不是一个错误；这是 V1 设计原则（“可审计性优于巧妙性”）遇到的第一个具体后果。修复方法在第 9 节中，它是一个 warp 专用的锦标赛选择器，而不是一个匆忙的重构。

再补充一个诚实的说明：在这次提交的快照中，GPU 的时钟未被锁定。这意味着绝对毫秒数会随着温度和 DVFS 稍微变化，但提速比保持稳定。仓库中提供了 scripts/lock_gpu_clocks.sh 脚本，任何想要在 GTX 1080 上使用锁定时钟重现表格的人都可以使用。不言而喻，结构上的发现没有改变。

6. “好吧，但这与 FAISS / cuVS / hnswlib 有什么不同？”

这是一个非常合理的问题，值得直接回答，因为向量搜索领域有许多重叠的原语，HPC 读者会在第一条评论中提出这个问题。

• FAISS（CPU 索引）。大多数智能体框架的默认选择。优秀的库。运行在 CPU 上。智能体的每次查询都要支付一次 PCIe 往返，而这篇文章的存在就是为了消除这种往返。如果你已经在使用 IndexFlatIP，并且在检索时受 CPU 限制，那么你就是目标用户。

• FAISS（GPU 索引）。解决了 GPU 常驻问题，其内核套件比这个仓库更加成熟。CUDA-TopK-Retrieval 的目的是“我比 FAISS-GPU 更聪明”——它从未如此，也从不试图如此。其目的是在 343 行代码中展示真正酷炫的检索原语是什么样子，以及为什么在没有它的情况下，智能体流水线会感觉很慢。如果你今天需要一个生产级索引，请使用 FAISS-GPU。如果你想了解真正起作用的小热路径——一次微小的 H→D 复制，三次内核启动，两次小的 D→H 复制——请阅读这个内核。

• NVIDIA cuVS / RAFT。一个严肃的、生产级的在 GPU 上的向量搜索堆栈。更大、更快、更多算法、更多依赖项。与 FAISS-GPU 相同的警告：这个内核是教学用的 / 单二进制版本，而不是竞争对手。

• hnswlib 和相关库（近似最近邻）。它们的权衡方式完全不同 —— 它们以牺牲精确性为代价，实现对大规模语料库的亚线性查询时间。CUDA-TopK-Retrieval 是精确的暴力评分 + 选择；其加速效果完全来自于驻留时间，而不是跳过工作。

这个仓库的目的并不是“用它构建你的生产向量数据库”。其目的是：代理的检索跳转希望保留在 GPU 上，一旦接受这一点，即使是一个微小的手写内核，在一张七年前的显卡上，对于大多数实际的 K 值，也比在 CPU 上运行的暴力方法要快得多。

7. 那么，我该如何实际尝试它呢？

克隆仓库，然后使用 CMake 进行构建（-DGGML_CUDA=ON 标志与系列早期部分中 llama.cpp 的构建方式一致）：

cmake -S . -B build -DGGML_CUDA=ON -DCMAKE_BUILD_TYPE=Release
cmake --build build -j
cd build && ctest --output-on-failure

然后按照 README 的方式运行演示和基准测试：

./build/topk_demo                 # 小规模的烟雾测试（需要 GPU）
./build/topk_bench --n 20000 --d 384 --k 32 --trials 3 --warmup 1 --seed 1 --metric 0

topk_demo 是一个小规模的烟雾测试 —— 4,096 个语料库行，128 维，K=8，打印出邻居的 ID。topk_bench 是用于生成 Python 脚本所消费的 TOPK_BENCH_JSON 行的工具。对于在标准硬件上的完整 45 个配置扫描：

python3 scripts/benchmark_campaign.py.example          # 完整扫描（需要 GPU；写入 examples/benchmark-campaign-runs/run-*)

在发布级运行之前，锁定 GPU 时钟（仓库提供了脚本）：

sudo bash scripts/lock_gpu_clocks.sh

要求：Linux、CUDA 工具包、NVIDIA GPU（Pascal 或更新的型号），以及阅读一次 CMake 文件的耐心。生成的文件将位于 examples/example-run-results/（快速路径）或 examples/benchmark-campaign-runs/run-<UTC>-<pid>/（完整扫描）下，README 明确指出禁止提交 .nsys-rep 数据库 —— 仅允许 PNG 时间线。

8. 剧情反转 —— 这只是披着 CUDA 外衣的 5G 波束选择

我应该在此坦白：我并不是一个受过专业训练的“GPU 专家”。我来自电信领域 —— 5G NR，同时脚已坚定地踏入 6G 研究 —— 我一直在注意到，代理 AI 中的每一个基础设施问题，其实早在二十年前就已经在无线电层解决了。

对于没有 3GPP 背景的读者：在现代的 5G 基站中，天线不会在每个方向上均匀辐射。它会形成一个方向性波束的代码本 —— 窄的无线电能量瓣 —— 在任何给定的时刻，你的手机都会被接收功率最高的那个波束（或少数几个波束）所服务。快速选择正确的波束是无线领域最被研究的检索问题之一。UE（用户设备）会测量 L1-RSRP（每个波束的接收功率评分），然后通过 CSI 反馈信道报告回最佳的几个波束。gNB（5G 基站）使用这些报告来决定要调度哪些波束（好吧，我尽量简化了，实际情况其实非常复杂！）。

这就是 Top-K 向量搜索，用无线电的术语来表达。候选波束就是语料库。瞬时信道测量就是查询。得分是接收功率。K 是报告回传的最优波束数量。UE 在基带 DSP 层进行评分 —— 它不会将 I/Q 样本回传到中央 CPU 集群，然后让 Python 脚本判断哪个波束最好，因为如果在每毫秒的循环中这么做，会把空中接口熔化掉。

请将这两者并排对比，然后直视着我说这是两个不同的问题：

5G NR 波束选择（在 UE / 基带）

CUDA-TopK-Retrieval（在 GPU）

候选波束的码本（在配置时固定）

预嵌入块的语料库（一次性上传）

瞬时信道测量

本次跳转的查询嵌入

每个候选波束的 L1-RSRP 得分

每条语料库行的余弦 / 点积得分

将最优波束报告回 gNB

将 Top-K 行索引返回给代理

每个波束的得分位于基带 DSP，而不是主机 CPU

每行的得分位于 VRAM，而不是主机 RAM

在 CPU 上进行往返操作会熔化空中接口

在 CPU 上进行往返操作会降低代理的吞吐量

#### 给两个截然不同的受众群体的简短插话

给那些以 HPC 和 CUDA 为主的朋友们：我明白你们的想法。这里提到的数学原语并不新颖。我们都知道 cuBLAS 能更快地运行矩阵乘法，cuVS 能在数据中心规模上处理 Top-K，一个高度优化的锦标赛选择算法可以轻松击败每个块的冒泡排序。但这里的目标并不是重新发明 NVIDIA 的企业库。其价值在于零依赖的打包方式。这是一个 343 行的架构证明 —— 包括一个严格的 CPU 模拟器和一个 45 个配置的基准测试 —— 它被设计为可以在一台 8 GB 的消费级 GPU 上运行。这是那种你构建的端到端工程制品，用来证明你真正理解了硬件内存瓶颈，而不仅仅是知道如何调用框架 API。

给我的通信朋友们：如果“Top-K 向量搜索”在十分钟前听起来像一种外语，你们不是落后了，而是走在了前面。二十年来，我们的世界是 FPGAs、ASICs、PRBs 和星座图。我们优化的是频谱，而不是硅。然后 AI-RAN、NWDAF、NVIDIA Aerial 和 3GPP Rel-20 研究项目在几个月内发生得太快，现在电信行业的下一个十年需要在频谱世界和 GPU 世界之间双语精通。直觉是完全一致的。从第一个 MIMO 码本开始，你们一直在硬实时约束下进行接收端的 Top-K 处理。是同一种动物，只是换了个动物园。

9. 诚实的注意事项（因为评论即将来临）

如果你来这里是为了找出这个项目的问题 —— 恭喜你，你是本文的第一位仔细读者。直接来自 README 的 LIMITATIONS 部分和内联代码注释：

• K=100 是 V1 失败的地方。部分 Top-K 路径使用每个块单通道选择以确保可审计性；它还不是 warp 专用的锦标赛选择内核。在 K=100 时，O(K²) 的冒泡排序占主导地位，CSV 上 CPU 在 15 个 K=100 行中有 14 行表现更优（有时快 2 倍，有时快 46 倍）。K=100 时 GPU 唯一的胜利 —— N=1M，D=1024，1.44 倍 —— 是评分工作终于足够大，以淹没选择上限，而不是选择器本身变得更优。这已经被记录；修复方法是已知的后续步骤。

• 提交的收据中未锁定 GPU 时钟。提交的 environment.json 文件报告 gpu_clocks_locked: false。在消费级显卡上，开启热管理时，绝对毫秒数会发生偏移；而表头中的比率是稳定的。仓库中提供了 scripts/lock_gpu_clocks.sh 脚本（用于将 GTX 1080 的应用时钟锁定在 1607 MHz 的持久模式），供需要发布级数据的用户使用。

• 数值容差，而非精确的浮点数相等。GPU 与 CPU 的得分比较使用每个得分的小 fp32 容差；平局仍通过索引确定性地解决。这是现实中的必要性 —— fp32 的归约在 GPU 上与 CPU 上的关联方式不同 —— 而且在索引完全匹配之前，基准测试不会开始计时。

• 合成嵌入向量。基准测试使用高斯随机向量（std::normal_distribution，种子为 1）来隔离驻留时间与往返时间信号，避免内容效应的影响，并确保试验逐位可复现。真实的嵌入向量会产生更嘈杂的每轮绝对时间；PCIe 空闲时间与设备上计算之间的结构比率不会改变。

• 仅使用一种 CUDA 架构类别。所有数据均来自一款 Pascal 类 GTX 1080。在 Ada / Hopper 上，两种路径的绝对毫秒数都会减少；结构发现（PCIe 往返成本主导 CPU 侧检索）在更快的 GPU 上变得更加重要，而不是更不重要，因为内核时间的减少速度比往返时间更快。

• RAG 切片，而非完整的向量数据库。这是一个相似度 + Top-K 切片。没有压缩（PQ、OPQ），没有过滤，没有多 GPU 分片，没有同一引擎实例内的查询并发。这是代理调用的检索跳转原语 —— 而不是 FAISS-GPU 或 cuVS 的替代品。

列表中的所有内容都在路线图中。没有任何一项会改变表头结果。将其写下来的目的在于，你不需要去挖掘这些信息 —— 一旦基准测试博客文章隐藏了它的注意事项，它的数据就不再可信。

10. 总结（以及最终部分的设置）

如果你以构建智能代理流水线为生：请去查看你的检索器。打开你信任的性能分析工具。对一个工具调用进行端到端计时。如果你的 GPU 利用率在 Python 主进程进行相似性搜索时降至零，那么你已经赢得了诊断战。修复方案在 GitHub 上。

如果你以编写 CUDA 代码为生：是的，O(K²) 的冒泡排序是有意为之的。一个 warp 专用的锦标赛选择器在路线图中。

如果你以构建电信基础设施为生：是的，你抓住我了。这正是你过去二十年在 DSP 代码中编写的相同基带检索原语。AI 行业只是改变了术语；数学没有改变。

#### 接下来的内容：如何防止你的代理互相创伤倾倒

CUDA-TopK-Retrieval 证明你可以停止将每个检索跳转都发送到 GPU。但如果你重新阅读注意事项 #1 以及第 5 节中 K=100 的行，你已经发现了下一个瓶颈：每个查询的工作是相互独立的。

每次检索跳转都从冷启动开始。确实，语料库在设备上。但代理的状态 —— 其先前决策的嵌入向量，以及它自然会回溯关注的键和值 —— 在每次交接时都会被丢弃并重新构建。GPU 保持温暖；代理的记忆却保持冰冷。

这在单次 RAG 步骤中是可行的。但一旦你运行这个系列所设计的工作负载——跨多个专业化代理的多跳推理——情况就会变得复杂。在那样的规模下，你不再关心“我们是否将检索保留在 GPU 上”，而是开始关注单次内核无法回答的问题：

• 当代理 A 将任务交给代理 B 时，B 是否可以继续使用 A 积累的上下文，而不是从零开始？

• 每次跳跃的持久状态可以小到什么程度，同时仍然有用？

• 在下一个代理的 GPU 上恢复该状态的延迟成本是多少？

• 我们如何确保任务交接时不丢失信息？

要回答这些问题，你必须像 CPU 在 cudaMemcpy 过程中所做的那样：耐心地坐着，等待下一部分的到来。

我们将在第 4 部分，也就是最后一部分再见。

免责声明：本文中的插图是使用 AI（Claude Opus 4.8）生成的。这些插图是示意性的，而非照片，图像中可见的任何标签都是风格化的，而非权威性的——请参阅文章正文和代码本身，以获取精确的功能名称、指标值和架构细节。

作者

查看 Anubhab Banerjee 的所有文章

,

人工智能

CUDA

深度学习

编辑推荐

分享这篇文章

• 在 Facebook 上分享

• 在 LinkedIn 上分享

• 在 X 上分享

Towards Data Science 是一个社区出版物。提交你的见解，以触达全球受众，并通过 TDS 作者支付计划获得报酬。

将 href 更新为你的实际提交 URL

为 TDS 写作

✦ end CTA ✦