SSD 嵌入算子

CUDA 算子

enum RocksdbWriteMode

rocksdb 写入模式

在 SSD 卸载中，每次训练迭代中有 3 次写入 FWD_ROCKSDB_READ：缓存查找会将未缓存的数据从 rocksdb 移动到 fwd 路径上的 L2 缓存

FWD_L1_EVICTION：L1 缓存驱逐会将数据驱逐到 fwd 路径上的 L2 缓存中

BWD_L1_CNFLCT_MISS_WRITE_BACK：L1 冲突未命中会将数据插入到 L2 中，以便在 bwd 路径上进行嵌入更新

以上所有 L2 缓存填充都可能在 L2 缓存满时触发 rocksdb 写入

此外，我们将在 L2 刷新时执行 ssd io

值

enumerator FWD_ROCKSDB_READ

enumerator FWD_L1_EVICTION

enumerator BWD_L1_CNFLCT_MISS_WRITE_BACK

enumerator FLUSH

inline size_t hash_shard(int64_t id, size_t num_shards)

用于 SSD L2 缓存和 rocksdb 分片算法的哈希函数

参数:

• id – 分片键

• num_shards – 分片范围

返回值:

分片 id 范围为 [0, num_shards)

void cuda_callback_func(cudaStream_t stream, cudaError_t status, void *functor)

用于 cudaStreamAddCallback 的回调函数

用于 cudaStreamAddCallback 的通用回调函数，即 cudaStreamCallback_t callback。此函数将 functor 转换为 void 函数，调用它，然后删除它（删除发生在另一个线程中）

参数:

• stream – cudaStreamAddCallback 在其上运行的 CUDA 流

• status – CUDA 状态

• functor – 将被调用的函子

返回值:

无

Tensor masked_index_put_cuda(Tensor self, Tensor indices, Tensor values, Tensor count, const bool use_pipeline, const int64_t preferred_sms)

类似于 torch.Tensor.index_put，但忽略 indices < 0 的情况

masked_index_put_cuda 仅支持 2D 输入 values。它使用 indices 中 >= 0 的行索引将 values 中的 count 行放入 self 中。

# Equivalent PyTorch Python code
indices = indices[:count]
filter_ = indices >= 0
indices_ = indices[filter_]
self[indices_] = values[filter_.nonzero().flatten()]

参数:

• self – 2D 输出张量（被索引的张量）

• indices – 1D 索引张量

• values – 2D 输入张量

• count – 包含要处理的 indices 长度的张量

• use_pipeline – 一个标志，指示此内核将与其他内核重叠。如果为 true，则使用一部分 SM 以减少资源竞争

• preferred_sms – 当 use_pipeline=true 时，内核首选使用的 SM 数量。当 use_pipeline=false 时，此值将被忽略。

返回值:

self 张量

Tensor masked_index_select_cuda(Tensor self, Tensor indices, Tensor values, Tensor count, const bool use_pipeline, const int64_t preferred_sms)

类似于 torch.index_select，但忽略 indices < 0 的情况

masked_index_select_cuda 仅支持 2D 输入 values。它将 indices（其中 indices >= 0）中指定的 count 行从 values 放入 self 中

# Equivalent PyTorch Python code
indices = indices[:count]
filter_ = indices >= 0
indices_ = indices[filter_]
self[filter_.nonzero().flatten()] = values[indices_]

参数:

• self – 2D 输出张量

• indices – 1D 索引张量

• values – 2D 输入张量（被索引的张量）

• count – 包含要处理的 indices 长度的张量

• use_pipeline – 一个标志，指示此内核将与其他内核重叠。如果为 true，则使用一部分 SM 以减少资源竞争

• preferred_sms – 当 use_pipeline=true 时，内核首选使用的 SM 数量。当 use_pipeline=false 时，此值将被忽略。

返回值:

self 张量

std::tuple<Tensor, Tensor> ssd_generate_row_addrs_cuda(const Tensor &lxu_cache_locations, const Tensor &assigned_cache_slots, const Tensor &linear_index_inverse_indices, const Tensor &unique_indices_count_cumsum, const Tensor &cache_set_inverse_indices, const Tensor &lxu_cache_weights, const Tensor &inserted_ssd_weights, const Tensor &unique_indices_length, const Tensor &cache_set_sorted_unique_indices)

为 SSD TBE 数据生成内存地址。

从 SSD 检索的数据可以存储在暂存区 (HBM) 或 LXU 缓存（也位于 HBM 中）。lxu_cache_locations 用于指定数据的位置。如果位置为 -1，则关联索引的数据位于暂存区中；否则，它位于缓存中。为了使 TBE 内核能够方便地访问数据，此算子为每个索引生成第一个字节的内存地址。访问数据时，TBE 内核只需要将地址转换为指针。

此外，此算子还生成后向驱逐索引列表，这些索引基本上是数据位于暂存区中的索引。

参数:

• lxu_cache_locations – 包含为完整索引列表存储数据的缓存槽的张量。-1 是一个哨兵值，表示数据不在缓存中。

• assigned_cache_slots – 包含唯一索引列表的缓存槽的张量。-1 表示数据不在缓存中

• linear_index_inverse_indices – 包含线性索引在排序前的原始位置的张量

• unique_indices_count_cumsum – 包含唯一索引计数排他前缀和结果的张量

• cache_set_inverse_indices – 包含缓存集在排序前的原始位置的张量

