tables-overview

发表于 分类于
NV Embedding Cache SDK 中 8 种 Table 实现的完整全览，涵盖 GPU 缓存表、CPU 哈希表、远程存储表的设计与实现。

Table 实现全览

概述

Table 是 NV Embedding Cache SDK 中所有存储后端的统一抽象基类，定义了以下纯虚接口：

class Table {
  virtual ~Table() = default;
  virtual void clear(context_ptr_t& ctx) = 0;
  virtual void erase(context_ptr_t& ctx, int64_t n, const void* keys) = 0;
  virtual void find(context_ptr_t& ctx, int64_t n, const void* keys,
                    max_bitmask_repr_t* hit_mask, int64_t value_stride,
                    void* values, int64_t* value_sizes) const = 0;
  virtual void insert(context_ptr_t& ctx, int64_t n, const void* keys,
                      int64_t value_stride, int64_t value_size, const void* values) = 0;
  virtual void update(...) = 0;
  virtual void update_accumulate(...) = 0;
  virtual context_ptr_t create_execution_context(...) = 0;
  virtual void reset_lookup_counter(...) = 0;
  virtual void get_lookup_counter(...) = 0;
  virtual bool lookup_counter_hits() = 0;
  virtual int32_t get_device_id() const = 0;
  virtual int64_t get_max_row_size() const = 0;
  virtual int64_t get_invalid_key() const = 0;
};

项目实现了 8 种表，按存储位置和底层技术可分为：

类别 表实现 存储位置 数据持久化 依赖
GPU 缓存表 GpuTable GPU 显存 ❌ cuEmbed + EmbedCacheSA
CPU 线性表 LinearHostTable CPU 内存（用户分配） ❌ 无（memcpy 直接读写）
CPU 哈希表（内置） STLContainerTable CPU 内存 ❌ std::unordered_map
CPU 哈希表（Abseil） AbseilFlatMapTable CPU 内存 ❌ Abseil Swiss Table
CPU 哈希表（PHMAP） PHMapFlatMapTable CPU 内存 ❌ Parallel Hashmap
CPU GPU 哈希表 NvhmMapTable CPU 内存（nvhashmap） ❌ NVIDIA nvhashmap
远程缓存 RedisClusterTable 远程 Redis 集群 ✅ 持久化 hiredis + redis++
磁盘存储 RocksDBTable 本地磁盘 ✅ 持久化 Facebook RocksDB

一、Table 基类接口详解

核心操作契约

find(ctx, n, keys, hit_mask, value_stride, values, value_sizes)

  • 输入：hit_mask 中 bit=0 的键需要被查找，bit=1 的键应跳过
  • 输出：对于 bit=0 且在表中找到的键，将其 bit 设为 1，并将值写入 values 的对应位置
  • 紧凑写入：如果 values 是 nullptr，则只做计数/键存在性检查

insert(ctx, n, keys, value_stride, value_size, values)

  • 向表中插入键值对。如果键已存在，行为取决于实现（替换或忽略）
  • 最佳努力操作：如果触发了超额淘汰，新插入的键可能立即被淘汰

update(ctx, n, keys, value_stride, value_size, values)

  • 只覆盖表中已存在的键，不存在的键被忽略
  • GpuTable 如果配置了 uvm_table 且未禁用 UVM 更新，会同时写入 UVM 后备表

update_accumulate(ctx, n, keys, update_stride, update_size, updates, update_dtype)

  • 梯度累积模式：table[key] += updates[i]
  • update_dtype 可以不同于表的存储类型（如 fp32 梯度累积到 fp16 表）

命中率计数

方法 作用
reset_lookup_counter(ctx) 重置计数器
get_lookup_counter(ctx, counter) 读取计数器值（GPU 表需先 cudaStreamSynchronize）
lookup_counter_hits() false = 计数的是 miss，true = 计数的是 hit

二、GpuTable — GPU 集合关联缓存表

头文件：include/gpu_table.hpp | 实现：src/gpu_table.cu

GpuTable 是 NVE 的核心 GPU 端表实现，内部封装了 EmbedCacheSA<KeyType, KeyType> ——一个集合关联（Set-Associative）软件管理缓存。

架构层次

