MInference 1.0: Acelerando o Pré-preenchimento para LLMs de Longo Contexto via Atenção Esparsa Dinâmica

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

July 2, 2024

Autores: 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

Resumo

Os desafios computacionais da inferência de Modelos de Linguagem Grandes (LLM) continuam a ser uma barreira significativa para sua implantação generalizada, especialmente à medida que os comprimentos das entradas continuam a aumentar. Devido à complexidade quadrática do cálculo de atenção, um LLM de 8B leva 30 minutos para processar uma entrada de 1M de tokens (ou seja, a etapa de pré-preenchimento) em uma única GPU A100. Métodos existentes para acelerar o pré-preenchimento frequentemente falham em manter precisão ou eficiência aceitáveis quando aplicados a LLMs de contexto longo. Para abordar essa lacuna, introduzimos o MInference (Inferência de Milhões de Tokens), um método de cálculo esparsa projetado para acelerar o processamento de sequências longas na etapa de pré-preenchimento. Especificamente, identificamos três padrões únicos em matrizes de atenção de contexto longo - o formato A, Vertical-Slash e Bloco-Esparsos - que podem ser aproveitados para cálculos esparsos eficientes em GPUs. Determinamos o padrão ideal para cada cabeça de atenção offline e construímos dinamicamente índices esparsos com base no padrão atribuído durante a inferência. Com o padrão e os índices esparsos, realizamos cálculos de atenção esparsos eficientes por meio de nossos kernels de GPU otimizados para reduzir significativamente a latência na etapa de pré-preenchimento de LLMs de contexto longo. Nossa técnica proposta pode ser aplicada diretamente a LLMs existentes sem modificações na configuração de pré-treinamento ou ajustes adicionais. Ao avaliar uma ampla gama de tarefas secundárias, incluindo InfiniteBench, RULER, PG-19 e Needle In A Haystack, e modelos como LLaMA-3-1M, GLM4-1M, Yi-200K, Phi-3-128K e Qwen2-128K, demonstramos que o MInference reduz efetivamente a latência de inferência em até 10 vezes para o pré-preenchimento em uma A100, mantendo a precisão. Nosso código está disponível em 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.