How latent sparse attention (LSA) tackles the search bottleneck in long-context AI

This article is part of our coverage of the latest in AI research.

The attention mechanism is one of the biggest bottlenecks in large language models (LLMs) applications, and it is also where some of the most innovation is happening.

The latest attention optimization technique that has caught my eye is Latent Sparse Attention (LSA), developed by researchers at Prefrontal AI and ETH Zurich. LSA dramatically reduces the computational cost and increases the speed of processing long contexts in large language models (LLMs).

Existing sparse attention methods reduce the number of tokens a model examines but still require expensive computations to determine which tokens deserve attention. LSA addresses this bottleneck by organizing attention keys into clusters and groups to make it easier to search through them when processing long contexts. 

Experiments with 10-million-token tasks show that LSA makes attention-layer prefill more than 20x faster than an existing sparse attention method and more than 70x faster than full attention. 

In a world where a lot of AI workloads depend on processing requests with hundreds of thousands of tokens (e.g., large code repositories or document collections), LSA could have a substantial impact making LLM applications and AI agents more efficient.

The compute bottleneck of sparse attention

The attention mechanism is the part that enables LLMs to compute the relation between the tokens in the input sequence and retrieve information from earlier parts of context when generating new tokens.

At a high level, attention has three components: queries, keys, and values (QKV). Q determines what the new token is searching for, K is the label or index that determines which tokens are relevant to each other, and V is the information that each token contains. When generating a new token, the model compares the Q against the K of preceding tokens to measure their relevance to each other, and combines the corresponding Vs to produce the output.

In LLMs with full attention, each token attends to all previous tokens. As a result, the total compute cost of processing an input sequence (the prefill stage) grows quadratically as the context becomes longer. This means that doubling the context requires roughly four times as many attention computations during prefill. 

To avoid recomputing all the attention values, LLM inference engines store the K and V values in a KV cache. This helps speed up the computation of new tokens but also creates memory requirements that quickly become prohibitively large as the context grows.

The main idea behind sparse attention is that every token doesn’t need to attend to all the tokens that come before them. Sparse attention cuts the compute costs of full attention by selecting a smaller subset of keys for each query. Sparse attention techniques come in different flavors. Some use an indexer component to score keys and identify the ones that are most relevant and need to be attended to. The model then performs attention only over the selected keys.

Sparse attention increases the speed of computing attention values and also reduces the memory costs of storing the KV cache.

However, existing indexers still examine all preceding keys or blocks of keys. Some methods compress multiple keys into smaller representations to reduce the cost of these comparisons, but the indexing operation still increases quadratically during prefill.

This creates its own bottleneck in tasks with very long contexts (though not as severe as full attention). In experiments with GLM-5.3-Flash, a model that already uses sparse attention, the researchers found that at 10 million tokens, the sparse attention layer spent approximately 98% of the prefill time on indexing the tokens it needed to attend to.

At this scale, most of the computation goes into finding which tokens the model should attend to rather than processing their information.

How latent sparse attention works

LSA addresses the indexing bottleneck by changing how keys are organized and searched. To prevent the model from scanning the entire context for every new token, LSA groups similar keys into clusters that can be searched more efficiently.

LSA starts by representing keys in a latent space where similar keys can be identified efficiently. LSA first pools consecutive keys into smaller representations and then groups them into clusters according to their similarity.

Each cluster is represented by a centroid value (a common practice if you’ve studied clustering algorithms) that is stored alongside references to the original keys.

When processing a query, instead of scanning the original keys, LSA compares it against the centroids to identify the most promising clusters. It then retrieves a fixed number of original keys from the clusters that the selected centroids represent. These keys, along with a window of recent tokens and a small number of initial sink tokens, are used to compute attention for the new token.

LSA comes with a few tweaks, including two strategies for searching clusters.

The first is “flat search,” where the query is compared against all cluster centroids. As the context and the number of clusters grow, the computational complexity of indexing during prefill remains quadratic (though it is still substantially cheaper than full attention).

The second strategy, “hierarchical search,” works like tree search. It organizes the clusters into multiple levels and instead of checking every cluster, the query starts at the highest level and progressively follows the most promising branches until it reaches the relevant keys.

With a fixed search budget at each level, this reduces the indexing complexity during prefill from quadratic to approximately O(T log T), where T is the context length.

The two strategies provide a tradeoff. Hierarchical search scales better as the number of clusters grows, but it can miss relevant keys when narrowing the search. On the other hand, flat search preserves accuracy but becomes more compute intensive as the context grows longer.

The researchers also propose two ways of applying LSA to language models.

Distilled LSA (D-LSA) applies the sparsification technique to existing pretrained models. It adds small projection matrices that are trained to select the keys that the model’s original attention mechanism would prioritize. The researchers trained these components on 352 million tokens while keeping the base models frozen.

The distillation data included scientific documents, code repositories, and agent interaction histories, with extra emphasis on long sequences. The researchers found that training on more long sequences improved retrieval accuracy at million-token contexts.

They also propose a training-free version of D-LSA that uses principal component analysis (PCA) to construct the indexing space without training additional projections. This provides a cheaper starting point for applying LSA to existing models, although its accuracy depends on the model.

On the other hand, Native LSA (N-LSA) bakes the latest sparse attention into the model when it is being trained from scratch. The model learns to organize and select keys through its normal next-token prediction objective.

LSA in action

To evaluate D-LSA, the researchers tested it on Qwen3.8-27B, an open model with full attention, and GLM-5.3-Flash, which uses a variant of DeepSeek Sparse Attention (DSA).

