3.14M Context on 3x 3090: The Power of an Architecture That Puts KV Cache in Host RAM

In most cases, the reason a long context cannot fit is the KV cache, not the model. There are probably plenty of people who downloaded a GGUF file that barely fit in VRAM and learned the hard way.

The KV cache is a troublesome. Many people pull up unreliable metrics like wikitext-2 @512 and terrible benchmarks detached from real-world use and still have the audacity to say 8-bit is lossless. Even worse, that is not the only annoying thing about the KV cache. Depending on the model, the KV cache size consumed per context varies a lot, and of course even the sensitivity to KV cache quantization is different.

This is a problem for people without enough VRAM, because they have to deal with the KV cache situation whether they like it or not. However, Qwen Sparse Attention (hereafter QSA) has solved this to some extent.

My environment has three RTX 3090s, with a total of 70.68 GiB of VRAM. Loading the EXL3 4.05bpw quant of Qwen3.8-Flash-Next takes 64.35 GiB, including weights, recurrent states, and workspace. That leaves only 6.33 GiB for the cache. Since fp16 KV takes 27.75 KiB per token, the limit is about 234,000 tokens.

I never quantize the cache. Quantizing increases the context you can fit, but you throw away noticeable quality. Offloading weights to the CPU also increases context, but you throw away speed. Having to choose one of these was the reality of previous Qwen architectures, but this qwen4exp architecture is different.

There is a third way. With Qwen Sparse Attention, you can extend the context with almost no loss in quality or speed. Most of the cache is now in host RAM, the limit is 3,145,728 tokens, and decode runs at 50.7 tok/s at a context of 1 million tokens.

Environment

Three RTX 3090s, 128 GB DDR4-3200 RAM, Arch Linux.

I usually use vLLM, but this time I used EXL3. The reason is that its quality at low bitrates is high.

Why

QSA has a mechanism called the indexer that chooses where in the history to read, and it has a fixed budget. In this model’s text_config, indexer_budget: 2048 and indexer_compress_ratio: 4. Once the context exceeds 4 * block_topk + 3 = 2051 tokens, QSA stops reading the whole history and only reads what the indexer selected. In other words, the K/V read to generate one token is exactly 2048 tokens, and never more than that. Whether there are 3,000 tokens or 3 million tokens behind it, it is the same.

That is pretty much the whole story. Calculating the bytes read from RAM in a single decode step:

K_pad(≈2048) × head_dim(256) × 2 B × 2 (K and V) × num_kv_heads(2)  = 4 MiB / layer
× 12 full-attention layers                                          = 48 MiB / step

It is 48 MiB per token, and it does not increase after that. The measured transfer speed from RAM to GPU on this machine over PCIe 4.0 x16 is 26.50 GB/s, so the transfer takes 1.9 ms/token. Decode was originally in the 15 to 20 ms range. There is room to fit it in.

This does not work for most other models. Since reading scales linearly with context length, even thinking about offloading the cache is a waste of time. Sparse attention was designed to reduce computation, but along the way, it allows the cache to be in RAM.

No Need to Rewrite the Kernel

The second piece of good fortune lies in how the gather kernel determines where to read. From _qsa_sparse_split_kernel:

k_ptrs = k_cache + ((tok[None, :] * n_kv_heads + kv_head) * head_dim + offs_d[:, None])
v_ptrs = v_cache + ((tok[:, None] * n_kv_heads + kv_head) * head_dim + offs_d[None, :])

It simply adds an offset to the base address. Because of this, even if k_cache points to host RAM mapped for direct GPU access rather than VRAM, the kernel runs without modifications. In fact, it just worked. That is because pinned_cuda_view() already existed for MoE CPU offloading.

The implementation only takes a single cache class. At alloc() time:

1. Allocate memory for K/V with anonymous mmap (page-aligned, and pre-zeroed so no 24 GiB memset is needed)
2. Set the target device as current and call cudaHostRegister(ptr, nbytes, PORTABLE | MAPPED)
3. Wrap the returned device pointer as a CUDA tensor

Downstream functions like sparse_attend, get_kv, and ext.paged_kv_cache_update all pass through without noticing. CUDA graph capture also survived. Looking with EXL3_BC_ATTN_TRACE=1, all 12 layers are built, and no layer failed graph capture or fell back to normal execution.