GpuTable
  └── EmbedCacheSA<KeyType, KeyType>  (GPU 集合关联缓存)
        ├── EmbedCacheSA 基类: 标签匹配 + 数据区 + 替换策略
        ├── CacheSADeviceModify: 修改操作在 GPU 上执行
        └── CacheSAHostModify: 修改操作在 CPU 上执行（调试用途）

查找的四种 kernel 模式

当 uvm_table 非空时，run_find_uvm() 根据键数量选择执行策略：

Kernel 模式 阈值条件 行为
LookupUVM 默认：num_keys < 1M 一次 cache->lookup 调用，cache miss 时直接从 UVM table 读取
SortGather 默认：num_keys ≥ 1M 先排序键（去重），大块连续读取 UVM，改善 page fault 局部性
PipelineGather 用户明确指定 流水线 gather kernel，分 tile 处理
DynamicKernel — 暂未实现

Insert 的直方图优化

GpuTable 的 insert 操作使用 DefaultGPUHistogram 在 GPU 上对键做频率统计和去重：

输入: [1, 2, 1, 3, 2, 1, 1]  (7 个键)
  ↓
DefaultGPUHistogram（GPU kernel）
  ↓
输出: 键 [1, 2, 3] 按频率排序 (4, 2, 1)
  ↓
cache->insert() → 高频键优先插入 cache

频率越高的键在 insert 中获得越高优先级，在集合关联缓存中更不容易被淘汰。

三、LinearHostTable — CPU 线性表

头文件：include/linear_host_table.hpp | 实现：src/linear_host_table.cpp

设计理念

LinearHostTable 是最简单的 CPU 端表实现：嵌入向量存储在用户提供的连续内存缓冲区中，通过 key * row_size 直接计算地址后 memcpy 读写。无哈希、无索引、无淘汰——这与 GPUEmbeddingLayer 在 GPU 端的角色完全对称。

实现

// 查找: 最简单的 memcpy
void find(...) {
    thread_pool->execute_n(0, num_tasks, [=](int64_t idx) {
        for each key 跳过 hitmask 已标记的:
            memcpy(output + k * output_stride,
                   src + keys[k] * row_size_in_bytes,
                   row_size_in_bytes);
    });
    memset(hit_mask, 0xff, ...);
}

适用场景

  • 与 GpuTable 配合，UVM 后备表的 CPU 端视图
  • 纯 CPU 推理场景，嵌入表可完全装入系统内存
  • 作为 HierarchicalEmbeddingLayer 的中间层

四、STLContainerTable — 内置 CPU 哈希表

头文件：include/stl_map_backed_table.hpp | 实现：src/stl_map_backed_table.cpp

架构

STLContainerTable 是所有基于 STL 风格容器的 CPU 哈希表实现的公共基类。它使用 std::unordered_map<KeyType, char*> 作为底层映射，核心价值在于它实现了一套完整的内存管理、并发控制、超额淘汰基础设施。

STLContainerTable  (模板基类)
  ├── 持有 std::vector<Partition>
  │     └── Partition:
  │           ├── std::shared_mutex  (读写锁)
  │           ├── std::unordered_map<KeyType, char*>  (键 → 槽指针)
  │           └── SlotBuffer (手动管理的内存池)
  │
  ├── 查找: 分片并发 + hitmask 过滤 + 软件预取
  ├── 插入: 分片加锁 + 槽分配 + 超额淘汰
  ├── 更新: 分片加锁 + 直接覆盖
  ├── 累积: 分片加锁 + 原子累加 + inline meta
  └── 淘汰: 三种策略 → EvictRandom / EvictLRU / EvictLFU

分片与并发

所有操作按 partitioner(key) & (num_partitions - 1) 分配到对应分片：

  • 查找：每个分片加共享锁（shared_lock），多线程并行
  • 修改：每个分片加独占锁（unique_lock），保证写一致性

内存槽管理

STLContainerTable 不把值直接存在 std::unordered_map 的 value 中，而是手动管理一个连续的 slot buffer：

