InfiniteHiP: Ampliando el Contexto del Modelo de Lenguaje Hasta 3 Millones de Tokens en una Única GPU

InfiniteHiP: Extending Language Model Context Up to 3 Million Tokens on a Single GPU

February 13, 2025

Autores: Heejun Lee, Geon Park, Jaduk Suh, Sung Ju Hwang

cs.AI

Resumen

En los modernos modelos de lenguaje de gran tamaño (LLMs), el manejo de longitudes de contexto muy largas presenta desafíos significativos, ya que provoca velocidades de inferencia más lentas y mayores costos de memoria. Además, la mayoría de los LLMs pre-entrenados existentes no logran generalizar más allá de las longitudes de secuencia de entrenamiento originales. Para permitir la utilización eficiente y práctica de contextos largos, presentamos InfiniteHiP, un marco de inferencia de LLM novedoso y práctico que acelera el procesamiento al eliminar dinámicamente tokens de contexto irrelevantes a través de un algoritmo de poda jerárquica modular. Nuestro método también permite la generalización a secuencias más largas mediante la aplicación selectiva de varios métodos de ajuste de RoPE según los patrones de atención interna dentro de los LLMs. Además, trasladamos la caché de clave-valor a la memoria del host durante la inferencia, reduciendo significativamente la presión de memoria de la GPU. Como resultado, InfiniteHiP permite el procesamiento de hasta 3 millones de tokens en una sola GPU L40s de 48 GB, 3 veces más grande, sin ninguna pérdida permanente de información de contexto. Nuestro marco logra una aceleración de 18.95 veces en la decodificación de atención para un contexto de 1 millón de tokens sin necesidad de entrenamiento adicional. Implementamos nuestro método en el marco SGLang y demostramos su eficacia y practicidad a través de evaluaciones extensas.

English

In modern large language models (LLMs), handling very long context lengths presents significant challenges as it causes slower inference speeds and increased memory costs. Additionally, most existing pre-trained LLMs fail to generalize beyond their original training sequence lengths. To enable efficient and practical long-context utilization, we introduce InfiniteHiP, a novel, and practical LLM inference framework that accelerates processing by dynamically eliminating irrelevant context tokens through a modular hierarchical token pruning algorithm. Our method also allows generalization to longer sequences by selectively applying various RoPE adjustment methods according to the internal attention patterns within LLMs. Furthermore, we offload the key-value cache to host memory during inference, significantly reducing GPU memory pressure. As a result, InfiniteHiP enables the processing of up to 3 million tokens on a single L40s 48GB GPU -- 3x larger -- without any permanent loss of context information. Our framework achieves an 18.95x speedup in attention decoding for a 1 million token context without requiring additional training. We implement our method in the SGLang framework and demonstrate its effectiveness and practicality through extensive evaluations.