L3: DIMM-PIM Integrated Architecture and Coordination for Scalable Long-Context LLM Inference

2025-04-24

AI Summary

To address the dual memory capacity and bandwidth bottlenecks imposed by KV caching in long-context LLM inference, this work proposes a heterogeneous in-memory computing (PIM) architecture co-integrating GPU and DIMM-level PIM. Focusing on multi-head attention computation during decoding, it introduces a hardware-level redesign to resolve the mismatch between PIM data layout and computational units—a first in the literature. The design enables efficient heterogeneous parallelism via three key techniques: communication-computation overlap, hierarchical KV cache offloading across memory tiers, and adaptive cross-device scheduling. Evaluated on real-world LLM inference traces, the architecture achieves up to 6.1× speedup over state-of-the-art HBM-based PIM approaches, while significantly increasing feasible batch sizes. Results demonstrate both scalability and practicality for production-grade long-context inference.

Technology Category

• Machine Learning: Hardware-aware ML
• Data Mining & Knowledge Management: Scalability, Parallel & Distributed Systems
• Cognitive Modeling & Cognitive Systems: (Computational) Cognitive Architectures

Application Category

• Economics, Online Markets and Human Computation: Architectures and workflows that use LLMs for crowd work
• Search and Retrieval-Augmented AI: Personalized, context-aware and across-device search
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

Abstract

Large Language Models (LLMs) increasingly require processing long text sequences, but GPU memory limitations force difficult trade-offs between memory capacity and bandwidth. While HBM-based acceleration offers high bandwidth, its capacity remains constrained. Offloading data to host-side DIMMs improves capacity but introduces costly data swapping overhead. We identify that the critical memory bottleneck lies in the decoding phase of multi-head attention (MHA) exclusively, which demands substantial capacity for storing KV caches and high bandwidth for attention computation. Our key insight reveals this operation uniquely aligns with modern DIMM-based processing-in-memory (PIM) architectures, which offers scalability of both capacity and bandwidth. Based on this observation and insight, we propose L3, a hardware-software co-designed system integrating DIMM-PIM and GPU devices. L3 introduces three innovations: First, hardware redesigns resolve data layout mismatches and computational element mismatches in DIMM-PIM, enhancing LLM inference utilization. Second, communication optimization enables hiding the data transfer overhead with the computation. Third, an adaptive scheduler coordinates GPU-DIMM-PIM operations to maximize parallelism between devices. Evaluations using real-world traces show L3 achieves up to 6.1$ imes$ speedup over state-of-the-art HBM-PIM solutions while significantly improving batch sizes.

Problem

• Overcoming GPU memory limits for long-context LLM inference
• Reducing data swapping overhead in host-side DIMM offloading
• Optimizing multi-head attention decoding phase bottlenecks

Innovation

• DIMM-PIM and GPU hardware-software co-design
• Hardware redesigns for DIMM-PIM compatibility
• Adaptive scheduler maximizes device parallelism