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;
}