Deep Recommender Models Inference: Automatic Asymmetric Data Flow Optimization

TL;DR

This paper introduces automatic asymmetric data flow optimization techniques for deep recommender model inference, significantly improving speed and efficiency on AI accelerators by tailoring data flows and embedding table mappings.

Contribution

It presents four strategies for embedding lookups and a framework for automatic asymmetric mapping of tables to multiple cores, enhancing performance and independence from query distribution.

Findings

Speed-up of 1.5x to 6.5x on real workloads

Over 20x speed-up on unbalanced distributions

Improved independence from query distribution

Abstract

Deep Recommender Models (DLRMs) inference is a fundamental AI workload accounting for more than 79% of the total AI workload in Meta's data centers. DLRMs' performance bottleneck is found in the embedding layers, which perform many random memory accesses to retrieve small embedding vectors from tables of various sizes. We propose the design of tailored data flows to speedup embedding look-ups. Namely, we propose four strategies to look up an embedding table effectively on one core, and a framework to automatically map the tables asymmetrically to the multiple cores of a SoC. We assess the effectiveness of our method using the Huawei Ascend AI accelerators, comparing it with the default Ascend compiler, and we perform high-level comparisons with Nvidia A100. Results show a speed-up varying from 1.5x up to 6.5x for real workload distributions, and more than 20x for extremely unbalanced distributions. Furthermore, the method proves to be much more independent of the query distribution than the baseline.