Architecting Long-Context LLM Acceleration with Packing-Prefetch Scheduler and Ultra-Large Capacity On-Chip Memories

2025-08-11

AI Summary

To address HBM bandwidth saturation and computational bottlenecks caused by KV cache management in long-context LLM inference, this work proposes a pack-and-prefetch co-scheduling architecture, integrated with 3D monolithic stacking and BEOL-embedded memory to realize ultra-large-capacity on-chip KV cache—enabling, for the first time, hardware-level spatial and temporal reuse of long-context KV caches. Leveraging dynamic request packing, fine-grained prefetching, and TPU-like system-level co-optimization, the design significantly alleviates memory access pressure. Evaluated on Llama3.1-8B, it achieves an 8.06× decoding speedup and a 1.83× reduction in end-to-end latency. Under multi-request workloads, throughput improves by 1.7–2.4× while HBM bandwidth consumption decreases by 1.5–2.4×.

Technology Category

• Machine Learning: Hardware-aware ML
• Natural Language Processing: (Large) Language Models
• Computer Vision: Large Vision Models

Application Category

• Economics, Online Markets and Human Computation: Architectures and workflows that use LLMs for crowd work
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Search and Retrieval-Augmented AI: Search Tool Learning with LLM: Teaching LLMs to invoke search and make use of retrieved information

Abstract

Long-context Large Language Model (LLM) inference faces increasing compute bottlenecks as attention calculations scale with context length, primarily due to the growing KV-cache transfer overhead that saturates High Bandwidth Memory (HBM). While prefetching techniques mitigate cache misses by fetching KV data in advance, their spatial and temporal benefits present new opportunities to exploit. This work proposes a packing-prefetch scheduling architecture with monolithic 3D (M3D) back-end-of-line (BEOL) compatible embedded memories with ultra-large on-chip capacity to accelerate long-context LLM inference. Our optimizations demonstrate 8.06x decode speedup and 1.83x overall latency reduction on Llama3.1-8B using TPUv6e-like hardware with additional 512MB BEOL memories over the serial execution. Evaluations of multi-request workloads on TPU-like architectures show 1.7x-2.4x throughput improvement and 1.5x-2.4x HBM bandwidth reduction compared to packing-only methods on Llama3.1-8B and Llama3.1-70B models. With the co-design of packing, prefetching, and BEOL memories, our approach alleviates HBM constraints and enables efficient long-context LLM inference.

Problem

• Reducing KV-cache transfer overhead in long-context LLM inference
• Mitigating High Bandwidth Memory saturation during attention calculations
• Accelerating long-context processing despite scaling computational bottlenecks

Innovation

• Packing-prefetch scheduling architecture for LLM acceleration
• Monolithic 3D BEOL compatible embedded ultra-large memories
• Co-design approach reducing HBM bandwidth constraints