Auto-regressive LLM System Reducing Latency via Key-Value Caching
Find Innovative SolutionsGenerate Solutions
Solution Overview
Problem
The existing Large Language Models (LLMs) face significant latency due to the increased input and output dimensions caused by the full read, process, and write of cache data during each iteration, leading to software overhead from repeated memory writes.
Innovation Solution
An auto-regressive system for LLMs is introduced, which reduces the input and output dimensions by caching and reusing key and value data, and utilizing a ring buffer mechanism to optimize memory usage and reduce redundant computations.
Engineering Contradictions & Design Principles
Engineering Contradiction Analysis
1Reliability
If cache data is fully read, processed, and written during each iteration period, then the LLM can maintain accurate token predictions, but the input and output dimensions increase significantly causing latency
Solution Approach 1:
The system performs preliminary actions by caching key data and value data before the actual attention computation. The key data and value data are computed and stored in advance in the key cache and value cache respectively, so that during the forward propagation phase, the model only needs to retrieve these pre-computed values rather than computing them from scratch, thereby reducing inference latency while maintaining prediction accuracy.
Solution Approach 2:
The patent segments the attention mechanism computation into distinct phases: key data computation, value data computation, and attention score calculation. By separating these operations and caching the key and value data independently, the system avoids redundant computations during each iteration while preserving the full attention mechanism's accuracy for token prediction.
2Reliability
If cache data is fully written during each iteration, then the LLM maintains up-to-date information, but repeated rewriting in the same memory address space introduces software overhead
Solution Approach 1:
The system performs preliminary computation of key data and value data and stores them in caches before they are needed. This preliminary action ensures that the cache contains up-to-date information without requiring repeated writing operations during each iteration, as the pre-computed data is simply retrieved and used in subsequent attention calculations.
Solution Approach 2:
The patent implements a copying mechanism where key data and value data are computed once and copied to cache memory structures. Rather than repeatedly writing to the same memory addresses, the system creates copies of the computed data in the key cache and value cache, reducing memory write operations and associated software overhead while maintaining data freshness.
Data Source
Figure 1
Figure 2
Figure 3
AI summary
An auto-regressive method for a large language model includes receiving a hidden state associated with at least one token (S801), generating key data, first value data, and query data according to a received hidden state (S802), generating first positionally encoded key data by encoding the key data positionally (S803), generating positionally encoded query data by encoding the query data positionally (S804), performing first element-wise dot product operations according to the first positionally encoded key data, the positionally encoded query data, and second positionally encoded key data to generate an attention score (S805), performing second element-wise dot product operations according to the first value data, the attention score, and second value data to generate an attention output(S806), and adding the attention output and the hidden state to generate an updated hidden state (S807).