Round Attention: A Novel Round-Level Attention Mechanism to Accelerate LLM Inference

2025-02-21

AI Summary

To address the GPU memory bottleneck caused by progressive KV cache expansion across dialogue turns in large language model (LLM) inference, this paper proposes a turn-level attention mechanism. We first identify and leverage a “watershed layer” in LLM reasoning—revealing a previously unrecognized regularity: cross-layer attention distributions exhibit strong similarity across dialogue turns. Building upon this insight, we introduce the first turn-aware relevance modeling and dynamic cache pruning framework, enabling selective reuse of key historical-turn KV caches and lightweight reweighting. Experiments on multiple long-context benchmarks demonstrate that our method preserves original model accuracy (within ±0.3% absolute change), reduces GPU memory consumption by 55%, and decreases end-to-end inference latency by 32%.

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Natural Language Processing: (Large) Language Models
• Computer Vision: Large Vision Models

Application Category

• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• 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

The increasing context window size in large language models (LLMs) has improved their ability to handle complex, long-text tasks. However, as the conversation rounds continue, it is required to store a large amount of KV cache in GPU memory, which significantly affects the efficiency and even availability of the model serving systems. This paper analyzes dialogue data from real users and discovers that the LLM inference manifests a watershed layer, after which the distribution of round-level attention shows notable similarity. We propose Round Attention, a novel round-level attention mechanism that only recalls and computes the KV cache of the most relevant rounds. The experiments show that our method saves 55% memory usage without compromising model performance.

Problem

• Reduce GPU memory usage
• Optimize LLM inference efficiency
• Enhance model serving systems

Innovation

• Round-Level Attention Mechanism
• KV Cache Optimization
• Memory Usage Reduction