Memorize Step by Step: Efficient Long-Context Prefilling with Incremental Memory and Decremental Chunk (2024)

Abstract: The evolution of Large Language Models (LLMs) has led to significant advancements, with models like Claude and Gemini capable of processing contexts up to 1 million tokens. However, efficiently handling long sequences remains challenging, particularly during the prefilling stage when input lengths exceed GPU memory capacity. Traditional methods often segment sequence into chunks and compress them iteratively with fixed-size memory. However, our empirical analysis shows that the fixed-size memory results in wasted computational and GPU memory resources. Therefore, we introduces Incremental Memory (IM), a method that starts with a small memory size and gradually increases it, optimizing computational efficiency. Additionally, we propose Decremental Chunk based on Incremental Memory (IMDC), which reduces chunk size while increasing memory size, ensuring stable and lower GPU memory usage. Our experiments demonstrate that IMDC is consistently faster (1.45x) and reduces GPU memory consumption by 23.3% compared to fixed-size memory, achieving comparable performance on the LongBench Benchmark.

## introduction the evolution of large language models (llms) has reached new frontiers, with models like claude @xcite and gemini @xcite capable of processing contexts spanning up to a 1 million tokens. however, the efficiency of processing long sequences with llm still faces significant challenges. the inference of llm can be divided into two parts: prefilling and decoding. llm inference for long documents faces significant challenges in both stages. in the prefill stage, the model needs to read long sequences and endure the quadratic complexity of attention calculations with respect * corresponding author the distribution of memory content across chunks, where we count the number of key-value pairs in memory originating from each chunk. for both @xcite and (b), we used kv cache pruner (snapkv @xcite and streamingllm @xcite ) to compress memory and chunk. to the sequence length. during the decoding stage, decoding each token requires accessing the substantial key-value (kv) cache generated in the prefill stage. most efforts to optimize the efficiency of llm for long sequence focus on the decoding stage, particularly on compressing the kv cache @xcite @xcite @xcite . however, when the input length during the prefilling stage exceeds the maximum length supported by gpu memory capacity, even prefilling cannot proceed. existing works @xcite @xcite tackle this problem by dividing the sequence into chunks with the same size and iteratively compress these chunks with a fixed-size buffer as memory. in this work, we made an empirical investigation on the chunked prefilling with compressed kv-cache, which we refer to simply as memory. our anlysis on the memory displayed in figure 1 reveals that: 1) the attention scores of memory starts at a relatively low value and gradually increases throughout the prefill process (figure 1a .), which suggests that early-stage memory has minimal in-fluence on the next-step computation; 2) once the prefill phase concludes, the memory distribution is primarily concentrated at the end of the sequence (figure 1b ), implying that most of the early-stage memory is not retained by the end of the prefill. overall, our findings suggest that the early-stage memory in the prefill phase is less impactful compared to the later-stage memory. it is unnecessary to maintain a large memory size at the early stage of prefilling. this implies that approaches @xcite @xcite ) that maintaining a fixed-size buffer to compress long sequences may result in wasted computational and memory resources. to avoid computational waste during the early stage of prefilling, we propose incremental memory (im), which starts with a small memory size and gradually increases it until the end of the prefilling phase. during this growth phase, the memory size of im remains smaller than the maximum length, resulting in greater efficiency compared to the commonly used fixed-size memory. while analyzing memory distribution across different layers 1 , we observed that higher layers exhibit a more uniform memory distribution compared to lower layers. consequently, we propose an adaptive memory growth strategy to set memory sizes for each layer based on the proportion of memory retained after compression, with layers retaining more memory being allocated larger memory sizes. although im is faster than fixed-size memory, it does not significantly reduce peak gpu memory usage, as the memory size of im is the same as that of fixed-size memory at the end of the prefilling phase. therefore, we propose decremental chunk based on incremental memory (imdc), which starts with a large chunk size that decreases as memory size increases. when the memory size is small, the chunk size is large, and vice versa. the incremental memory and decremental chunk strategies complement each other, maintaining stable gpu memory usage that is lower than fixed-size memory, which is illustrated in figure 2 . our experiments show that imdc is consistently faster (1.45x) than fixed-size memory and consumes less gpu memory (23.3% reduction) during the prefill stage, yielding comparable results on longbench benchmark @xcite . contributions our main contributions include: ## related works the long-context efficiency of llm has been widely studied, which can be classified into two categories: prefilling and decoding. prefilling the prefilling of llm encounters quadratic complexity in attention calculations with respect to sequence length. numerous research efforts have sought to reduce this quadratic complexity through methods such as low-rank approximation @xcite @xcite and sparsification @xcite @xcite . @xcite provided a comprehensive review of these approaches. these methods modify the computation mode of attention, often resulting in a trade-off with model performance. in contrast, flash attention @xcite identified that the efficiency bottleneck lies primarily in input/output (i/o) operations rather than computational processes. by implementing cuda operations, they significantly accelerated attention calculations without altering the fundamental computation of attention. rmt @xcite proposed an iterative compression scheme for long texts, maintaining and dynamically updating a fixed-size memory, which is followed by @xcite @xcite . autocompressors @xcite ) also introduced incremental memory, but different from our method, they increase memory size to enhance the model performance, which results in significant overhead. 2 decoding most efforts to optimize the efficiency of long-context decoding have focused on kv cache compression. research in this area can be categorized into kv cache pruning @xcite @xcite , low-rank approximation @xcite @xcite , quantization @xcite @xcite , and layer sharing @xcite @xcite . key works in kv cache pruning include h2o @xcite and streamingllm @xcite . h2o selects important kvs based on cumulative attention scores, while streamingllm retains only the kvs closest to the end of the sequence. subsequent works @xcite @xcite @xcite proposed several improvements to h2o, all of which determine kv importance based on attention scores. notable approaches for low-rank approximation include multi-query attention @xcite and grouped query attention (ainslie et al., 2 we demonstrate the superiority of our method compared to autocompressors empirically in appendix b.4. 2023), where different queries share the same kvs. layer sharing methods @xcite identify redundancy among the kv caches of different layers, retaining only the kvs of certain layers. quantization compression @xcite reduces kv cache precision from fp16 to int8 through various quantization methods @xcite . our research adopted the iterative compression method from rmt. however, unlike rmt @xcite , which compresses sequences into soft tokens, we used streamingllm and snapkv to compress kv cache, because they do not require training and can maintain a constant memory size during the iteration.