MInference 1.0: 동적 희소 주의 메커니즘을 통한 장문맥 LLM의 사전 채우기 가속화

MInference 1.0: Accelerating Pre-filling for Long-Context LLMs via Dynamic Sparse Attention

July 2, 2024

저자: Huiqiang Jiang, Yucheng Li, Chengruidong Zhang, Qianhui Wu, Xufang Luo, Surin Ahn, Zhenhua Han, Amir H. Abdi, Dongsheng Li, Chin-Yew Lin, Yuqing Yang, Lili Qiu

cs.AI

초록

대규모 언어 모델(LLM) 추론의 계산적 과제는 여전히 광범위한 배포에 있어 상당한 장애물로 남아 있으며, 특히 프롬프트 길이가 계속 증가함에 따라 더욱 두드러지고 있습니다. 어텐션 계산의 2차 복잡도로 인해, 단일 A100 GPU에서 8B 규모의 LLM이 1백만 토큰의 프롬프트(즉, 프리필링 단계)를 처리하는 데 30분이 소요됩니다. 기존의 프리필링 가속화 방법들은 장문맥 LLM에 적용할 때 허용 가능한 정확도나 효율성을 유지하지 못하는 경우가 많습니다. 이러한 격차를 해결하기 위해, 우리는 장문맥 처리를 위한 프리필링을 가속화하기 위해 설계된 희소 계산 방법인 MInference(Milliontokens Inference)를 소개합니다. 구체적으로, 우리는 장문맥 어텐션 행렬에서 A-모양, 수직 슬래시, 블록 희소성이라는 세 가지 독특한 패턴을 식별하여 GPU에서 효율적인 희소 계산을 수행할 수 있도록 합니다. 우리는 각 어텐션 헤드에 대해 오프라인에서 최적의 패턴을 결정하고, 추론 중에 할당된 패턴을 기반으로 동적으로 희소 인덱스를 구축합니다. 패턴과 희소 인덱스를 통해, 우리는 최적화된 GPU 커널을 통해 효율적인 희소 어텐션 계산을 수행하여 장문맥 LLM의 프리필링 단계에서의 지연 시간을 크게 줄입니다. 우리가 제안한 기술은 사전 학습 설정을 수정하거나 추가적인 미세 조정 없이 기존 LLM에 직접 적용할 수 있습니다. InfiniteBench, RULER, PG-19, Needle In A Haystack와 같은 다양한 다운스트림 작업과 LLaMA-3-1M, GLM4-1M, Yi-200K, Phi-3-128K, Qwen2-128K와 같은 모델을 평가함으로써, MInference가 A100에서 프리필링에 대한 추론 지연 시간을 최대 10배까지 효과적으로 줄이면서도 정확도를 유지한다는 것을 입증합니다. 우리의 코드는 https://aka.ms/MInference에서 확인할 수 있습니다.

English

The computational challenges of Large Language Model (LLM) inference remain a significant barrier to their widespread deployment, especially as prompt lengths continue to increase. Due to the quadratic complexity of the attention computation, it takes 30 minutes for an 8B LLM to process a prompt of 1M tokens (i.e., the pre-filling stage) on a single A100 GPU. Existing methods for speeding up prefilling often fail to maintain acceptable accuracy or efficiency when applied to long-context LLMs. To address this gap, we introduce MInference (Milliontokens Inference), a sparse calculation method designed to accelerate pre-filling of long-sequence processing. Specifically, we identify three unique patterns in long-context attention matrices-the A-shape, Vertical-Slash, and Block-Sparsethat can be leveraged for efficient sparse computation on GPUs. We determine the optimal pattern for each attention head offline and dynamically build sparse indices based on the assigned pattern during inference. With the pattern and sparse indices, we perform efficient sparse attention calculations via our optimized GPU kernels to significantly reduce the latency in the pre-filling stage of long-context LLMs. Our proposed technique can be directly applied to existing LLMs without any modifications to the pre-training setup or additional fine-tuning. By evaluating on a wide range of downstream tasks, including InfiniteBench, RULER, PG-19, and Needle In A Haystack, and models including LLaMA-3-1M, GLM4-1M, Yi-200K, Phi-3-128K, and Qwen2-128K, we demonstrate that MInference effectively reduces inference latency by up to 10x for pre-filling on an A100, while maintaining accuracy. Our code is available at https://aka.ms/MInference.