H2EAL: Hybrid-Bonding Architecture with Hybrid Sparse Attention for Efficient Long-Context LLM Inference

TL;DR

<3-5 sentence high-level summary> H2EAL addresses the edge inference bottleneck in long-context LLMs by marrying hybrid-bonding memory with a hybrid static-dynamic sparse attention algorithm. The approach partitions attention heads into streaming (static) and retrieval (dynamic) categories and implements a hardware-software co-design featuring memory-compute co-placement and adaptive tiling to balance workload on HB’s distributed banks. Experimental results across multiple edge-models demonstrate substantial speedups (up to ~48x) and energy-efficiency gains (up to ~73x) with negligible accuracy loss (~0.9%), validating the practicality of HB-based edge LLM inference with long contexts. The work advances edge AI by tackling KV-cache pressure and memory-distribution challenges through a cohesive algorithmic and architectural framework.

Abstract

Large language models (LLMs) have demonstrated remarkable proficiency in a wide range of natural language processing applications. However, the high energy and latency overhead induced by the KV cache limits the edge deployment, especially for long contexts. Emerging hybrid bonding (HB) technology has been proposed as a promising alternative to conventional near-memory processing (NMP) architectures, offering improved bandwidth efficiency and lower power consumption while exhibiting characteristics of distributed memory. In this paper, we propose H2EAL, a hybrid bonding-based accelerator with sparse attention algorithm-hardware co-design for efficient LLM inference at the edge. At the algorithm level, we propose a hybrid sparse attention scheme with static and dynamic sparsity for different heads to fully leverage the sparsity with high accuracy. At the hardware level, we co-design the hardware to support hybrid sparse attention and propose memory-compute co-placement to address the distributed memory bottleneck. Since different attention heads exhibit different sparse patterns and the attention structure often mismatches the HB architecture, we further develop a load-balancing scheduler with parallel tiled attention to address workload imbalance and optimize the mapping strategy. Extensive experiments demonstrate H2EAL achieves 5.20~48.21x speedup and 6.22~73.48x energy efficiency improvement over baseline HB implementation, with a negligible average accuracy drop of 0.87% on multiple benchmarks.

Figures (13)

• Figure 1: Trends in long-context processing, including the evolution of maximum context length of LLMs and datasets.
• Figure 2: Overview of the challenges and contributions in H2EAL.
• Figure 3: (a) The hybrid-bonding architecture comprises stacked memory dies on top and a logic die at the bottom, featuring distributed memory. The Basic LLM implementation on HB architectures utilizes head parallelism. (b) Full attention may lead to an overflow in the KV cache. When using sparse attention, it is observed that the sparsity varies across different heads. (c) Implementation methods of MHA, GQA, and MQA.
• Figure 4: Attention sparsity across heads of multiple LLMs. Colder color indicates higher sparsity.
• Figure 5: Hybrid static-dynamic sparse attention algorithm. (a) Head Identification: an optimization-based method is employed to determine the retrieval and streaming head, where the gating parameter $\alpha$ is the only trainable parameter. (b) Streaming Head: static sparse attention retains sink and local tokens. (c) Retrieval head: dynamic sparse attention utilizes page selection based on the current query.