Hardware-Software Co-design for 3D-DRAM-based LLM Serving Accelerator

This paper presents Helios, a hybrid-bonding-based hardware-software co-designed accelerator for 3D-DRAM that overcomes limitations in existing near-memory processing designs by introducing spatially-aware KV cache allocation and distributed tiled attention execution to significantly improve the speed and energy efficiency of serving highly dynamic large language model workloads.

The Gist

Imagine you are running a massive, 24/7 customer service call center for a super-intelligent AI. This AI (a Large Language Model, or LLM) talks to thousands of people at once. Some people ask short questions ("What's the weather?"), while others ask for a 10,000-word story. Some minutes are quiet; other minutes, the phone lines are jammed.

The problem is that the "brain" of this AI is huge, and the "memory" it needs to remember the conversation (the KV Cache) is constantly changing size and shape.

The Old Way: The Rigid Warehouse

Existing computer chips (like GPUs) and even some new "Near-Memory" chips are like rigid, pre-fabricated warehouse shelves.

• The Problem: If a customer asks for a short story, the system still has to reserve a giant shelf for a 10,000-word story, just in case. This wastes huge amounts of space (memory fragmentation).
• The Bottleneck: If 100 people call at once, but only 10 shelves are free, the other 90 people have to wait, even if there's plenty of total space in the building. The system is too clumsy to rearrange the shelves on the fly.
• The Result: The AI gets slow and inefficient, especially when the workload is unpredictable.

The New Solution: Helios (The Dynamic Lego City)

The paper introduces Helios, a new hardware design that acts like a dynamic, self-organizing Lego city built with a special super-glue called Hybrid Bonding.

Here is how Helios fixes the problems, using simple analogies:

1. The Glue: Hybrid Bonding (The Super-Highway)

Traditional chips are like a city where the library (memory) and the office (processor) are in different buildings connected by a slow, winding road.

• Helios uses Hybrid Bonding to glue the library and the office together, floor-by-floor, with millions of tiny, ultra-fast elevators (wires) connecting them directly.
• Analogy: Instead of driving to the library, the workers are now standing inside the library. They can grab books (data) instantly without leaving their desks. This solves the "traffic jam" of data moving between memory and the processor.

2. The Shelves: Fine-Grained Block Management (The Puzzle Pieces)

Old systems treat a conversation like a single, giant block of concrete. If you need to add one word, you have to move the whole block.

• Helios breaks conversations into tiny Puzzle Pieces (Blocks).
• Analogy: Imagine you are building a wall. Old systems say, "You must use a 10-foot brick." If you only need 2 feet, you waste 8 feet. Helios says, "Use 1-foot bricks." You can fit exactly as many bricks as you need, no matter how big or small the request is.
• The Magic: These puzzle pieces are scattered across all the workers (Processing Engines) in the city. If one worker has space, they take a piece. If another is busy, they take a different piece. No space is ever wasted.

3. The Traffic Cop: Spatially-Aware Allocation (The Smart GPS)

Just having puzzle pieces isn't enough; you need to know where to put them so the workers don't have to run across the city to talk to each other.

• Helios has a Smart GPS (Spatially-Aware Allocation).
• Analogy: When a new conversation starts, the GPS doesn't just throw the puzzle pieces anywhere. It looks at the map and says, "Put these pieces on the desks of the workers who are sitting closest to each other."
• Result: The workers can pass notes to each other instantly because they are neighbors. This minimizes the time spent shouting across the room (data transfer overhead).

4. The Assembly Line: Distributed Tiled Attention (The Assembly Line)

When the AI calculates the answer, it has to look at every word the user said so far.

• Old Way: Everyone stops to look at the whole book at once, then stops to calculate. It's chaotic.
• Helios: It uses a Tiled Assembly Line.
• Analogy: Imagine a team of chefs making a huge pizza. Instead of one chef trying to put cheese on the whole pizza at once, they split the pizza into slices. Each chef puts cheese on their slice, passes the slice to the next chef for sauce, and so on. They work in parallel, overlapping their tasks so no one stands still.
• Result: The AI generates words much faster, even for very long conversations.

The Bottom Line: Why Should You Care?

The paper shows that Helios is a game-changer:

• Speed: It is 3.25 times faster than the best current chips.
• Efficiency: It uses 3.36 times less energy to do the same job.
• Responsiveness: It handles the "rush hour" of AI requests much better, meaning you won't have to wait as long for the AI to reply, even when thousands of people are using it at once.

In short: Helios turns a clumsy, rigid warehouse into a nimble, self-organizing city where data flows instantly, space is never wasted, and the AI can talk to you faster than ever before.

Technical Summary

Here is a detailed technical summary of the paper "Hardware-Software Co-design for 3D-DRAM-based LLM Serving Accelerator" (Helios).

1. Problem Statement

Large Language Model (LLM) serving faces significant challenges due to the highly dynamic nature of workloads, characterized by fluctuating request arrival rates and variable sequence lengths (ranging from short prompts to 10K+ tokens). Existing Near-Memory Processing (NMP) accelerators, which integrate Processing Engines (PEs) into memory modules (e.g., HBM) to boost bandwidth, suffer from two critical limitations when handling these dynamic workloads:

• Coarse-Grained KV Cache Management: Existing NMP designs statically assign full-length Key-Value (KV) caches of attention heads to fixed NMP modules (e.g., specific HBM cubes). This leads to:
  • Compute Load Imbalance: If the number of active requests or KV heads does not perfectly match the number of NMP modules, resources remain idle.
  • Memory Fragmentation: Requests with varying lengths cannot be efficiently packed. Modules often reserve space for the maximum context window even if the current request is short, leading to severe capacity waste. Conversely, a single long request might not fit in one module even if total system capacity is sufficient.
