KV-streams：通过直接操作KV缓存加速Agentic RL训练

An interesting new approach to agentic RL:

Researchers found a way to train AI agents nearly 2x faster.

Agentic RL usually produces long rollouts, and when a rollout reaches its context budget, a compaction policy removes older turns and keeps the prompt, recent turns, or a summary.

Most systems use re-prefill compaction.

They create a shorter sequence and run every retained token through the model again. This rebuilds the KV cache after each compaction, so the same tokens may be processed several times.

KV-streams is a new approach to solve this.

Instead of rebuilding the shorter sequence, KV-streams removes discarded turns directly from the live KV cache. Generation continues from the entries that survived.

In practice, you cannot just reuse those entries as is. It requires two fixes.

> First, RoPE has already rotated each cached Key using its original token position. A Key may move from physical slot 1,000 to slot 500 after compaction, but the model must still treat it as position 1,000.

KV-streams keeps physical storage and logical position separate. Retained Keys keep their rotations. New Queries and Keys continue from the original position count.

> Second, vLLM does not store each KV entry independently. It groups them into fixed 16-token blocks and uses a block table to find those blocks in GPU memory.

KV-streams cannot leave empty slots inside that table because the attention kernel may read the wrong memory. It therefore removes complete blocks and joins the remaining block references together.

This introduces another problem.

A conversation turn may end halfway through a 16-token block. In their vLLM setup, that partial block is not committed to the prefix cache and must be recomputed when the next turn arrives.

If compaction happens first, those tokens may have originally attended to context that is no longer present.

For the vLLM experiments, the researchers pad each completion to a multiple of 16 tokens. This commits every generated token to the cache before the next turn or compaction event. The SGLang setup uses single-token blocks and needs no padding.

Separately, the researchers found that token-level eviction produced degenerate text and required substantial supervised fine-tuning.

They suspect that removing message-boundary tokens disrupted the chat format learned during training. The agent experiments therefore remove complete turns, including their boundary tags.

Once inference works this way, training must follow the same process.

A normal trainer would either process the full rollout or rebuild the shorter sequence after compaction. Both options produce different KV entries from those used during generation.

KV-streams records every compaction event during generation. The trainer then reproduces that history: retained tokens keep the states they formed before compaction, while later tokens are blocked from attending to spans that inference had already removed.

In the repository’s validation runs, the difference between training and generation stayed near a mismatch KL of 0.001, indicating that both paths produced almost identical token probabilities.

On SWE-bench Verified, the Markovian Thinker setup with KV-streams reached peak performance in about 32 hours.

The same compaction strategy with re-prefill took roughly 60 hours, while both reached comparable final scores.

The speedup mainly applies to RL training, where one rollout may be compacted several times.

At production inference, batching can hide some prefill cost, and keeping caches for inactive users reduces serving capacity.

There is one more detail.

The retained KV entries were created while the removed turns were still visible. Future tokens cannot attend to those turns directly, but some of their influence may remain in the entries that were not evicted.

In a controlled test, RL recovered an assignment after its original message was removed. Recall reached 100% when enough cache remained.

KV-streams does not decide what the model should forget. Instead, it makes existing compaction policies cheaper to train by preserving the KV state they choose to retain.

Repo: https://github.com/Emilianopp/KV-streams

If you want to dive deeper into the broader KV-cache engineering stack, I wrote an article that covers eviction, paging, quantization, offloading, and architectural compression.

Read it below.