Presentation
Bringing Graphs Closer to Memory: Compiler-Driven Data Transfer Reduction in the PIM Accelerator
DescriptionProcessing-in-Memory (PIM) alleviates the memory bottleneck of modern AI workloads; however, its limited computational capability often necessitates hybrid architectures that integrate PIM with nearby processing units. In such heterogeneous systems, communication and host-side overheads remain dominant performance bottlenecks. Our profiling of SK hynix's AiMX reveals that CPU–AiMX communication and CPU-side data reordering account for 50.07% of end-to-end LLM prefill and 78.36% of decode execution time. Notably, 54.03% of reordering operations occur adjacent to AiMX-executable nodes, indicating significant untapped potential for near-memory execution to eliminate unnecessary data movement.
We propose a graph compiler that jointly optimizes computation and data movement for AiMX-based acceleration. The compiler performs (1) graph simplification to maximize AiMX-kernel coverage, (2) DMA-assisted tensor reordering and layout redefinition to offload host-side preprocessing, and (3) operator fusion to suppress redundant memory transfers. Integrated into the ONNX Runtime, our approach improves 2.00~2.73x speedup in the prefill phase and 18.50~38.29x in the decode phase, demonstrating the substantial benefits of graph-level co-optimization for PIM-enabled LLM inference and achieving performance comparable to hand-tuned AiMX execution.
We propose a graph compiler that jointly optimizes computation and data movement for AiMX-based acceleration. The compiler performs (1) graph simplification to maximize AiMX-kernel coverage, (2) DMA-assisted tensor reordering and layout redefinition to offload host-side preprocessing, and (3) operator fusion to suppress redundant memory transfers. Integrated into the ONNX Runtime, our approach improves 2.00~2.73x speedup in the prefill phase and 18.50~38.29x in the decode phase, demonstrating the substantial benefits of graph-level co-optimization for PIM-enabled LLM inference and achieving performance comparable to hand-tuned AiMX execution.
Event Type
Research Manuscript
TimeWednesday, July 2910:43am - 10:56am PDT
LocationMtg Room 101B