Chunk indexed sparse attention speeds up long context large language model inference
CISA: Retrofitting Dense Attention Models with Chunk-Indexed Sparse Attention
Machine Learning
Summary
Processing very long text in large language models is slow because the model has to repeatedly read a growing memory of past information. The authors present a way called Chunk-Indexed Sparse Attention (CISA) that speeds this up, without retraining the models. CISA smartly picks small groups of these past information chunks to focus on, reducing the work needed. Their tests show CISA can make the attention part of the model up to 13 times faster while keeping results nearly as good as before.
What this means in practice
- •For machine learning engineers: Speed up inference on large language models when processing very long input texts without retraining the model.
- •For cloud ai service operators: Improve throughput and GPU efficiency when serving many concurrent long-prompt requests with large language models.
Authors
Mohammed Dabbah, Izhak Golan, Oren Tropp, Daniel Afrimi, Vladimir Anisimov, Tomer Asida, Nave Assaf, Ran El-Yaniv, Ido Galil, Yonatan Geifman, Omri Puny, Ido Shahaf, Akhiad Bercovich
Abstract
Long-context LLM inference is costly because attention repeatedly reads a growing key-value (KV) history. We introduce Chunk-Indexed Sparse Attention (CISA), which retrofits pretrained dense attention models without gradient updates. CISA selects fixed-size KV chunks using compact summaries of existing keys, avoiding token-level indexing and a separate full-length index cache. It uses one indexing query per KV group and shares selected chunks across queries in that group. During prefill, QShare reuses selections across adjacent queries to reduce indexing cost and improve KV reuse and GPU utilization. We evaluate CISA on NVIDIA Nemotron 3 Nano and Ultra, and Inkling-Small. In one-layer BF16 benchmarks at batch size 32 using Ultra's per-device attention dimensions, core attention is 11.7x faster for prefill queries near the end of a 1M-token prompt and 13.0x faster during decode. Including chunk scoring, selection, and all other measured GPU work, speedups are 8.6x and 3.0x. On Ultra NVFP4 with a 4,096-chunk budget, aggregate scores stay within 1.5 percentage points of dense attention on MRCR, LongBench v2, and GraphWalks. RULER sweeps on Nano and Ultra BF16 up to one million tokens show scores within 2.05 percentage points of dense attention with 8-16x fewer KV reads. Serving million-token prompts with 32 concurrent requests and 4K output tokens per request, CISA achieves 2.29x the end-to-end throughput of dense attention using Nemotron 3 Ultra NVFP4 on four NVIDIA B300 GPUs. Selecting fewer chunks increases throughput to 2.61x; in separate GraphWalks evaluations, the smaller budget widens the aggregate score gap from 1.49 to 3.59 percentage points. These results expose a practical quality-speed tradeoff without sparse training.