[Eurosys‘22] Fleche: An Efficient GPU Embedding Cache for Personalized Recommendations

本文针对深度学习推荐系统推理中CPU端DRAM访问和GPU计算之间的速度鸿沟，以及现有的GPU端cache策略的低效性，提出了一种整体的cache策略，来进行高效的GPU端embedding缓存。

1 Introduction

对于深度学习推荐系统，关注的点不仅在于算法优化层面的accuracy，同时也要关注系统的性能表现（如延迟和吞吐量）。
新兴的DLRM，由两部分组成。一是庞大的embedding层（几百GB级别大小）；二是感知机（MLP）层，运行在GPU上（几百MB大小）。其中，由于GPU显存大小有限，embedding table通常以哈希表的方式存储在CPU端的DRAM中。之前的工作表明，CPU端DRAM带宽是严重不足的，原因就在于embedding table的不规律和稀疏访问。这使得embedding table成为限制推荐系统模型性能表现的主要瓶颈。
GPU拥有高带宽的HBM，所以考虑在GPU的HBM中做缓存。现有的缓存方法（static per-table
cache structure）为每一个embedding table维护一个固定大小的cache table。本文通过对NVIDIA HugeCTR推荐系统模型分析发现，现有方法有两个关键缺陷：

• cache的低利用率——来源于所设计的cache数据结构的缺陷
• kernel维护的开销太大

基于以上发现的两个缺陷，本文提出了Fleche。

2 Background & Motivation

下图给出了基于深度学习的推荐系统通用的Embedding-MLP计算模式。

其中，embedding层占据了主要的推荐系统推理延迟。原因在于：

• 随机访问embedding vector造成了大量的CPU cache miss
• 同时由多个线程访问多个embedding table会很快耗尽DRAM带宽

为了解决这个问题，之前的工作提出了在GPU端缓存热的embedding，但是经分析认为其存在缺陷。

Analysis of Existing Cache Schemes

上图展示的是现有embedding cache策略的工作方式，还是比较清楚的。注意现有的方法为每个embedding table加载一个kernel，并附带相关的参数，如对该table进行look up需要的ID_list等。
Issue 1: cache under-utilization

上图展示了HugeCTR在Avazu和Criteo-Kaggle两个数据集、不同cache size下的cache命中率情况。文章分析认为，造成cache命中率低下的原因是因为HugeCTR为所有的embedding table指定相同的cache size。但是不同table热点数据的规模不同，并且会随着时间变化。这样，现有方法这种static的cache结构会倾向于缓存局部的热点数据，而不是所有table范围内全局的热点数据，这会导致低的命中率。
Issue 2: overhead of kernel maintenance
这一部分首先提到了kernel execution time和maintenance time，前者是GPU实际执行的时间，后者是此外的所有时间，包括CPU launching, context initialization, CPU synchronization, communication between CPU and GPU等。实际推荐系统的embedding table个数超过60是很常见的事情，但是从下图来看，在n达到60的时候，maintenance time已经超过了两倍的execution time。

由于现有的方法为每一个table分配一个kernel来处理，所以文章分析上述延迟现象，认为内在的原因是过量的kernel。当所有的请求id被分配到每一个table上时，实际上每个kernel处理的id个数是很小的，这样以来，maintenance time无法被execution time所隐藏，并且总的maintenance time直接正比于query kernel的个数，而query kernel的个数又正比于cache table的个数。（此部分理解不深）

3 Design

3.1 Flat Cache

一句话来说，当请求到达GPU端cache的时候，Fleche会先将所有的ID（来自不同的embedding table）进行重新编码，得到一个统一格式的flat key，然后启动一个统一的kernel去请求cache后端。
Cache table structure
由Fig.5c可以看到，FC的物理存储结构由两部分组成，memory pool和index。其中memory pool存储所有table中的所有embedding，index则存储flat key到embedding的对应关系。
Re-encoding IDs to flat keys with size-aware coding
通过flat key的抽象，所有的cache table能分享一个全局的cache，而不需要感知到多embedding table的存在，从而提高cache利用率。但是为了保持不同embedding table的不同语义，我们需要以一种方式将ID重新编码为FC的key。直接的方法是留出固定长度的bit作为table ID来定位不同的table。但是这种方法没有考虑到embedding table之间大小的差异。
Fleche基于这个问题提出了一种size-aware encoding方法。具体来说，不同table其table ID所占用的bit位数，是满足其剩下feature ID位数能表示所包含的所有key条件下的最长bit位数。
Cache replacement & eviction

