SAI: Virtualizing Shared Memory of GPU for AI workload acceleration
TL;DR
This work proposes SAI, a mechanism that virtualizes shared memory into the L2 cache to improve GPU performance for AI applications and introduces an L2 cache management strategy that integrates associativity-based virtual page allocation and a replacement information table, reducing page-swapping overhead while preserving L2 cache performance.
Abstract
AI workloads increasingly demand high parallelism and efficient on-chip memory utilization on modern GPUs. However, the limited capacity of shared memory often constrains thread-level parallelism, while portions of the L2 cache remain underutilized. In this work, we propose SAI, a mechanism that virtualizes shared memory into the L2 cache to improve GPU performance for AI applications. SAI dynamically activates virtualization based on runtime resource usage, enabling additional CTAs with minimal architectural modifications. We further introduce an L2 cache management strategy that integrates associativity-based virtual page allocation and a replacement information table, reducing page-swapping overhead while preserving L2 cache performance. Experimental results demonstrate that SAI achieves an 18.6% performance improvement over the baseline design and outperforms SMILE-opt, our reproduced state-of-the-art baseline, by 8.8%. Moreover, SAI captures 60.2% of the performance gain delivered by the idealized Double SMEM design while reducing energy consumption to 87.7% of the baseline. These results highlight the effectiveness of SAI in enhancing thread-level parallelism and optimizing on-chip memory for AI workloads.