InfiniteHiP: Ampliando o Contexto do Modelo de Linguagem até 3 Milhões de Tokens em uma Ú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

Resumo

Nos modelos de linguagem grandes (LLMs) modernos, lidar com comprimentos de contexto muito longos apresenta desafios significativos, pois causa velocidades de inferência mais lentas e custos de memória aumentados. Além disso, a maioria dos LLMs pré-treinados existentes falha em generalizar além dos comprimentos de sequência de treinamento originais. Para permitir a utilização eficiente e prática de contextos longos, apresentamos o InfiniteHiP, um novo e prático framework de inferência LLM que acelera o processamento ao eliminar dinamicamente tokens de contexto irrelevantes por meio de um algoritmo de poda de token hierárquico modular. Nosso método também permite a generalização para sequências mais longas, aplicando seletivamente vários métodos de ajuste de RoPE de acordo com os padrões de atenção internos nos LLMs. Além disso, transferimos o cache de chave-valor para a memória do host durante a inferência, reduzindo significativamente a pressão de memória da GPU. Como resultado, o InfiniteHiP permite o processamento de até 3 milhões de tokens em uma única GPU L40s 48GB - 3 vezes maior - sem qualquer perda permanente de informações de contexto. Nosso framework alcança uma aceleração de 18,95 vezes na decodificação de atenção para um contexto de 1 milhão de tokens sem exigir treinamento adicional. Implementamos nosso método no framework SGLang e demonstramos sua eficácia e praticidade por meio de extensas avaliações.

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.