• lxu_cache_weights – LXU 缓存张量

• inserted_ssd_weights – 暂存区张量

• unique_indices_length – 包含唯一索引数量的张量（GPU 张量）

• cache_set_sorted_unique_indices – 包含排序后的唯一缓存集的关联唯一索引的张量

返回值:

张量元组（SSD 行地址张量和后向驱逐索引张量）

void ssd_update_row_addrs_cuda(const Tensor &ssd_row_addrs_curr, const Tensor &inserted_ssd_weights_curr_next_map, const Tensor &lxu_cache_locations_curr, const Tensor &linear_index_inverse_indices_curr, const Tensor &unique_indices_count_cumsum_curr, const Tensor &cache_set_inverse_indices_curr, const Tensor &lxu_cache_weights, const Tensor &inserted_ssd_weights_next, const Tensor &unique_indices_length_curr)

更新 SSD TBE 数据的内存地址。

当启用流水线预取时，当前迭代暂存区中的数据可以在预取步骤中移动到 L1 或下一个迭代的暂存区。此算子更新已重定位到正确位置的数据的内存地址。

参数:

• ssd_row_addrs_curr – 包含当前迭代的行地址的张量

• inserted_ssd_weights_curr_next_map – 包含当前迭代中每个索引在下一个迭代的暂存区中的位置之间的映射的张量。(-1 = 数据尚未移动)。inserted_ssd_weights_curr_next_map[i] 是位置

• lxu_cache_locations_curr – 包含为完整的当前迭代索引列表存储数据的缓存槽的张量。-1 是一个哨兵值，表示数据不在缓存中。

• linear_index_inverse_indices_curr – 包含当前迭代线性索引在排序前的原始位置的张量

• unique_indices_count_cumsum_curr – 包含当前迭代唯一索引计数排他前缀和结果的张量

• cache_set_inverse_indices_curr – 包含当前迭代缓存集在排序前的原始位置的张量

• lxu_cache_weights – LXU 缓存张量

• inserted_ssd_weights_next – 下一个迭代的暂存区张量

• unique_indices_length_curr – 包含当前迭代唯一索引数量的张量（GPU 张量）

返回值:

无

void compact_indices_cuda(std::vector<Tensor> compact_indices, Tensor compact_count, std::vector<Tensor> indices, Tensor masks, Tensor count)

压缩给定的索引列表。

此算子根据给定的掩码（包含 0 或 1 的张量）压缩给定的索引列表。该算子移除与其对应掩码为 0 的索引。它仅对 count 个元素（而不是完整张量）进行操作。

示例

indices = [[0, 3, -1, 3, -1, -1, 7], [0, 2, 2, 3, -1, 9, 7]]
masks = [1, 1, 0, 1, 0, 0, 1]
count = 5

# x represents an arbitrary value
compact_indices = [[0, 3, 3, x, x, x, x], [0, 2, 3, x, x, x, x]]
compact_count = 3

参数:

• compact_indices – 压缩索引列表（输出索引）。

• compact_count – 包含压缩后元素数量的张量

• indices – 要压缩的输入索引列表

• masks – 包含 0 或 1 以指示是否移除/保留元素的张量。0 = 移除对应的索引。1 = 保留对应的索引。@count count 包含要压缩的元素数量的张量

class CacheLibCache

#include <cachelib_cache.h>

用于 Cachlib 交互的 Cachelib 包装器类。

它用于维护所有缓存相关操作，包括初始化、插入、查找和驱逐。它对于驱逐逻辑是有状态的，调用者必须专门获取和重置驱逐相关状态。Cachelib 相关优化将在此类内部捕获，例如，fetch 和 delayed markUseful 以提高 get 性能

注意

此类仅处理单个 Cachelib 读取/更新。并行性在调用者端完成

class EmbeddingParameterServer : public EmbeddingKVDB

#include <ps_table_batched_embeddings.h>

用于训练参数服务 (TPS) 客户端的 EmbeddingKVDB 的实现。

class CacheContext

#include <kv_db_table_batched_embeddings.h>

它保存 l2cache 查找结果。

num_misses 是 l2 缓存查找中未命中的次数 cached_addr_list 预先分配了查找次数，以获得更好的并行性，无效位置（缓存未命中）将保留为哨兵值

struct QueueItem

#include <kv_db_table_batched_embeddings.h>

用于后台 L2/rocksdb 更新的队列项

indices/weights/count 是对应的 set() 参数

read_handles 是 cachelib 抽象的索引/嵌入对元数据，稍后将在更新 cachelib LRU 队列时使用，因为我们将其与 EmbeddingKVDB::get_cache() 分离

mode 用于监视 rocksdb 写入，请查看 RocksdbWriteMode 以获取详细说明

class EmbeddingKVDB : public std::enable_shared_from_this<EmbeddingKVDB>

#include <kv_db_table_batched_embeddings.h>

用于与不同缓存层和存储层交互的类，公共调用在 cuda 流上执行。

目前，TBE 使用它将键（嵌入索引）值（嵌入）卸载到 DRAM、SSD 或远程存储，以提供更好的可扩展性，而不会耗尽 HBM 资源

子类：EmbeddingParameterServer, EmbeddingRocksDB

class EmbeddingRocksDB : public EmbeddingKVDB

#include <ssd_table_batched_embeddings.h>

用于 RocksDB 的 EmbeddingKVDB 的实现。