In performance experiments on Nvidia H200 GPUs, D-LSA with flat search achieved 22.7-fold and 72.4-fold attention-layer prefill speedups at 10 million tokens on GLM and Qwen respectively (the difference is partly because GLM already uses sparse attention). It also accelerated attention-layer decoding by 4.6 times for GLM and 42.5 times for Qwen.

The improvements carried over to complete model inference, although they were smaller because other model operations remained unchanged. On eight H200 GPUs, D-LSA reduced time per output token by 2.1 times for GLM and 2.9 times for Qwen at 10 million tokens.

But at what context length do these optimizations become worth the added complexity?

“The crossover depends on the baseline attention, hardware and serving setup. In our measurements, D-LSA is roughly break-even at 512k tokens against other sparse models such as GLM-5.3-Flash, and starts to really pay off from about 1M,” James C.R. Whittington, co-author of the paper, told TechTalks.

Whittington attributes the smaller gains at shorter contexts to two factors. First, the existing sparse indexer only becomes a major computational bottleneck at around 512k tokens. Second, LSA accelerates attention but doesn’t speed up other operations in the model. This limits the impact on overall inference time, a constraint known as Amdahl’s law.

An important question is whether these improvements affect model accuracy.

On knowledge and reasoning benchmarks, including MMLU-Pro, GPQA-Diamond, AIME25, and MATH500, D-LSA remained within approximately two percentage points of the original models.

The researchers also evaluated the models on RULER, a benchmark that tests their ability to find and manipulate information across long contexts.

At 10 million tokens, Qwen with D-LSA achieved 63.9% accuracy on RULER, compared with 59.9% for the original full-attention model. GLM showed more mixed results, trailing its existing sparse-attention baseline by up to eight percentage points at context lengths between two and five million tokens.

PaperTrail inserts synthetic facts into scientific documents and asks questions that require finding and connecting information across one or more papers. Some questions require tracing relationships across three or more documents.

At 10 million tokens, Qwen with D-LSA scored 29.1% on PaperTrail, compared with 19.8% for full attention. However, GLM with D-LSA trailed its original sparse-attention implementation by approximately six to seven percentage points at long contexts.

There were also differences between LSA’s search strategies. At 10 million tokens, GLM with hierarchical search scored 55.8% on RULER, compared with 66.8% for flat search. This highlights the accuracy tradeoff that comes with more aggressive search pruning.

The researchers separately tested N-LSA by training a hybrid model with approximately 500 million parameters. At 128,000 tokens, N-LSA matched the validation loss of full attention and maintained similar performance on long-context retrieval benchmarks. It also made attention-layer training 1.6 times faster and reduced whole-model training time by 1.24 times during the long-context training phase.

What LSA means for developers

The researchers have released the LSA code under Apache 2.0 license. The code includes implementations of D-LSA and N-LSA, flat and hierarchical search, and evaluation tools. They have also published distilled projection weights and an integration with SGLang for running supported Qwen and GLM models.

For developers who want to apply LSA to their own pretrained or fine-tuned models, Whittington recommends starting with the training-free approach before investing in distillation.

“For D-LSA, we’d recommend starting with the training-free PCA version and test its accuracy on your target benchmark. If the accuracy is acceptable, that’s the low-hanging fruit. If not you can distill,” he said.

The accuracy gap between the two approaches varies. In the researchers’ experiments, training-free LSA performed nearly as well as distilled LSA on Qwen, but scored approximately 10 points lower on RULER at 10 million tokens on GLM.

For teams that need to distill an indexer, the process requires considerably less compute than retraining the base model. Whittington estimates around 50 GPU-hours for a 27-billion-parameter dense model such as Qwen and 200 GPU-hours for a 300-billion-parameter-class mixture-of-experts model such as GLM. Most of that time goes into running the frozen model to extract training targets. Once those are available, training the projection matrices takes only a few GPU-hours.

LSA is still a research implementation. The published serving setup relies on a custom SGLang fork and GPU-specific kernels, and it is not a drop-in optimization for arbitrary LLMs or hosted model APIs.

If the technique becomes easier to integrate into inference engines, it could have important implications for applications that work with large amounts of information.

Coding agents could process larger portions of code repositories and maintain longer working histories. Research and enterprise assistants could examine extensive document collections without relying as heavily on context truncation.

LSA could also complement retrieval-augmented generation (RAG), which is commonly used to retrieve relevant information from large external knowledge bases. While RAG is effective at finding relevant documents, some tasks require connecting facts scattered across many sources. PaperTrail is designed to test this kind of multi-document reasoning.

The researchers haven’t directly compared LSA against RAG pipelines, but Whittington sees the two approaches as complementary.

“For example, RAG could select around 10M relevant tokens from billions in an external database, and LSA could then process them,” he said.

There are remaining challenges, especially around GPU memory. LSA makes it cheaper to find relevant keys, but it doesn’t eliminate the need to store them.

“The KV cache is unchanged apart from a small index, about 2% of the KV cache and smaller than the indexes of existing sparse attention. Today all keys stay in GPU memory,” Whittington said.

At even longer contexts, such as 100 million tokens, the KV cache could exceed the available GPU memory. One possible solution would be to move less frequently accessed keys to CPU memory or storage while keeping frequently accessed clusters on the GPU. LSA’s buckets could help identify which keys to prefetch, although such a system would require additional engineering and could introduce memory-transfer overhead.

Whittington expects sparse attention to become an increasingly important part of future model architectures. As models process larger contexts, memory management may become as important as efficient key selection.

“We expect sparse attention to become standard in frontier models,” he said.