HierarchicalKV: NVIDIA的高性能分层键值存储解决方案

HierarchicalKV

HierarchicalKV简介

HierarchicalKV是NVIDIA Merlin项目的一个重要组成部分,专门为推荐系统设计的高性能分层键值存储解决方案。它的核心功能是能够在GPU的高带宽内存(HBM)和主机内存中存储键值对(特征-嵌入),以满足大规模推荐系统的需求。同时,HierarchicalKV也可以作为通用的键值存储使用。

HierarchicalKV架构图

主要特点和优势

支持HBM和主机内存混合存储
通过绕过CPU来提高性能,减少通信开销
实现基于LRU或自定义策略的表大小限制
高工作状态负载因子,接近1.0

HierarchicalKV的设计理念使NVIDIA GPU更适合训练大型和超大型的搜索、推荐和广告模型。它简化了构建、评估和提供复杂推荐模型的常见挑战。

核心设计理念

HierarchicalKV的核心设计理念包括:

桶本地有序
分别存储键和值
将所有键存储在HBM中
内置和可定制的淘汰策略

这些设计理念使HierarchicalKV能够高效地管理大规模的键值数据,同时保持灵活性和可扩展性。

API文档和使用方法

HierarchicalKV提供了丰富的API,主要包括以下几个核心类和结构体:

class HashTable：哈希表的主要实现
class EvictStrategy：淘汰策略的实现
struct HashTableOptions：哈希表的配置选项

详细的API文档可以在官方文档中找到。

使用示例

以下是一个简单的使用示例:

#include "merlin_hashtable.cuh"

using TableOptions = nv::merlin::HashTableOptions;
using EvictStrategy = nv::merlin::EvictStrategy;

int main(int argc, char *argv[])
{
  using K = uint64_t;
  using V = float;
  using S = uint64_t;
  
  // 1. 定义表并使用LRU淘汰策略
  using HKVTable = nv::merlin::HashTable<K, V, S, EvictStrategy::kLru>;
  std::unique_ptr<HKVTable> table = std::make_unique<HKVTable>();
  
  // 2. 定义配置选项
  TableOptions options;
  options.init_capacity = 16 * 1024 * 1024;
  options.max_capacity = options.init_capacity;
  options.dim = 16;
  options.max_hbm_for_vectors = nv::merlin::GB(16);
  
  // 3. 初始化表内存资源
  table->init(options);
  
  // 4. 使用表进行操作
  
  return 0;
}

淘汰策略

HierarchicalKV提供了多种淘汰策略,包括:

LRU (最近最少使用)
LFU (最不经常使用)
EpochLRU
EpochLFU
自定义策略

这些策略通过score来定义每个键的重要性,分数越大,越不容易被淘汰。

性能评测

HierarchicalKV在不同场景下展现了优秀的性能。以下是一些性能测试结果:

纯HBM模式 (dim = 64, 容量 = 64 Million-KV, HBM = 16 GB, HMEM = 0 GB)

负载因子	insert_or_assign	find	find_or_insert	assign	find*	find_or_insert*
0.50	0.834	1.982	1.113	1.499	3.950	1.502
0.75	0.801	1.951	1.033	1.493	3.545	1.359
1.00	0.621	2.021	0.608	1.541	1.965	0.613

注: 吞吐量单位为十亿KV/秒

HBM+HMEM混合模式 (dim = 64, 容量 = 512 Million-KV, HBM = 32 GB, HMEM = 96 GB)

负载因子	insert_or_assign	find	find_or_insert	assign	find*	find_or_insert*
0.50	0.049	0.069	0.049	0.069	3.484	1.370
0.75	0.049	0.069	0.049	0.069	3.116	1.242
1.00	0.047	0.072	0.047	0.070	1.771	0.607

这些性能数据展示了HierarchicalKV在不同负载和存储模式下的优秀表现。

构建和使用限制

HierarchicalKV主要是一个仅头文件的库,但也提供了基准测试和单元测试的二进制文件。构建环境需要满足以下要求:

CUDA版本 >= 11.2
NVIDIA GPU with compute capability 8.0, 8.6, 8.7 or 9.0
GCC支持C++17标准或更高
Bazel版本 >= 3.7.2 (仅限Bazel编译)

使用限制:

key_type必须是int64_t或uint64_t
score_type必须是uint64_t

贡献和支持

HierarchicalKV是由NVIDIA Merlin团队和NVIDIA产品最终用户共同维护的开源项目。它欢迎公众贡献、错误修复和文档改进。如果您遇到任何问题或有任何疑问,可以在GitHub Issues页面提交问题。

总结

HierarchicalKV为大规模推荐系统提供了一个强大的分层键值存储解决方案。通过利用GPU的高带宽内存和主机内存,它能够有效地管理海量的特征-嵌入对,同时保持高性能和灵活性。其丰富的API、多样的淘汰策略和优秀的性能使其成为构建现代推荐系统的重要工具。

随着推荐系统规模的不断增长和复杂性的提高,HierarchicalKV将继续发挥重要作用,帮助开发者更轻松地构建和优化大规模推荐模型。未来,我们可以期待看到更多基于HierarchicalKV的创新应用和性能优化,进一步推动推荐系统技术的发展。