MInference 1.0: Acelerando el prellenado para LLMs de largo contexto a través de Atención Dinámica Esparsa.

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

Resumen

Los desafíos computacionales de la inferencia de Modelos de Lenguaje Grandes (LLM, por sus siglas en inglés) siguen siendo una barrera significativa para su implementación generalizada, especialmente a medida que las longitudes de las indicaciones continúan aumentando. Debido a la complejidad cuadrática de la computación de atención, un LLM de 8B tarda 30 minutos en procesar una indicación de 1M de tokens (es decir, la etapa de pre-llenado) en una sola GPU A100. Los métodos existentes para acelerar el pre-llenado a menudo no logran mantener una precisión o eficiencia aceptables cuando se aplican a LLMs de contexto largo. Para abordar esta brecha, presentamos MInference (Inferencia de Millontokens), un método de cálculo disperso diseñado para acelerar el procesamiento de secuencias largas en la etapa de pre-llenado. Específicamente, identificamos tres patrones únicos en las matrices de atención de contexto largo: la forma A, Vertical-Slash y Block-Sparse, que pueden aprovecharse para una computación dispersa eficiente en GPUs. Determinamos el patrón óptimo para cada cabeza de atención de forma offline y construimos dinámicamente índices dispersos basados en el patrón asignado durante la inferencia. Con el patrón e índices dispersos, realizamos cálculos de atención dispersa eficientes a través de nuestros núcleos de GPU optimizados para reducir significativamente la latencia en la etapa de pre-llenado de LLMs de contexto largo. Nuestra técnica propuesta puede aplicarse directamente a LLMs existentes sin ninguna modificación en la configuración de pre-entrenamiento o ajuste adicional. Al evaluar en una amplia gama de tareas secundarias, incluyendo InfiniteBench, RULER, PG-19 y Needle In A Haystack, y modelos que incluyen LLaMA-3-1M, GLM4-1M, Yi-200K, Phi-3-128K y Qwen2-128K, demostramos que MInference reduce efectivamente la latencia de inferencia hasta 10 veces para el pre-llenado en una A100, manteniendo la precisión. Nuestro código está disponible en 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.