SSD 嵌入算子

CUDA 算子

enum RocksdbWriteMode

RocksDB 写入模式

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

FWD_L1_EVICTION：L1 缓存驱逐将在前向路径上将数据驱逐到 L2 缓存

BWD_L1_CNFLCT_MISS_WRITE_BACK：L1 冲突未命中将在后向路径上将数据插入 L2 以进行嵌入更新

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

此外，我们将在 L2 刷新时执行 SSD I/O

值

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 – 一个标志，指示此内核将与其他内核重叠。如果为真，则使用一部分 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 仅支持二维输入 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 – 二维输出张量

• indices – 1D 索引张量

• values – 二维输入张量（被索引的张量）

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

• use_pipeline – 一个标志，指示此内核将与其他内核重叠。如果为真，则使用一部分 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>

用于 Cachelib 交互的 Cachelib 包装类。

用于维护所有与缓存相关的操作，包括初始化、插入、查找和驱逐。它对于驱逐逻辑是状态化的，调用方必须专门获取和重置与驱逐相关的状态。与 Cachelib 相关的优化将在此类中捕获，例如获取和延迟标记为有用以提高获取性能

注意

此类仅处理单个 Cachelib 读取/更新。并行处理在调用方一侧完成

class EmbeddingParameterServer : public EmbeddingKVDB

#include <ps_table_batched_embeddings.h>

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

class CacheContext

#include <kv_db_table_batched_embeddings.h>

它保存 L2 缓存查找结果。

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 实现。