Unified memory management improves XPU performance and resource efficiency
GMM: Generalized OS Memory Management for Programmable XPUs
Operating Systems
Summary
Managing memory for specialized AI processors, called XPUs, is hard because each manufacturer uses different software. The authors offer a unified system called GMM that makes it easier to handle memory across different XPUs and take advantage of advanced operating system memory features. Their system works well with XPUs from big companies like Nvidia and Huawei and even improves performance in some cases. It also helps real-world AI applications use memory more efficiently and switch tasks faster.
What this means in practice
- •For xpu driver developers: Create lightweight, unified memory management drivers that improve compatibility and leverage OS memory features across different XPU hardware.
- •For ai service operators: Reduce memory use and speed up model inference switching in AI systems by employing GMM-enabled memory management for XPUs.
Authors
Weixi Zhu, Bin Wang, Zhiyuan Fang, Jiahua Wang, Yao Guo, Gingfung Yeung, Adam Barker, Xinwei Hu
Abstract
Since the AI boom, the fast evolving programmable XPU devices utilize high-bandwidth local memory to accelerate domain-specific tasks offloaded from the host CPU, resulting in a proliferation of XPU memory management systems implemented in proprietary drivers. However, existing OS core memory management (MM) continues to lack a high-level software abstraction, excluding the XPUs from leveraging OS advanced memory features. The XPU manufacturers have to reimplement independent XPU MM drivers, exacerbating the fragmented XPU software ecosystem. This paper presents GMM, generalized memory management that provides unified MM subsystems for XPUs. By adding an explicit logical mapping abstraction, GMM 1) provides universal NUMA-aware memory API for XPUs; 2) allows ease of XPU MM delegation; 3) enables sharing of advanced XPU memory features. Lightweight XPU MM drivers are implemented for two representative XPU manufacturers, Nvidia and Huawei. Both GMM-based drivers obtain competitive performance compared with their vendor-supplied drivers. GMM also obtains up to 41% higher fault-driven page migration throughput. Two real-world industrial case studies are exercised, where GMM reduces XPU memory utilization by 10% with up to 22.6x faster inference serving switch latency.