▐ 背景
阿里妈妈展示广告召回大多采用 Tree-based Deep Model(以下简称TDM)模型，它通过对候选广告的聚类，构造了深达十余层的二叉树索引，并使用 beam search 在此索引上进行检索[1]。由于线上服务的 latency 约束及当时的硬件算力限制使我们不能直接对整个候选集（百万甚至千万量级）进行模型打分。
随着近几年硬件（GPU & ASIC）的发展，我们在模型上的计算能力有了很大提升。于是，算法同学从扩大打分量和提升模型复杂度的两个方向对目前的模型进行了升级。第一个方向是进行全库打分，即抛弃索引结构，用相对简单的双塔模型一次性给所有广告打分；第二个方向诞生了索引扁平化项目，保留树结构索引的同时将其扁平化，减少 beam search 的深度并增加宽度，同时引入 attention 计算，增加了打分模型的复杂度。
这些给工程优化提出了很高的要求，也带来很大挑战。首先，模型大小和计算量的暴增，使得打分的广告数成倍增长，模型的复杂度也大大增加，例如索引扁平化模型中，单次打分消耗的浮点计算次数（FLOPs）就要比原 TDM 模型增长约十余倍。其次，模型中引入了一些非典型的非数值计算（比如TopK和beam search），这些计算如果只使用TensorFlow原生算子会带来性能热点，使 latency 过长（数十上百毫秒）无法满足线上服务的要求。
本文主要分享我们是如何通过系统和算法的协同设计，逐一解决这两个模型在上线过程中的工程难点，为广告召回业务迭代打开新的空间。
▐ 路径一： 全库模型
1. 模型结构介绍
由于每一次用户 page view 模型计算需要处理所有的广告，对于在线系统来说打分量过于巨大（百万至千万量级），只能采用相对简单的模型结构，这里我们选用了经典的双塔模型结构：用户特征和广告特征分别经过DNN得到特征向量，用这两个向量的内积表示匹配程度，并由TopK从中选出最大的那部分。由于广告特征都是固定的，所以广告所在那一路的 DNN 可以在离线提前算完，线上模型直接拿到所有广告的向量作为常量模型参数。如下图所示。
此模型最大的性能挑战来自于 TopK。由于规模达到了数百万，如使用 TF 原生 TopK 算子，其实现并不适合我们的数据特点，latency 将高达数十 ms，不满足线上服务的要求。第二个挑战来自于内积，它是全模型计算量和访存量最大的部分，对于它的处理也相当重要。
2. 对于 TopK 的优化
总体思路
首先，我们调研了 TF 的原生 TopK 算子(即tf.math.topk())的 GPU 实现。它是由若干 RadixSort 组成，其时间复杂度也的确与输入规模n成线性，如下表所示。
观察到在我们问题中 k(500or1000) 是远远小于 n(100W-1000W) 的，可以设定一个阈值，并通过筛选的方式过滤掉大部分数据来减少计算量。例如，我们可以选取千分之一的分位数过滤，这样就相当于将问题规模缩减了1000倍。此筛选操作主要是一个 elementwise 比较加很小规模数据的搬运，在GPU上并行会非常快（在1000W的规模上也不会超过1ms），而这个分位数可以通过采样+排序来进行粗略的估计。虽然这一过程会带来一定的 overhead，但相比于其带来的收益，此 overhead 可以说微乎其微。
筛选算法
对于数据采样的数量，既不能太多，也不能太少。太多会给排序带来很大的 overhead；太少则分位数粒度太粗且不准。通常，我们选取的采样数会在满足分位精度条件的基础上尽可能小，同时要求采样数大于等于k，原因如下：
我们选作阈值的分位数是估计出来的，并不是准确值，所以可能存在经过此阈值筛选出的个数实际上小于k，从而导致后续 TopK 出错。当这种 badcase 出现时，我们的做法是不断选取下一个分位数（即经排序的样本的下一个）再进行筛选，直到筛选出的个数大于等于k。为保证此循环能够正确的结束，我们需要保证样本个数s大于等于k。这样，只要当阈值取到第k个之后，就能保证一定有足够的数被筛选出来(即前k个样本)。
同样，为了降低这一 badcase 出现的概率，我们可以适当放宽此阈值。在实践中，我们的做法是将阈值的分位数放宽到了 k/n 的两倍，即选取第 2_ceil(k/n_s)+1 大的样本作为阈值。根据 order statistics，第k个的样本的概率分布满足如下公式：
下图显示了均匀分布上的1000个样本，1倍分位阈值(第二大样本)和2倍分位阈值(第三大样本)筛选出来的个数的分布情况。左图为概率密度函数 PDF，右图为累积分布函数 CDF。横坐标是筛出的比例，即筛选出的个数除以n。可计算出在 k=1000，n=1000W 的情况下，即横坐标为 k/n=0.0001 时，此方法能将badcase发生概率降低30倍。
与原生TopK性能对比
我们通过TF原生op搭建了一个子图，它不依赖任何 custom op 的实现，并且能在 GPU 上高效的运行。为了匹配 TF 原生的 TopK 算子的语义，我们通过 tf.while() 实现了对高维输入的支持。
下图显示了 batchsize=8 的不同尺寸下，原生 TF 实现和我们的优化算法的性能对比。我们的优化算法相比原生算计性能有了成倍的提升，且在越大的尺寸上优势越明显。
相较于现在常规使用的向量检索框架（例如 Faiss），这种基于 TF 的做法相对灵活，在百万千万的数量级上也有不错的性能。在一千万的规模上，我们的 TopK 算法能达到显存带宽理论上限的10%左右，并且随着问题规模的扩大这一指标能继续提升（Faiss中的 WarpSelection 算法能做到16%，但k的大小受到GPU寄存器数量的限制[2]）。当然，对于 GPU 上类似 WarpSelection 的 state-of-art 的 TopK 算法，我们也能很容易的迁移进TF里来。
3. 对于内积的优化
Batching
当 batchsize 大于1的时候，这里的内积表现为矩阵乘。TF (v1.15)原本用的 cublasSgemmEx 接口并不能在 fp16 类型的 GEMM 上启用 TensorCore，修改成 cublasGemmEx 可以解决这一问题。在 m=8 时，获得约 20%~40% 不等的性能提升，实测 latency 见下表。
观察到在 TensorCore 的加持下，当 batch 从1到8增大时，latency 几乎没怎么变，也就是增大 batchsize 可以显著提升内积这部分的计算效率。不过考虑到线上 latency 的约束，batchsize 还是需要控制在合理范围之内，实践中，我们就将其控制在8以下。
我们测试了 batching 带来的收益。对于一百万左右规模的模型，在线上可以容许的 latency 范围内，按 batchsize=8 做 batching 可以将有效服务容量提升约3倍。
关于 INT8 量化
内积部分的访存在整个模型计算中的占比达到60%以上，不做 batching 时，更是高达90%。为进一步减少访存，我们尝试将这个 GEMM 用 INT8 量化。这里有一个问题需要注意，在 cublasLt 接口中，INT8 GEMM 的输入矩阵需要遵循特殊的内存排列格式，A矩阵和C矩阵遵循 CUBLASLT_ORDER_COL32 格式而B矩阵需要遵循 CUBLASLT_ORDER_COL4_4R2_8C 格式，因此需要再计算前后进行 transform 操作。
幸运的是，B矩阵（即广告向量的矩阵），是一个常量，它的 transform 过程可以被 constant fold，从而避免很多访存开销。对与C矩阵，可以将它的 transform 与一个定制的近似 TopK 的算子进行融合，降低 transform 的开销。而A矩阵的 transform 是无法避免的。
需要注意的时，上文中所描述的利用分位数筛选的 TopK 算法对数据的区分度是有要求的。如果数据的区分度不够大，这个筛选算法就不能有效地约减问题规模，会拖慢后面寻找真实 TopK 的过程的效率。而内积部分使用 INT8 量化的话有可能会导致这个问题，需要对量化算法进行调整和校准。
关于 cache 的利用
另一个节省访存的做法就是充分利用 cache。在内积计算的过程中，大部分的访存都集中在对广告向量的读取上，如果我们能让这个 constant 常驻高速缓存的话，就可以大幅减少内存访问。
GPU 的缓存太小，不太适合进行缓存常驻，而目前涌现的众多人工智能ASIC就拥有相对来说大得多的缓存，例如平头哥的含光800芯片，更适合这样的优化。同时这些芯片在矩阵乘等密集计算上的算力相比 GPU 也毫不逊色。但由于这些芯片支持的算子可能不那么全面，某些定制的量化算法和 TopK 之类的计算就需要被 offload 到 CPU 端进行，此过程中会造成 PCIE 传输的 overhead。
