Combating the Memory Walls: Optimization Pathways for Long-Context Agentic LLM Inference

By Haoran Wu ¹, Can Xiao ², Jiayi Nie ¹, Xuan Guo ¹, Binglei Lou ², Jeffrey T. H. Wong ², Zhiwen Mo ², Cheng Zhang ², Przemyslaw Forys ², Wayne Luk ², Hongxiang Fan ², Jianyi Cheng ³, Timothy M. Jones ¹, Rika Antonova ¹, Robert Mullins ¹, Aaron Zhao ^2
1 University of Cambridge
² Imperial College London
³ University of Edinburgh

Abstract

LLMs now form the backbone of AI agents for a diverse array of applications, including tool use, command-line agents, and web or computer use agents. These agentic LLM inference tasks are fundamentally different from chatbot-focused inference -- they often have much larger context lengths to capture complex, prolonged inputs, such as entire webpage DOMs or complicated tool call trajectories. This, in turn, generates significant off-chip memory traffic for the underlying hardware at the inference stage and causes the workload to be constrained by two memory walls, namely the bandwidth and capacity memory walls, preventing the on-chip compute units from achieving high utilization.

In this paper, we introduce PLENA, a hardware-software co-designed system that applies three core optimization pathways to tackle these challenges. PLENA includes an efficient hardware implementation of compute and memory units supporting an asymmetric quantization scheme. PLENA also features a novel flattened systolic array architecture that has native support for FlashAttention to tackle these memory walls in the scenario of inference serving for long-context LLMs. Additionally, PLENA is developed with a complete stack, including a custom ISA, a compiler, a cycle-emulated simulator, and an automated design space exploration flow. The simulated results show that PLENA achieves up to 8.5x higher utilization than existing accelerators, and delivers 2.24x higher throughput than the A100 GPU and 3.85x higher throughput than the TPU v6e, under the same multiplier count and memory settings. The full PLENA system will also be open-sourced.

To read the full article, click here