• Inflexible Attention Execution Flow: Current NMP architectures execute attention as two standalone General Matrix Multiplications (GEMMs) with static data partitioning. This rigid flow prevents the fine-grained, block-wise KV cache management (used in software systems like vLLM) from being effectively utilized, hindering the ability to handle variable-length requests dynamically.

2. Methodology: The Helios Architecture

The authors propose Helios, a hybrid-bonding-based LLM serving accelerator that bridges the gap between dynamic software management and distributed hardware abstraction.

A. Hardware Architecture (Hybrid Bonding)

• Technology: Helios utilizes Hybrid Bonding (HB) to 3D-stack multiple DRAM dies (up to 4 layers) onto a logic die. This provides high I/O density (110,000/mm²), low latency, and high bandwidth (0.66 pJ/bit).
• Device Structure: The logic die contains a $4 \times 4$ mesh of Processing Engines (PEs). Each PE has direct access to local DRAM banks via HB controllers.
• PE Design: Each PE includes a Matrix Unit (for GEMM), an Online Softmax Unit, a Vector Unit, and a Reduction Unit. The design supports Distributed Tiled Attention, allowing attention computation to be split across PEs based on KV cache blocks rather than fixed heads.

B. Operator Execution & Communication

• Intra-PE Execution: Helios employs Tiled Attention with Online Softmax. It merges multiple KV blocks into "Macro Blocks" to overlap vector operations (softmax/accumulation) with matrix operations (GEMM), hiding latency.
• Inter-PE Communication:
  • Query Replication: Uses 1D All-Reduce and All-Gather to replicate query vectors across all PEs.
  • Key/Value Transfer: Uses a specialized Reduce-Scatter pattern to route KV vectors only to the PEs holding the corresponding request's last block, minimizing unnecessary data movement.
  • Output Accumulation: Uses 2D Reduce-Scatter followed by All-Gather to aggregate partial attention sums.
  • MLA Support: The architecture supports Multi-Head Latent Attention (MLA) by absorbing reconstruction matrices into FC layers, avoiding the reconstruction of full KV tensors.

C. System Design: Spatially-Aware KV Cache Allocation

To address the mismatch between dynamic software needs and distributed hardware, Helios introduces a Spatially-Aware KV Cache Allocation mechanism:

• Blockwise Management: KV caches are divided into fine-grained blocks (e.g., 64 tokens) and distributed across all PEs.
• Allocation Algorithm: A central scheduler allocates blocks based on two criteria:
  1. Load Balancing: Prioritizes PEs with fewer total tokens to balance compute/storage load.
  2. Communication Optimization: Prioritizes PEs closer to the "center" of the mesh or those that minimize the Manhattan distance for future KV vector transfers (Reduce-Scatter paths).
• Dynamic Adaptation: This allows requests of any length to be served without pre-allocating the full context window, eliminating memory fragmentation.

3. Key Contributions

1. Analysis of NMP Limitations: Identified that existing NMP designs fail under dynamic LLM serving due to static KV cache assignment and rigid attention flows.
2. Helios Architecture: Proposed a hybrid-bonding-based accelerator enabling NMP-native distributed tiled attention, allowing KV cache blocks to be dynamically distributed across all PEs.
3. Spatially-Aware Scheduling: Developed a system-level allocation strategy that balances compute load and minimizes inter-PE communication overhead by considering the physical topology of the PE mesh.
4. Full-Stack Co-Design: Integrated hardware execution flows (online softmax, reduction units) with software management (blockwise KV cache) to support variable-length requests efficiently.

4. Experimental Results

The authors evaluated Helios against state-of-the-art baselines (A100 GPUs, In-die NMP, Duplex, and non-tiled HB designs) across diverse models (OPT, LLaMA3, Mixtral, Qwen, DeepSeek) and workloads.

• Performance Speedup: Helios achieves a 3.25× geometric mean speedup in decoding latency compared to the best-performing GPU/NMP baselines.
• Energy Efficiency: It delivers 3.36× better energy efficiency (geometric mean).
• Serving Performance:
  • Time-Between-Tokens (TBT): Helios reduces P50/P99 TBT degradation by up to 72%/76% compared to baselines under high load.
  • Throughput: Helios achieves 1.19× to 7.33× higher decoding throughput than prior NMP solutions.
• Fragmentation & Load Balance:
  • Memory Fragmentation: Helios reduces peak memory fragmentation ratios by 4.67× to 23.10× compared to coarse-grained NMP designs.
  • Load Imbalance: The load imbalance ratio (max-min latency gap) is kept within 7.1%–13.3%, whereas baselines suffer from imbalance ratios up to 100% under low load.
• Overhead: Inter-PE communication accounts for only ~7% of total decoding latency, validating the efficiency of the communication primitives.

5. Significance

Helios represents a paradigm shift in LLM serving hardware. By moving away from static, coarse-grained memory assignments to a fine-grained, spatially-aware, distributed architecture, it effectively solves the "dynamic workload vs. static hardware" mismatch.

• Scalability: It enables NMP accelerators to handle the extreme variability of real-world LLM serving (fluctuating batch sizes and context lengths) without the memory waste inherent in current designs.
• Efficiency: The co-design of hybrid bonding hardware with tiled attention algorithms maximizes both compute utilization and memory bandwidth, offering a path toward more energy-efficient and lower-latency AI inference services.
• Future Impact: The work demonstrates that emerging 3D-DRAM technologies (Hybrid Bonding) can be fully leveraged for LLMs only when accompanied by a complete rethinking of system scheduling and operator execution flows.