Slot Buffer 布局:
┌────────────────────────────────────────────────┐
│ Row 0  │  value_bytes  │  meta_bytes  │ padding│
│ Row 1  │  value_bytes  │  meta_bytes  │ padding│
│ ...... │  .............│ .............│ .......│
│ Row N  │  value_bytes  │  meta_bytes  │ padding│
└────────────────────────────────────────────────┘
  • value_bytes：嵌入向量数据
  • meta_bytes：inline metadata（LRU 时间戳或 LFU 频率计数）
  • std::unordered_map 只存 key → char*（指向 slot 起始位置）

这种设计将哈希表索引和数据存储分离，元数据跟随数据一起移动，淘汰时只需排序 meta 即可决定谁被淘汰。

超额淘汰（Overflow）

当插入导致分片容量超过阈值时触发淘汰：

overflow_margin = 512（额外条目数）
resolution_margin = 0.2（淘汰比例 20%）

触发条件: partition.size() > initial_capacity + overflow_margin
淘汰过程:
  1. 收集所有 key + meta
  2. 按 meta 排序:
     - EvictRandom: 随机打乱
     - EvictLRU: 按时间戳排序（最久未使用优先）
     - EvictLFU: 按频率排序（最少使用优先）
  3. 淘汰 resolution_margin 比例的条目
  4. 释放对应的 slot 内存

模板分发

produce() 方法按 mask_size → key_size → overflow_handler → partitioner 四维 switch 分发，生成约 192 种编译时类型特化。

内置别名

系统中预注册了 6 个别名指向同一个基于 std::unordered_map 的 STLMapTableFactory：

{"stl_map", "map", "stl_unordered_map", "unordered_map", "stl_umap", "umap"}

五、AbseilFlatMapTable — Abseil Swiss Table 插件

头文件：plugins/abseil/include/abseil_flat_map_table.hpp | 注册名：abseil_flat_map

继承 STLContainerTable，将底层容器替换为 absl::flat_hash_map<KeyType, char*>。absl::flat_hash_map 使用 Google 的 Swiss Table 算法（开放寻址 + 元数据字节 + SIMD 查找），比 std::unordered_map 有更好的内存局部性和查找性能。

维度 STLContainerTable AbseilFlatMapTable
底层容器 std::unordered_map（链式哈希） absl::flat_hash_map（Swiss Table）
额外配置 无 initial_capacity 用于 reserve()
链接依赖 无 absl_raw_logging_internal + absl_hash

六、PHMapFlatMapTable — Parallel Hashmap 插件

头文件：plugins/phmap/include/phmap_flat_map_table.hpp | 注册名：phmap_flat_map

实现方式与 Abseil 插件完全相同——继承 STLContainerTable，底层替换为 phmap::flat_hash_map<KeyType, char*>。PHMAP 使用与 Abseil 类似的开放寻址算法，但实现更轻量，无外部链接依赖（header-only）。

七、NvhmMapTable — NVIDIA nvhashmap 插件

头文件：plugins/nvhm/include/nvhm_map_table.hpp | 注册名：nvhm_map

与 STLContainerTable 的区别

NvhmMapTable 不继承 STLContainerTable，而是直接继承 HostTable<NvhmMapTableConfig>，直接使用 NVIDIA nvhashmap 的 map 接口：

template <typename MaskType, typename MapType, typename PartitionerType>
class NvhmMapTable final : public HostTable<NvhmMapTableConfig> {
    struct Partition {
        mutable std::shared_mutex read_write;
        map_type map;  // nvhm::map<KeyType, ValueType, ...>
    };
    std::vector<Partition> parts_;
};

配置

字段 默认 说明
kernel_size nvhm 默认 内核大小 11024（对应 81024B entry）
key_fetch_queue_length 8 软件预取管道深度 {0,1,2,4,8}
prefetch_values true 是否发出值预取指令
minimize_psl false 缩短 probe search length
auto_shrink false 大量淘汰后自动缩容

查找中的软件预取

nvhashmap 的 ring_prefetch_queue 实现了软件流水线预取：

for (int i = 0; i < KeyFetchQueueLength; i++) {
    __builtin_prefetch(next_key_address);
}

在 find_() 模板方法中，通过 KeyFetchQueueLength 和 PrefetchValues 两个编译时参数控制预取行为，switch 生成 5 × 2 = 10 种变体。

模板分发

NvhmMapTable 的模板分发维度最多：

mask_size × key_size × overflow_handler × kernel_size × minimize_psl × auto_shrink × partitioner

