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

TL;DR

InfiniteHiP tackles the bottlenecks of very long-context LLMs by unifying modular hierarchical token pruning, host-memory KV-cache offloading, and RoPE-based out-of-length generalization into a training-free framework. It uses multi-stage chunk-based pruning to produce a sparse attention mask, enabling block sparse attention with minimal recomputation, while offloading rarely used KV pairs to host memory via an LRUscheme. Dynamic RoPE strategies across stages allow accurate extrapolation beyond the pretrained context, and an efficient implementation in SGLang with Triton enables practical deployment on single GPUs up to $3\times 10^6$ tokens. Experimental results on LongBench and $\infty$Bench show major speedups (e.g., up to $7.24\times$ end-to-end decoding on 3M tokens) and reduced VRAM usage compared to FlashAttention2 and prior long-context methods, with robust out-of-length generalization and minimal degradation in long-context understanding.

Abstract

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.

Figures (12)

• Figure 1: Overview of InfiniteHiP.(a) Infinitely growing KV cache: In InfiniteHiP, the context keys and values are stored in a unified memory space, where some of the keys and values are loaded on GPU memory. (b) Configurable modular pruning: Each pruning stage narrows down the candidate key indices based on the current query block. During pruning, if a cache miss is encountered, the missing tokens are dynamically loaded and the GPU cache is updated. (c) Paged block sparse attention: The selected key indices are used to perform efficient paged block sparse attention.
• Figure 2: Design of our Context Pruning Algorithm.
• Figure 2: $\infty$Bench Results. The average score of each category is the mean of dataset performance, and the average score of the whole benchmark is the relative performance compared to the best-performing result. In the 'Window' column, 'fast' and 'flash' indicates refreshing the sparse attention mask less frequently (see \ref{['sec:exp_setting']}). See the caption on \ref{['tab:longbench']} on 'Abs. Rel.'.
• Figure 3: Results with Llama3.1 8B.
• Figure 4: Results with Short Context Models. Star ($\filledstar$)-shaped markers indicate out-of-length generalization results.