InfiniteHiP : Étendre le contexte du modèle de langue jusqu'à 3 millions de jetons sur un seul GPU

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

February 13, 2025

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

cs.AI

Résumé

Dans les modèles de langage de grande taille (LLM) modernes, la gestion de longueurs de contexte très étendues pose des défis importants car elle entraîne des vitesses d'inférence plus lentes et des coûts de mémoire accrus. De plus, la plupart des LLM pré-entraînés existants échouent à généraliser au-delà de leurs longueurs de séquence d'entraînement d'origine. Pour permettre une utilisation efficace et pratique des contextes longs, nous introduisons InfiniteHiP, un cadre d'inférence LLM novateur et pratique qui accélère le traitement en éliminant dynamiquement les jetons de contexte non pertinents grâce à un algorithme de taille de jeton hiérarchique modulaire. Notre méthode permet également la généralisation à des séquences plus longues en appliquant sélectivement diverses méthodes d'ajustement de RoPE en fonction des schémas d'attention internes des LLM. De plus, nous déchargeons le cache clé-valeur vers la mémoire hôte pendant l'inférence, réduisant ainsi considérablement la pression sur la mémoire GPU. En conséquence, InfiniteHiP permet le traitement jusqu'à 3 millions de jetons sur un seul GPU L40s 48 Go - 3 fois plus grand - sans aucune perte permanente d'informations de contexte. Notre cadre réalise un gain de vitesse de 18,95 fois dans le décodage de l'attention pour un contexte d'un million de jetons sans nécessiter de formation supplémentaire. Nous implémentons notre méthode dans le cadre SGLang et démontrons son efficacité et sa praticité à travers des évaluations approfondies.

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.