NVIDIA researchers built a new transformer variant.

One small change to the layers made:

• decoding 1.7x faster
• long-reasoning accuracy up 6.5 points

In a typical transformer architecture, every attention layer computes Q, K, and V.

NVIDIA’s tweak adds a fourth projection, which predicts what the next layer will need.

To understand why they did this, let’s first see what happens in a Transformer architecture during inference right now.

Sparse attention was an attempt to handle long-context inference. Instead of attending to every cached token, modern designs score the KV cache in blocks, keep the top-k, and attend only to those.

That cuts attention compute and bandwidth, but this still leaves us with two problems.

First, the KV cache still grows with every generated token.

At 100K+ context, it no longer fits in GPU memory and gets offloaded to CPU RAM.

Now every layer must first copy its selected KV blocks from CPU memory back to the GPU. That copy is slow, the GPU sits idle while it waits, and the stall repeats at every layer of every decode step.

Second, the selection step itself is not free.

Standard selectors score every candidate block with every query head in a GQA group (grouped-query attention, where several query heads share one KV head), then softmax each head’s scores and sum them across the group.

During decode, the sparse attention itself is cheap because there is only one query token.

But the expensive part is deciding which blocks to attend to, and that cost keeps growing with context length.

Both problems trace back to the same design in today’s sparse attention methods, i.e., the attention query drives the block selection.

Selection needs the query vector Q, and Q only exists once its layer is already running. By then, it’s too late to fetch anything early.

The query also drags its multi-head layout into selection, so all that scoring computation runs just to make one top-k decision.

A recent paper from NVIDIA and MIT called SparDA breaks this coupling with one architectural change.

Each layer now emits four projections instead of three:

↳ Q, K, V, and a Forecast.

The Forecast from layer L predicts which KV blocks layer L+1 will need.

Layer L+1’s own query performs the sparse attention over those selected blocks.

This one change fixes both problems.

Since the next layer’s block set is known while the current layer is still computing, the runtime fetches those blocks from CPU memory on a separate CUDA stream.

The copy overlaps with the current layer’s compute, so the GPU no longer waits for it.

And since the Forecast is separate from the attention query, it doesn’t need one score per query head.

SparDA uses one Forecast head per GQA group, which removes the per-query-head scoring loop and skips the softmax step entirely.

DeepSeek did something similar in DSA, where a small indexer picks important tokens instead of the query doing it.

SparDA applies the same idea to blocks and adds the prefetch angle that DSA doesn’t touch.

The cost of the change is small.

The Forecast adds just 33.5M parameters on an 8B model (0.41%), and only those projections are trained, using a KL loss that matches the original selector’s block distribution.

On MiniCPM4.1-8B and NOSA-8B, accuracy matches or beats the sparse baseline, with NOSA-8B gaining +6.5 on long reasoning.

Prefill runs up to 1.25x faster and decode up to 1.7x faster than the sparse offload baseline.

There’s one more benefit.

Because prefetch hides the offload cost, most of the KV cache can live in CPU RAM, and the freed GPU memory fits much bigger batches, pushing decode throughput up to 5.3x over the non-offload sparse baseline.

That said, this lookahead will only pay off during decode with CPU offload. During prefill, all keys already live on the GPU, so the gain there comes purely from the cheaper selection.

Here’s the paper: https://arxiv.org/abs/2606.04511

I wrote a first-principles breakdown of how the KV cache works. It walks through why the model stores keys and values at all, why the cache grows with every token, and a comparison of LLM generation speed with and without KV caching.

Read it below.

SparDA: Sparse Decoupled Attention for Efficient Long-Context LLM Inference

Source: https://arxiv.org/abs/2606.04511 View PDF

Abstract:Sparse attention reduces compute and memory bandwidth for long-context LLM inference. However, two key challenges remain: (1) KV cache capacity still grows with sequence length, and offloading to CPU memory introduces a PCIe transfer bottleneck; (2) the sparse selection step itself retains O\(T^2\) complexity and can dominate attention cost at long contexts. We propose SparDA, a decoupled sparse attention architecture that introduces a fourth per-layer projection, the Forecast, alongside Query, Key, and Value. The Forecast predicts the KV blocks needed by the next layer, enabling lookahead selection that overlaps CPU-to-GPU prefetch with current-layer execution. Because Forecast is decoupled from the attention query, our GQA implementation uses one Forecast head per GQA group, reducing selection overhead versus the original multi-head selector. SparDA adds <0.5% parameters and trains only the Forecast projections by matching the original selector’s attention distribution. On two sparse-pretrained 8B models, SparDA matches or slightly improves accuracy and delivers up to 1.25\\times prefill speedup and 1.7\\times decode speedup over the sparse-attention offload baseline. By enabling larger feasible batch sizes on a single GPU, SparDA further reaches up to 5.3\\times higher decode throughput than the non-offload sparse baseline. Our source code is available atthis https URL.

Submission history

From: Yaosheng Fu [view email] **[v1]**Wed, 3 Jun 2026 06:42:05 UTC (328 KB)