LazyLLM: 효율적인 장문 맥락 LLM 추론을 위한 동적 토큰 프루닝

LazyLLM: Dynamic Token Pruning for Efficient Long Context LLM Inference

July 19, 2024

저자: Qichen Fu, Minsik Cho, Thomas Merth, Sachin Mehta, Mohammad Rastegari, Mahyar Najibi

cs.AI

초록

트랜스포머 기반 대형 언어 모델의 추론 과정은 두 가지 순차적 단계로 구성됩니다: 1) 프롬프트의 KV 캐시를 계산하고 첫 번째 토큰을 생성하는 프리필링(prefilling) 단계, 2) 이후 토큰을 생성하는 디코딩(decoding) 단계. 긴 프롬프트의 경우, 프리필링 단계에서 모든 토큰에 대한 KV 캐시를 계산해야 하므로 첫 번째 토큰 생성에 필요한 시간이 크게 증가할 수 있습니다. 결과적으로 프리필링 단계가 생성 과정의 병목 현상이 될 수 있습니다. 첫 번째 토큰 생성에 모든 프롬프트 토큰이 필수적인지에 대한 의문이 남아 있습니다. 이를 해결하기 위해, 우리는 LazyLLM이라는 새로운 방법을 제안합니다. 이 방법은 프리필링 및 디코딩 단계에서 다음 토큰 예측에 중요한 토큰에 대한 KV를 선택적으로 계산합니다. 프롬프트를 한 번에 정적으로 제거하는 기존의 가지치기(pruning) 접근법과 달리, LazyLLM은 언어 모델이 이전 단계에서 제거되었을지라도 다양한 생성 단계에서 컨텍스트의 다른 토큰 하위 집합을 동적으로 선택할 수 있도록 합니다. 다양한 작업에 걸친 표준 데이터셋에서의 광범위한 실험을 통해 LazyLLM이 기존 언어 모델에 원활하게 통합되어 미세 조정 없이도 생성 속도를 크게 가속할 수 있는 일반적인 방법임을 입증했습니다. 예를 들어, 다중 문서 질의응답 작업에서 LazyLLM은 LLama 2 7B 모델의 프리필링 단계를 정확도를 유지하면서 2.34배 가속했습니다.

English

The inference of transformer-based large language models consists of two sequential stages: 1) a prefilling stage to compute the KV cache of prompts and generate the first token, and 2) a decoding stage to generate subsequent tokens. For long prompts, the KV cache must be computed for all tokens during the prefilling stage, which can significantly increase the time needed to generate the first token. Consequently, the prefilling stage may become a bottleneck in the generation process. An open question remains whether all prompt tokens are essential for generating the first token. To answer this, we introduce a novel method, LazyLLM, that selectively computes the KV for tokens important for the next token prediction in both the prefilling and decoding stages. Contrary to static pruning approaches that prune the prompt at once, LazyLLM allows language models to dynamically select different subsets of tokens from the context in different generation steps, even though they might be pruned in previous steps. Extensive experiments on standard datasets across various tasks demonstrate that LazyLLM is a generic method that can be seamlessly integrated with existing language models to significantly accelerate the generation without fine-tuning. For instance, in the multi-document question-answering task, LazyLLM accelerates the prefilling stage of the LLama 2 7B model by 2.34x while maintaining accuracy.