Logo image
Back
I/O-Aware PIM Acceleration for Long-Sequence LLM Inference with Hybrid Sparse Attention
Conference proceeding

I/O-Aware PIM Acceleration for Long-Sequence LLM Inference with Hybrid Sparse Attention

Xiaoyang Lu, Lihan Hu, Hongrui Huang, Peng Jiang and Xian-He Sun
Proceedings - IEEE International Parallel and Distributed Processing Symposium, pp.445-457
05/25/2026
DOI: 10.1109/IPDPS65963.2026.00046

View Online

Abstract

Attention mechanisms are crucial for enabling large language models (LLMs) to process long-sequence inference tasks. However, efficient LLM inference remains challenging due to the quadratic computational complexity of attention in the prefill stage and the memory-bound nature of attention in the decoding stage. Sparse attention techniques partially reduce complexity in the prefill stage, yet substantial data movement between processors and memory (external I/O) continues to limit GPU efficiency in practice. To address these challenges, we propose PILOT, a Processing-in-Memory (PIM)-based, I/O-aware software-hardware co-design that targets memory-bound attention computations. At the hardware level, PILOT strategically integrates processing units within memory to efficiently support hybrid sparse attention in the prefill stage and both dense and sparse attention in the decoding stage, reducing external I/O overhead. At the software level, a detailed I/O analysis guides tiling and scheduling strategies that optimize buffer utilization and minimize internal I/O within the PIM architecture. Finally, we establish a GPU-PILOT heterogeneous system, delivering an efficient end-to-end solution for long-sequence LLM inference. Evaluation results demonstrate that PILOT achieves up to 19.96× and 57.11× speedups over the GPU baseline for sparse attention in the prefill stage and attention in the decoding stage, respectively. The full GPU-PILOT heterogeneous system further improves end-to-end long-sequence LLM inference by up to 5.31× compared to a conventional GPU system.
 Banking Computer Architecture data movement Decoding Graphics processing units I/O analysis Memory Modeling Printing Processing-in-memory (PIM) Sequences Sequential analysis sparse attention acceleration Tiles

Details

Metrics

1 Record Views
Logo image