一句话:提出 Balance 线程调度器,通过核分配最小化、MLP 感知调度与工作窃取,加速 CPU 上的 DLRM 推理。
推荐模型是数据中心最常用的机器学习负载之一,DLRM 推理在 CPU 上比 GPU 更具成本效益,故其 CPU 推理优化具有重要价值。本文指出已有 DLRM 并行化技术的缺陷:CPU chiplet 间负载不均衡、嵌入表核分配次优、内存级并行(MLP)资源利用不足。作者提出线程调度器 Balance:最小化每张嵌入表的核分配以最大化核利用率,基于嵌入表特征进行 MLP 感知任务调度以更好利用内存带宽,并结合工作窃取与表重排序缓解 chiplet 间负载不均。在真实硬件上用生产 DLRM 轨迹评估,96 核下相比此前最优技术最高加速 1.67×,在多种 batch size 下稳定保持 1.22× 提升。
Recommendation models can enhance consumer experiences and are one of the most frequently used machine learning models in data centers. The deep learning recommendation model (DLRM) is one such key workload. While DLRMs are often trained using GPUs, CPUs can be a cost-effective solution for inference. Therefore, optimizing DLRM inference for CPUs is an important research problem with significant business value. In this work, we identify several shortcomings of existing DLRM parallelization techniques, which can include load imbalance across CPU chiplets, suboptimal core allocation for embedding tables, and inefficient utilization of memory- level parallelism (MLP) resources. We propose a novel thread scheduler, called ''Balance,'' that addresses those shortcomings by (1) minimizing core allocation per embedding table to maximize core utilization, (2) using MLP-aware task scheduling based on the characteristics of the embedding tables to better utilize memory bandwidth, and (3) combining work stealing and table reordering mechanisms to reduce load imbalance across CPU chiplets. We evaluate Balance on real hardware with production DLRM traces and demonstrate up to a 1.67× higher speedup over prior state-of-the-art DLRM parallelization techniques with 96 cores. Further, Balance consistently achieves 1.22× higher performance over a range of batch sizes.
@inproceedings{JainCKKLKD25,
title = {Load and MLP-Aware Thread Orchestration for Recommendation Systems Inference on CPUs},
author = {Rishabh Jain and Teyuh Chou and Onur Kayiran and John Kalamatianos and Gabriel H. Loh and Mahmut T. Kandemir and Chita R. Das},
booktitle = {ASPLOS},
year = {2025},
doi = {10.1145/3676641.3716003}
}