InfiniteHiP: Расширение контекста языковой модели до 3 миллионов токенов на одном графическом процессоре

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

February 13, 2025

Авторы: Heejun Lee, Geon Park, Jaduk Suh, Sung Ju Hwang

cs.AI

Аннотация

В современных больших моделях языка (LLM) обработка очень длинных контекстов представляет существенные вызовы, так как это приводит к замедлению скорости вывода и увеличению затрат памяти. Кроме того, большинство существующих предварительно обученных LLM не способны обобщать за пределы их исходных длин последовательностей обучения. Для обеспечения эффективного и практичного использования длинных контекстов мы представляем InfiniteHiP, новую и практичную инферентную платформу LLM, которая ускоряет обработку путем динамического устранения нерелевантных токенов контекста с помощью модульного иерархического алгоритма обрезки токенов. Наш метод также позволяет обобщать к более длинным последовательностям путем выборочного применения различных методов коррекции RoPE в соответствии с внутренними образцами внимания в LLM. Кроме того, мы выгружаем кэш ключ-значение в оперативную память хоста во время вывода, что значительно снижает давление на память GPU. В результате InfiniteHiP позволяет обрабатывать до 3 миллионов токенов на одном GPU L40s 48GB - в 3 раза больше - без постоянной потери контекстной информации. Наша платформа достигает ускорения в 18,95 раз в декодировании внимания для контекста из 1 миллиона токенов без дополнительного обучения. Мы реализуем наш метод в рамках SGLang и демонстрируем его эффективность и практичность через обширные оценки.

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.