Zero-copy (the method where the GPU reads RAM directly) was chosen here because prefetching is impossible during decode. Which tokens layer i+1 reads is not known until the output of layer i is produced, so transfers cannot be started in advance across layers. On the other hand, if an explicit transfer approach were used, synchronization to read back information about where to read to the CPU would happen for every layer, 12 times per token. In addition, CUDA graph capture would break. With zero-copy, the transfer wait time is hidden behind other operations that the kernel runs in parallel.

What Needs to Stay in VRAM

There is only one place in the KV cache where the amount read is not constant. If this were moved to the host as well, everything would be over.

The indexer scores all blocks at every step. Only the pooled keys used for that scoring (keys pooled per block) have a read volume that grows proportionally with context length. At 1 million tokens, that is 250,000 blocks × 128 × 2 B = 64 MiB/layer, scanning 732 MiB across 12 layers on every single token. At 835 GB/s in VRAM, that takes 0.92 ms. At 26 GB/s over PCIe, that would take 28 ms, which is slower than the entire decode step.

So, K, V, and indexer raw keys go to host RAM, and pooled keys stay in VRAM no matter what. The final split is 0.75 KiB/token in VRAM and 27.00 KiB/token in RAM. That is how the original 27.75 KiB/token entirely in VRAM was divided.

The Downside to Keep in Mind

“Constant read” is only true for a single token. During prefill, thousands of tokens are processed at the same time, and each token independently picks its own 2048 tokens. When you add them all up, you end up reading almost the entire history. If implemented naively, a single 4096-token chunk would read 204 GiB across PCIe. Prefilling 1 million tokens would take 30 minutes just for the transfer.

The solution is to swap the order of the loops. Instead of fetching what is needed token by token, divide the history into fixed-size chunks and copy each chunk into a small working buffer in VRAM just once. Compute all tokens for that chunk together, merge intermediate softmax results, and move on to the next chunk. The outer loop copies from RAM to VRAM, and the inner loop handles tokens. Doing the opposite would re-transfer the same chunk over and over for every batch of tokens.

You can see how well this worked from the statistics output when running with EXL3_QSA_KVO_STATS=1. For a 60,000-token prompt with -cs 262144:

That is a 103-fold reduction. You can also confirm that the number of selections per token does not change with context length, as expected.

There was one unexpected issue here. The original plan was to make the working buffer 2 GiB by default. However, because the model occupies 91% of VRAM on this machine, loading failed even at 128 MiB, and 64 MiB ran out of memory during prefill. The adopted value is 32 MiB. This turned out to be almost harmless. Having a smaller buffer only increases the number of times the history is partitioned, while the number of bytes crossing PCIe does not change. The only thing that increases is how many times the kernel re-reads the list of which tokens to read.

Results

Decode (tok/s):

The baseline column cuts off halfway because it cannot go any further. VRAM runs out around 234,000 tokens. The offload column represents the entire achievement of this work. Going from 2,048 to 1,048,320 tokens, the context grows by 512 times while decode speed degradation is kept to 6.2%. The remaining drop comes from reading all pooled keys every time, which grows proportionally with context length. This is right in line with the initial estimate.

The fixed cost is a drop of about 13% compared to the baseline at short contexts. This is for the PCIe transfer, paid once at the start and not increasing further.

Memory (GiB):

The 3.14 million token figure was reached by stopping the n-gram table from being loaded into RAM and reading it from disk instead, freeing up RAM. From here on, the bottleneck is RAM, not VRAM. The GPU side still has 3.4 GiB left over, which is enough for 4.5 million tokens at 0.75 KiB/token.

For a sense of scale, prefill runs at 1,640 tok/s at 32k and 592 tok/s at 1M. That means building a 1-million-token context from scratch takes about 29.5 minutes. Once built, it can decode at 50 tok/s.

Notes

• Beyond 262,144 tokens, YaRN is needed (rope_scaling in config.json). The 1M / 2M / 3M numbers here are measurements of memory and speed, and say nothing about output quality at those lengths.
• -kvo cannot be used together with the CPU page cache (-ccs). The original plan assumed that if K/V is already on the host, evicting only what remains in VRAM would be enough, but this does not work. Restoring from saved data that does not include K/V leaves behind the K/V from whatever sequence was used previously. To make them coexist properly, everything including K/V must be evicted together and copied from RAM to RAM. That was tedious, so I did not do it.

Usage

Get qsa-kv-offload.patch, then run:

cd exllamav3
git apply qsa-kv-offload.patch

This modifies 8 files and adds 2 files.

To use it, just add -kvo.