• Cache replacement: 一种基于概率的过滤策略。考虑到换入-换出操作的开销，每个missing的embedding以一定的概率p被换进cache。这样会使那些出现频率小于1/p的embedding不会进行换入-换出操作（bypass）
• Cache eviction: 当memory pool的利用率超过一定阈值，FC会扫描整个table，选择将那些冷的embedding逐出cache。考虑到read-after-delete这种情况的存在（释放的embedding被其他线程访问），FC采用了epoch-based space reclamation的策略，具体来说，被驱逐的embedding首先只是在逻辑上被收回（标记），实际的收回会推迟一段时间，直到所有的reader都不访问它。

3.2 Self-identified Kernel Fusion

这个技术的目的是为了将对cache进行请求的kernel个数缩减至1，从而降低kernel maintenance的开销。
为了对多个embedding table进行访问请求，现有的方法为每一套参数分配一个独立的kernel，这带来很大的开销。由于不同的kernel运行相同的函数，不同的只是传入的参数，所以为kernel融合提供了机会。具体融合方法分为三个阶段：
1) Initialization phase：CPU初始化一个array，将所有n个kernel的参数依次存放，并且对kernel包含的线程个数计算一个前缀和数组scan

2) Identification phase：在这个阶段，每一个GPU线程都需要获取其在多kernel场景下的定位。对于一个线程号为tid的线程，首先它会对scan进行一个二分查找，找到小于tid的最大的元素（及其下标$\varphi$），这样该线程就对应第tid-scan[$\varphi$]个kernel。由于所有的线程共享一个scan和args array，所以Fleche将其存放在GPU SM的shared memory里面以减少访问延迟。至于branch divergence，如果每个kernel包含的线程个数都是32的倍数，那么对于一个warp内部的32个线程来说，其进行二分查找的路径是完全一样的，所以这个问题不会存在。
3) Execution phase：每个线程获得相应的参数之后，就可以执行相应的任务。

3.3 Optimizing theWorkflow of Cache Query

这一节给出了进一步优化cache请求工作流的两个技术点。
（1）将copying和indexing解耦合。现有方法为了减少加载的kernel个数，会将index和copy操作耦合在kernel内部。这会受到关键路径copy的影响。如下图Naive所示。

通过Fleche解耦合的设计，cache操作的工作流也发生了变化。请求时，indexing kernel负责定位所有命中的embeddings地址，copying kernel从这些地址复制内容到输出矩阵中。这种设计会带来三个好处

• copy操作从关键路径中去除
• copying时间缩短。之前的copy仅限于用一个warp来完成，但是现在的copying kernel可以根据embedding的维度加载更多的线程，从而提升GPU SM的利用率。
• 如果发生cache miss，Fleche可以预先去请求CPU-DRAM，而不需要等待copying kernel的完成。因为一旦index kernel完成了它的工作，我们就已经知道哪些key发生了miss。

（2）Bypassing indexing of CPU-DRAM layer

这个技术就是将一部分embedding的indexing从CPU-DRAM卸载到GPU上做。具体来说，我们在FC中记录这些embedding的位置，并保存其在CPU-DRAM中的指针。这样对于cache miss的embedding在CPU-DRAM上的索引过程（慢），就替换为快速且并行的GPU上的请求过程。
存在的问题在于，牺牲了GPU cache中存放embedding的空间来存放CPU-DRAM的索引。这里文章提出了一种简单方法来进行trade-off。
（后面是实现和实验部分）

---