→ 约 6144 种编译时类型组合。

八、RedisClusterTable — Redis 远程表

头文件：plugins/redis/include/redis_cluster_table.hpp | 注册名：redis_cluster

架构

RedisClusterTable 是 NVE 中唯一基于网络通信的表实现。它通过 hiredis（C 客户端）和 redis++（C++ 客户端）连接到外部 Redis Cluster，将嵌入向量以序列化形式存储在 Redis 中。

配置

字段 默认 说明
address localhost:6379 Redis 集群任意节点地址
user_name default 用户名
password 空 密码
connections_per_node 5 每节点最大并行连接数
connection_lifetime 180s 连接生命周期
max_batch_size 16384 批量操作最大大小
num_partitions 1 并行分片数（2 的幂）

查找实现

auto reply = redis_cluster_->pipeline(keys.begin(), keys.end(),
    [](const auto& key) { return sw::redis::command("GET", key); });

对于每个分片，使用 redis++ 的 pipeline 功能批量发送 GET 命令，减少网络往返次数。

链接依赖

hiredis + hiredis_ssl + redis++_static + ssl + crypto

九、RocksDBTable — 磁盘持久化表

头文件：plugins/rocksdb/include/rocksdb_table.hpp | 注册名：rocksdb

架构

RocksDBTable 通过 Facebook RocksDB 提供本地磁盘持久化存储。嵌入向量以 key → value 形式存储在 RocksDB 的 Column Family 中。

配置

字段 默认 说明
column_family default RocksDB Column Family
max_batch_size 16384 批量操作最大大小
verify_checksums true 是否校验 checksum
path /tmp/rocksdb 数据库文件路径
read_only false 只读模式（允许多进程并发读取）
num_threads 16 后台线程数

查找实现

db_->MultiGet(read_opts_, col_families_, keys_slices, &values, &statuses);

使用 RocksDB 的 MultiGet 实现批量读取，单次调用返回多个键的状态。

十、所有 Table 的对比总表

维度 GpuTable LinearHostTable STLContainer AbseilFlatMap PHMapFlatMap NvhmMapTable RedisCluster RocksDB
存储位置 GPU 显存 CPU 线性内存 CPU 堆内存 CPU 堆内存 CPU 堆内存 CPU 堆内存 远程集群 本地磁盘
数据持久化 ❌ ❌ ❌ ❌ ❌ ❌ ✅ ✅
查找算法 集合关联 tag key * stride unordered_map O(1) Swiss Table O(1) PHMAP O(1) nvhashmap O(1) Redis GET LSM-Tree
并发粒控 CUDA stream 线程池并行 分片读写锁 分片读写锁 分片读写锁 分片读写锁 pipeline + 连接池 RocksDB 内部
淘汰策略 集合关联 LRU 无（固定数组） EvictRandom/LRU/LFU 同左 同左 同左 无（Redis 自管） 无（磁盘存储）
是否插件 核心库内置 核心库内置 核心库内置 ✅ abseil 插件 ✅ phmap 插件 ✅ nvhm 插件 ✅ redis 插件 ✅ rocksdb 插件
启动依赖 CUDA + nvhashmap 无 无 absl PHMAP header-only nvhashmap hiredis + redis++ librocksdb
JSON 实现名 — — stl_map/umap abseil_flat_map phmap_flat_map nvhm_map redis_cluster rocksdb
典型扮演角色 层级首表 层级中/末表 层级中层 层级中层 层级中层 层级中层 层级末表 层级末表

十一、在嵌入层中的角色总结

HierarchicalEmbeddingLayer:
  tables[0] ─── GpuTable (GPU cache, 最快层)
  tables[1] ─── NvhmMapTable / AbseilFlatMapTable / STLContainerTable (CPU hash, 中间层)
  tables[2] ─── RedisClusterTable / RocksDBTable (远程/持久化, 末层)

LinearUVMEmbeddingLayer:
  GpuTable.cache_ = EmbedCacheSA (GPU cache)
  GpuTable.config_.uvm_table = UVM 线性内存 (硬件透明回退)

GPUEmbeddingLayer:
  无 Table (直接 cuEmbed Forward kernel, 线性 GPU 内存)