GMM: Generalized OS Memory Management for Programmable XPUs

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.