一句话：首个面向 GPU 的 ML 编译器，将海量嵌入列融合为单内核并做符号化形状与子图优化，显著加速推荐模型推理。
嵌入列对深度推荐模型的精度至关重要，但在推理阶段非常耗时。尽管 ML 编译器已广泛用于自动优化模型，但尚无工作用编译器自动加速推荐模型推理中繁重的嵌入列计算。为此，作者提出 RECom——首个在 GPU 上优化推荐模型海量嵌入列的 ML 编译器，并解决三个挑战：其一，为嵌入列内大量算子生成高效 GPU 调度很困难，现有方案会产生大量小内核且缺乏子图间并行，RECom 采用新的代码生成策略，将海量嵌入列融合进单个内核，并把每一列映射到独立的线程块；其二，动态形状下的复杂形状计算阻碍了进一步的图优化，作者设计了基于符号表达式的模块来重建全部形状计算；其三，ML 框架出于鲁棒性考虑会引入冗余计算，作者开发了子图优化模块，基于整个嵌入列上下文进行图级简化。在自研与开源模型上的实验表明，RECom 在端到端推理延迟和吞吐上分别比当前最优基线提升 6.61 倍和 1.91 倍。
Embedding columns are important for deep recommendation models to achieve high accuracy, but they can be very time-consuming during inference. Machine learning (ML) compilers are used broadly in real businesses to optimize ML models automatically. Unfortunately, no existing work uses compilers to automatically accelerate the heavy embedding column computations during recommendation model inferences. To fill this gap, we propose RECom, the first ML compiler that aims at optimizing the massive embedding columns in recommendation models on the GPU. RECom addresses three major challenges. First, generating an efficient schedule on the GPU for the massive operators within embedding columns is difficult. Existing solutions usually lead to numerous small kernels and also lack inter-subgraph parallelism. We adopt a novel codegen strategy that fuses massive embedding columns into a single kernel and maps each column into a separate thread block on the GPU. Second, the complex shape computations under dynamic shape scenarios impede further graph optimizations. We develop a symbolic expression-based module to reconstruct all shape computations. Third, ML frameworks inevitably introduce redundant computations due to robustness considerations. We develop a subgraph optimization module that performs graph-level simplifications based on the entire embedding column context. Experiments on both in-house and open-source models show that RECom can achieve 6.61X and 1.91X over state-of-the-art baselines in terms of end-to-end inference latency and throughput, respectively. RECom's source code is publicly available at https://github.com/AlibabaResearch/recom.
@inproceedings{paperbot10117,
title = {RECom: A Compiler Approach to Accelerating Recommendation Model Inference with Massive Embedding Columns},
author = {Zaifeng Pan and Zhen Zheng and Feng Zhang and Ruofan Wu and Hao Liang and Dalin Wang and Xiafei Qiu and Junjie Bai and Wei Lin and Xiaoyong Du},
booktitle = {ASPLOS},
year = {2023},
doi = {10.1145/3623278.3624761}
}