[SIGMOD'27] HierarchicalKV: A GPU Hash Table with Cache Semantics for Continuous Online Embedding Storage

#摘要

传统的 GPU 哈希表会保留每一个插入的键 – 这种字典式的设计方式会浪费宝贵的 HBM 资源,因为当哈希表的大小超过单个 GPU 的容量时,这种处理方式就会变得不可行。我们打破了这一假设,采用了基于缓存语义的存储方式,其中策略驱动的淘汰操作成为了核心功能。我们提出了 HierarchicalKV(HKV)这一通用 GPU 哈希表库,它的正常运行模式基于缓存语义:每次完整的更新或插入操作都会通过淘汰或拒绝操作来直接修改数据,而不是通过重新哈希处理或因容量限制而导致失败。 HKV 结合了四种核心机制:与缓存行对齐的桶结构、基于评分的在线更新操作、基于评分的动态双桶选择机制,以及三重组并发处理机制。此外,HKV 还采用了分层键值分离的设计,从而能够在超出 HBM 容量的情况下实现扩展。在 NVIDIA H100 NVL GPU 上,HKV 每秒能够处理高达 39 亿个键值对。在负载因子为 0.50–1.00 的情况下,性能表现有 5% 的波动。与 WarpCore 相比,HKV 在吞吐量上高出 1.4 倍;在基于间接处理的 GPU 基准测试中,性能提升可达 2.6–9.4 倍。自 2022 年 10 月开源以来,HKV 已被集成到多个开源推荐系统中。

#1. 引言

现代的推荐、搜索和广告系统(Naumov 等人,2019 年)依赖于深度学习模型。这些模型中的嵌入表负责将稀疏特征映射到密集向量上,这一过程占据了大量的内存资源。通常,这些嵌入表占用的内存远远超过单颗 GPU 所能承受的最大内存容量(例如,1B 个 64-dim float32 key 会占用 ~256GB 的内存,而一块 NVIDIA H100 NVL 卡上只有 96GB)。持续的在线训练进一步增加了内存使用的压力,因为需要不断处理新的键信息,而内存预算却始终有限。

所有被广泛使用的 GPU 哈希表 – 如 cuCollections(NVIDIA,2020 年)、WarpCore(Jünger 等人,2020 年)、BGHT(Awad 等人,2023 年)、WarpSpeed(McCoy 和 Pandey,2025 年)以及 Hive(Polak 等人,2025 年)-- 都采用字典语义模型:每个插入的键都必须被保留下来。当负载因子接近 1.0 时,查询链会延长,吞吐量会下降,系统不得不进行重新哈希处理,否则就会失败。这些限制是字典语义模型固有的:只要要求保留每个键,那么在高负载情况下的性能下降以及外部维护工作就难以避免。我们的测量结果显示,当 λ 接近 1.0 时,查找吞吐量会下降 31% 到 100%(参见第 5.2 节),这使得字典语义的哈希表不适合用于需要持续填充的嵌入缓存系统。

我们注意到,将存储功能建模为 cache 比建模为 dictionary 更为合适:幂律访问模式意味着只需保留高价值的条目,而淘汰低价值的条目则有助于保持模型的稳定性。由于内存资源有限,淘汰低价值条目成为必然且频繁发生的操作。这种从字典语义到缓存语义的转变,从根本上改变了设计约束:完整的表不再被视为错误状态,而是正常的运行状态。就像 B 树索引(字典语义)与缓冲池(缓存语义)之间的区别导致了数据库系统中不同的并发处理与替换策略一样,同样的语义划分也导致了 GPU 哈希表设计的根本差异。采用缓存语义开辟了新的设计空间,但同时也带来了四个具体挑战(参见第 2.4 节):

  1. 将每个键的查找操作限制在固定的 GPU 内存事务数量之内。
  2. 在不重复处理的情况下完成全部更新操作(C2)
  3. 在有限关联性的情况下保留高价值的数据条目(C3)
  4. 在混合工作负载环境下将结构化的写入操作与非结构化的写入操作分开处理(C4)。

为了应对这些挑战,我们提出了 HierarchicalKV (HKV) – 首个 general-purpose GPU 哈希表库。该库的满容量操作契约采用缓存语义:每次对满 Bucket 进行插入或更新操作时,都会通过驱逐或拒绝来直接修改数据,而无需进行 Rehash 处理或外部维护。 HKV 将四种协同设计的机制整合为一套完整的架构体系。

  1. 单桶限制型 缓存行对齐的 桶结构(§3.2):这种结构将每个键的所有候选值压缩为 128 字节的摘要,这些摘要仅占用一个 GPU 的 L1 缓存行空间。这样一来,在单次内存事务中就能实现每个桶的完整检索,无论是在单桶模式下还是双桶模式下都是如此。
  2. 行内的 评分驱动的 upsert 操作语义(§3.3):桶级局部评分比较、并发控制以及比较与交换(CAS)操作直接集成到插入路径中,因此整个表的插入操作可以在内部完成,无需额外的驱逐流程。
  3. 基于得分的动态双桶选择算法(§3.4):该算法扩展了基于二选一策略的优化方法(Azar 等人,1994 年;Mitzenmacher,2001 年),将这种策略从负载均衡扩展到 GPU 哈希表中的淘汰策略优化。该算法能够解决在单桶模式下因生日悖论导致提前淘汰未使用的内存块的问题,同时能在 λ = 1.0 情况下实现 99.4%的最高得分保留率。
  4. 三重组并发机制(§3.5):将读取器、更新器和插入器角色分离,并采用 CPU-GPU 双层锁定机制。所有结构性的修改都集中在插入器内核中完成,因此读取器和更新器内核无需进行 CAS 操作或驱逐逻辑处理——这大大降低了每个内核的复杂性,并使得每种访问模式都能得到独立优化。

作为一种可扩展性的实现机制,分层键值分离结构(§3.6)将容量扩展到了 HBM 之外:键、摘要和分值分别存储在 HBM 中,而值则通过基于位置的地址映射被存储到固定的主机内存(HMEM)中。这四种核心机制协同工作,共同实现了三种特性,而这些特性是之前所有 GPU 哈希表所不具备的:在每次加载时,每个桶的查找操作都会执行固定的操作;能够在不重新哈希的情况下实现全容量就地 upsert 操作;以及能够同时执行读取、更新和插入操作,并且这些操作由具有角色隔离功能的 Kernel 来完成。分层键值分离结构还适用于混合 HBM 和 HMEM 的部署场景,同时仍然将键值处理任务保留在 GPU 上。

我们对 NVIDIA H100 NVL GPU 的评估表明,HKV 的查找吞吐量高达 3.9 B-KV/s,其在负载因子为 0.50–1.00 的情况下表现稳定,变化率仅为 5%。其查找吞吐量相当于 WarpCore 的 1.4×,而在基于间接寻址的设计中,吞吐量可达到 2.6–9.4×(即使用(键,索引)对并结合单独的值收集方式:BP2HT、BGHT、cuCollections))。自 2022 年 10 月开源以来,HKV 已集成到 NVIDIA Merlin HugeCTR、TFRA 以及 NVIDIA RecSys 示例中。本文介绍了 HKV 缓存语义架构的设计原理、正式特性以及全面的评估结果。

本文的贡献包括:

#