HierarchicalKV-ascend

    

前言

本项目参考NVIDIA HierarchicalKV(Beta)，实现了Ascend950系列架构下HKV HashTable的移植样例。

概述

HierarchicalKV-ascend（下称HKV）是一个基于昇腾平台，面向推荐系统的高性能key-value存储加速库。 在推荐系统中，HKV提供了大容量、高性能的动态Embedding表的增删改查能力。

关键特性

• 支持HBM内动态扩容
• 支持可定制的淘汰策略
• keys和values存储分离，keys仅存储于HBM

分级存储，使用HOST DDR存储values的能力正在开发中，敬请期待。

API说明

核心类、结构体如下（命名空间均为npu::hkv::）

• struct EvictStrategy : 淘汰策略
• class HashTableBase : HashTable接口类
• class HashTable : HashTable实现类

核心功能均由HashTable各接口调用算子完成。

API矩阵

接口名称	描述
init	使用HashTableOptions初始化HashTable
insert_or_assign	插入或更新keys，若桶内数据已满则会淘汰分数最低的key
insert_and_evict	插入keys，并返回被淘汰的keys
find_or_insert*	查找或插入keys，并返回values的HBM地址
assign_scores	更新keys对应的scores
find*	查找keys，并返回values的HBM地址
find_and_update*	查找keys并更新scores，返回values的HBM地址
reserve	扩容至指定大小
export_batch	批量导出有效keys
save	导出keys并写入文件
load	读取keys并插入
clear	清空所有keys
set_global_epoch	设置epoch淘汰策略的epoch
set_max_capacity	设置最大容量
empty	返回当前表是否为空
size	返回有效keys数
load_factor	返回负载率
capacity	返回当前容量
dim	返回value的dim值
max_bucket_size	返回最大桶深度
bucket_count	返回当前桶数

以上为当前支持的全量接口，其中unique_key参数当前不支持设置为false，请保证保证输入的key已去重。其他HashTableBase中的接口正在开发中，敬请期待。

接口参数说明与约束请参考HierarchicalKV API

淘汰策略

引入分数（score）来定义每个key的重要性，分数越大，key越重要，被淘汰（evict）的可能性就越小。只有当一个桶（bucket）满时才会发生淘汰操作。score_type必须是uint64_t类型。更多详情请参考class EvictStrategy。

策略名称	score意义
Lru	设备时钟（纳秒级），与主机时钟略有差异
Lfu	调用插值类的API，通过scores参数指定的频率分数，累加到对应key的频率分数上
EpochLru	高32位是通过set_global_epoch设置的全局epoch， 低32位等于(device_clock >> 20) & 0xffffffff，精度约为1毫秒
EpochLfu	高32位是通过set_global_epoch设置的全局epoch， 低32位是频率分数， 达到最大值0xffffffff后频率将保持恒定
Customized	由调用者调用插值类的API，通过scores传入的分数

• 注意：
  • 插值类API：insert_or_assign, insert_and_evict, find_or_insert, accum_or_assign, assign_scores以及find_or_insert。

配置选项

建议保留以*结尾的选项的默认配置。

配置名称	类型	默认值	描述
init_capacity	size_t	0	初始容量
max_capacity	size_t	0	最大容量
max_hbm_for_vectors	size_t	0	vectors可占用的最大HBM（单位字节）
max_bucket_size*	size_t	128	桶深度
dim	size_t	64	vectors的维度
max_load_factor*	float	0.5f	触发rehash的最大负载率
device_id	int	-1	设备id，为-1时，设置为aclrtGetDevice返回的设备id
num_of_buckets_per_alloc	int	1	申请桶内存时，单次申请的桶数

以上为当前支持的全量选项，部分HashTableOptions中的选项正在支持中，敬请期待。

补充说明：

• 由于当前不支持扩容至HOST DDR特性，请将max_hbm_for_vectors值设置足够大以存储全部vectors。
• max_bucket_size需满足max_bucket_size >= 16。

如何使用

#include "hkv_hashtable.h"

using TableOptions = npu::hkv::HashTableOptions;
using EvictStrategy = npu::hkv::EvictStrategy;

int main(int args, char *argv[])
{
  using K = uint64_t;
  using V = float;
  using S = uint64_t;

  // 1. 指定淘汰策略，并创建表对象
  using HKVTable = npu::hkv::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 = npu::hkv::GB(16);

  // 3. 调用init接口进行初始化
  table->init(options);

  // 4. 使用table进行相关操作
  // 5. 结束
  return 0;
}

使用约束

• key_type必须为uint64_t或int64_t
• score_type必须为uint64_t

目录结构介绍

├── hkv                       // HKV HashTable根目录
|   ├── 3rdparty              // 第三方依赖库
│   ├── benchmark             // 算子性能评估工程
│   ├── cmake                 // 编译工程文件
│   ├── hkv_hashtable         // HKV算子kernel实现
│   ├── include               // HKV HashTable头文件目录
│   │   ├── hkv_hashtable.h   // HKV HashTable对外主接口头文件
│   ├── tests                 // HKV算子测试工程和测试代码
│   ├── CMakeLists.txt        // 编译工程文件
│   ├── main.cpp              // HKV应用最小demo
│   └── run.sh                // 编译运行脚本

编译构建

构建产物

构建产物包含头文件（*.h）、动态库（*.so）以及benchmark、demo和单元测试的可执行文件。

构建依赖

• CANN开发套件包

• googletest

快速构建

下载代码

git clone --recurse-submodules https://gitcode.com/Ascend/HierarchicalKV-ascend
cd HierarchicalKV-ascend

依赖配置

• 配置CANN环境变量

  请根据当前环境上CANN开发套件包的安装方式，选择对应配置环境变量的命令。

  • 默认路径，root用户安装CANN软件包

    export ASCEND_INSTALL_PATH=/usr/local/Ascend/ascend-toolkit/latest

  • 默认路径，非root用户安装CANN软件包

    export ASCEND_INSTALL_PATH=$HOME/Ascend/ascend-toolkit/latest

  • 指定路径install_path，安装CANN软件包

    export ASCEND_INSTALL_PATH=${install_path}/ascend-toolkit/latest

编译运行

bash run.sh -v [SOC_VERSION] -d [DEVICE_ID]

• SOC_VERION：AI处理器的型号，默认值为Ascend950PR_9579。在服务器执行npu-smi info -t board -i id -c chip_id命令进行查询，获取Chip Name和 NPU Name信息，实际配置值为{Chip Name}_{NPU Name}。例如Chip Name取值为Ascend950PR，NPU Name取值为957c，实际配置值为Ascend950PR_957c。其中：
  • id：设备id，通过npu-smi info -l命令查出的NPU ID即为设备id。
  • chip_id：芯片id，通过npu-smi info -m命令查出的Chip ID即为芯片id。
• DEVICE_ID：样例执行的npu的卡号，默认值为0。
• 其他参数说明请使用-h获取帮助信息。

示例：

bash run.sh -v Ascend950PR_9579 -d 3