Source-linked AI summary
KV-streams for Efficient Compaction in Agentic Reinforcement Learning
Emiliano Penaloza, Dane Malenfant, Dheeraj Vattikonda, Roger Creus Castanyer, Siddarth Venkatraman, Abhay Puri, Jonathan Light, Matthew James Sargent, Augustine N. Mavor-Parker, Massimo Caccia, Lucas Caccia, Glen Berseth, Esmeralda S. Whitammer, Alessandro Sordoni, Minseon Kim, Marc-Alexandre Côté, Laurent Charlin, Guillaume Lajoie
TL;DR
Long agentic RL traces consume GPU memory, while compaction controls memory growth but can require repeated prefill that hinders training throughput. KV-streams preserves retained KV entries and reproduces eviction with a trainer attention mask. Across evaluated tasks, it speeds training, and streamed KVs can carry information beyond the visible context without preliminary SFT in the described experiments.
Problem
Long-running agent rollouts grow KV caches that can monopolize GPU memory, while compaction bounds memory but incurs repeated prefill compute.
Source: §1 · p. 2Method
KV-streams carries retained KV entries forward and uses a custom trainer attention mask to reproduce eviction without reprocessing tokens.
Source: §1 · p. 2Results
5× speedup on TextWorld and 1.3× on ALFWorld; software-engineering tasks reach peak performance 3× faster than full context at comparable performance.
Source: §1 · p. 2Takeaways & Limitations
KV-streams is a plug-and-play addition that reduces training wall-clock, matches or exceeds re-prefill compaction performance, and can yield pseudo-recurrent KV behavior without preliminary SFT.
Source: §8 · p. 10
Abstract
from arXiv · showhide
Scaling the horizon of agentic LLMs is bottlenecked by the need to fit ever longer context traces in GPU memory. Context compaction has been the most popular mechanism to alleviate this issue, keeping GPU memory constant for a given trace. Unfortunately, most compaction strategies rely on prefilling the LLM context many times over, hindering training throughput. To alleviate this bottleneck and enable efficient trainable compaction, we propose KV-streams, a plug-and-play strategy compatible with any compaction strategy that substantially increases throughput while showing no evidence of hindering performance. KV-streams enable scalable compaction by streaming the KV cache forward rather than flushing it after each compaction. We show that KV-streams enable three different compaction strategies, achieving a 2.6 to 5x wall-clock speedup in training. Beyond efficiency, we find that the streamed KV cache can act as a recurrent state, carrying forward information that has long since disappeared from the context. Specifically, in a controlled setting we show that, contrary to prior work, RL alone is all that is needed for this behavior to emerge. Overall, we show KV-streams to be an efficient and lightweight plug-and-play addition to any post-training pipeline.
B Compaction strategies
Figure 2 contrasts re-prefill compaction with KV-streams in both training traces and context transformations. KV-streams preserves retained KVs and uses an attention mask to reproduce eviction.
B Compaction strategies
- Re-prefill compaction recomputes retained tokens, while KV-streams preserves their existing KVs.Panel A contrasts the two strategies’ training costs.Source: Figure 2 · p. 2
- KV-streams trains on one continuous trace with an attention mask that reproduces eviction.Source: Figure 2 · p. 2
- Panel B shows context transformations for each strategy, labeling prompt as P and summary as S.Source: Figure 2 · p. 2
1 Introduction
Long agentic RL traces consume GPU memory through growing KV caches, while compaction limits memory but incurs repeated prefill work. KV-streams avoids that repeated prefill; evaluations report faster training and findings on performance and recurrent-state behavior.
1 Introduction
- Long-running rollouts’ KV caches grow with their length, and the longest rollouts can monopolize GPU memory and throttle throughput.Reducing long-trace memory can increase training throughput, allowing more optimizer steps and larger batch sizes.Source: §1 · p. 2
- Compaction bounds trace memory but increases training compute because retained context is prefilled again after each compaction.This repeated work requires the model to make forward passes over the same tokens.Source: §1 · p. 2
- KV-streams evicts KV entries in the inference engine and uses a custom trainer attention mask to reproduce eviction without reprocessing tokens.The approach preserves a continuous KV stream and is compatible with the four tested strategies and others.Source: §1 · p. 2
- 5× speedup on TextWorld and 1.3× on ALFWorld; software-engineering tasks reach peak performance 3× faster than full context at comparable performance.With sufficient training, KV-streams can close the performance gap with full context and even surpass it; models also generalize better to unseen text-based games.Source: §1 · p. 2
- Synthetic experiments show pseudo-recurrence can emerge in the KV cache without preliminary SFT, unlike prior work’s reported requirement.The authors report that extra SFT wall-clock time does not improve performance in their test.Source: §1 · p. 2
- KV-streams is presented as a plug-and-play addition that reduces training wall-clock and matches or exceeds re-prefill compaction performance.The conclusion also describes effective generalization and no need for an expensive SFT initialization stage.Source: §1 · pp. 2–3
2 Compaction
Compaction operators shorten context when a trigger is met, but typical re-prefill pipelines recompute retained entries in fresh traces. KV-streams reuses retained KVs and masks deleted spans during training, avoiding that repeated prefill work.
2 Compaction
- The paper describes compaction at the entry level, although its experiments use turn-wise compaction; the process generalizes to tokens or turns.The authors state that KV-streams is amenable to other long-context use cases.Source: §2 · p. 3
- A compaction operator maps context C to shorter context C′ when a token budget or model decision triggers it.The operator deletes at least one token span and may generate a bridge, such as a summary, in its place.Source: §2 · p. 3
- Typical re-prefill compaction resets retained tokens’ positions and KVs in a fresh trace, forcing those tokens to be prefetched again.The recomputed cache differs from the original because the token is processed under a new position and context.Source: §2 · p. 3
- Under a budget B, re-prefill compaction repeatedly processes the retained r entries, with extra work increasing as r approaches B.A rollout of N entries compacts after generating B − r new entries; the repeated prefill cost diverges as retained context approaches the budget.Source: §2 · p. 3
- Figure 3 compares extra prefills over a 32k-token rollout with tokens kept per compaction, across memory budgets B; KV-streams prefills none again.As retained tokens approach the budget, repeated prefill grows by a factor of N B−(o+s).Source: Figure 3 · p. 3
- KV-streams carries retained overlap and summary KVs forward, then uses a custom attention mask to block attention to deleted spans in one trainer pass.This lets a compacted rollout avoid re-prefilling and remain a single trace rather than splitting into separate traces.Source: §2 · p. 4
3 Compaction Strategies
The section studies four compaction strategies spanning low to high context retention. They differ in whether they summarize prior context, retain a fixed or selected subset, or keep nearly the full window.
3 Compaction Strategies
- The four strategies span from retaining almost none of the context to retaining nearly the full budget.Source: §3 · p. 4
- Summary replaces prior context with a variable-length summary and keeps no entries, with o = 0.The summary length s is bounded by a preset limit.Source: §3 · p. 4
- Markovian Thinker deletes half the context at each compaction and replaces nothing, retaining o = B/2 and s = 0.Source: §3 · p. 4
- Markovian Pick lets the model choose which b entries to retain and deletes all others at each compaction.Source: §3 · p. 4
- Sliding Window retains almost the full window, with o close to B and s = 0, probing the high-retention regime where repeated prefill is most expensive.The passage states this is where KV-streams helps most.Source: §3 · p. 4
4 Experimental Setup
The experiments compare compaction strategies across text-based games and software-engineering tasks, tracking evaluation performance against GPU-hours and wall-clock time.
4 Experimental Setup
- Runs report evaluation score against cumulative GPU-hours and stop early only when their wall-clock time exceeds the full-context run.Across benchmarks, runs use 8 rollouts per prompt and a learning rate of 10^-6, with shared optimization settings within each benchmark.Source: §4 · p. 4
- TextWorld uses 256 evaluation tasks and up to 200-turn episodes; ALFWorld evaluates success on 70 seen and 67 unseen games with episodes up to 100 turns.TextWorld training uses 500 gradient steps; ALFWorld trains for up to 200 gradient steps.Source: §4 · p. 4
- Software-engineering runs train on filtered SWE-rebench and ScaleSWE tasks, filter by pass@3, evaluate on SWE-bench Verified, and compare Markovian Thinker and Sliding Window against full context.Training uses 1,028 SWE-rebench and 443 ScaleSWE tasks; episodes run for up to 100 turns within a 64k-token budget, with compaction every 30 turns and batch size 256.Source: §4 · pp. 4–5
- Figure 4 plots success against GPU-hours and wall-clock hours, while Figure 5 reports per-step timing and rollout statistics on TextWorld.Figure 5 separates generation and forward/backward time and also shows mean turns per episode and sequence length.Source: Figure 4 · p. 5, Figure 5 · p. 6
- Only ALFWorld receives multi-seed experiments; TextWorld and software-engineering results use single seeds because of compute limits.Some experiments take over a week to execute.Source: §4 · p. 5
- Text-game experiments use modified vLLM and prime-RL components, while software-engineering experiments use modified SGLang and Slime components.The vLLM setup pads completions to 16-token multiples; SGLang commits individual tokens and needs no padding.Source: §4 · p. 5
5 Experimental Results
Across text-based games, KV-streams reaches re-prefill compaction's final performance with substantially less compute. On SWE-bench Verified, it performs comparably to full context and re-prefill compaction while reaching peak performance sooner.
5 Experimental Results
- The results cover two text-based games and software-engineering tasks, with only a subset of compaction strategies tested for software engineering because of compute limits.The software-engineering evaluation compares two KV-streams configurations against full context.Source: §5 · p. 6
5.1 Text-Based Games
- 3.7×–11.3× fewer GPU-hours on TextWorld and 1.3×–3.0× higher throughput on ALFWorld let KV-streams reach re-prefill's final success without performance loss.On TextWorld, KV-streams performs similarly or better; on ALFWorld it again reaches re-prefill's final success rate across four strategies.Source: §5.1 · p. 6
- Repeated recomputation of retained context makes forward/backward the re-prefill bottleneck, while KV-streams reduces both forward/backward and generation time.Re-prefill's early sampling-time advantage from shorter episodes diminishes as training progresses and is erased by added training time.Source: §5.1 · p. 6
- On every unseen text-based game, a KV-streams variant performs on par with the best strategy and improves on the base model.No single compaction strategy performs best across all games.Source: §5.1 · p. 6
5.2 Software Engineering Agents
- 52.7 ± 2.1% on SWE-bench Verified is statistically within full context's 51.5 ± 1.2% and re-prefill Markovian Thinker's 53.4 ± 0.2%.Sliding-Window with KV-streams reaches peak performance in about 20 hours versus more than 65 hours for full context; Markovian Thinker reaches peak in about 32 hours versus about 60 hours with re-prefill.Source: §5.2 · p. 7, Figure 6 · p. 7
6 Is SFT Necessary for Recurrent KVs?
This section tests whether KV-streams can carry information beyond the visible context without a preliminary SFT stage. It finds that RL alone can produce this behavior, with sufficient KV capacity in the controlled retrieval tasks.
6.1 Prior SFT on TextWorld
- Once SFT compute is counted, RL alone is more compute efficient and shows no observable performance drop versus an SFT warm-start.The comparison covers Sliding-Window and Markovian Thinker with KV-streams; SFT does not justify its cost for Markovian Thinker, while both strategies reach peak performance in the same time with Sliding-Window.Source: §6.1 · p. 8
6.2 Recalling Evicted Content
- The synthetic task evicts the assignment turn before asking the model to retrieve the assigned object, testing whether decoded tokens’ KV entries carry that information forward.The study varies the number of decoded tokens and also tests new object classes and a target with zero initial probability.Source: §6.2 · p. 8
- With sufficient KV-token budget, RL recall reaches 100% on evicted fruit assignments and matches SFT, becoming unreliable only at 16 and 32 tokens.Because the candidate list always contains the assigned fruit, random RL sampling can produce initial successes to bootstrap learning.Source: §6.2 · p. 8
- At larger token budgets, RL and SFT perform similarly when recalling actors after training on fruits; RL trails at smaller budgets.This tests generalization beyond the training context and object domain.Source: §6.2 · p. 9
- As KV-token budget grows, RL learns to recall a television with 100% accuracy even though the recall choices list only fruits.The television has almost no probability under the base model after eviction, making this the hardest tested case; SFT can force the target response by construction.Source: §6.2 · p. 9
- Taken together, the empirical and synthetic experiments find that RL alone can carry information forward through the KV cache given sufficient capacity, without an expensive SFT stage.The authors conclude this makes KV-streams a plug-and-play mechanism for post-training without an additional pipeline stage.Source: §6.2 · p. 9
7 Related Work
Prior approaches either modify model architecture to manage long contexts or compact a fixed context window, often rebuilding it after compaction. KV-streams targets that repeated prefill cost while remaining orthogonal to the choice of compaction rule.
7 Related Work
- Other long-context approaches modify the model or carry and compress segment memories, including recurrent and state-space architectures, Transformer-XL, Compressive Transformer, and GTrXL.Recurrent and state-space architectures hold cost flat only after retraining the model to use them.Source: §7 · pp. 9–10
- Fixed-size summarization and recent-turn methods rebuild the context after compaction, repeatedly prefilling retained tokens; KV-streams instead drops evicted KVs in the inference engine and mirrors eviction with a trainer attention mask.The authors characterize this as producing steady per-token cost independent of the compaction rule, and distinguish it from attention sinks’ fixed retention pattern.Source: §7 · p. 10
- Learned context-management methods choose what survives compaction, but still rebuild and prefill the window; KV-streams is orthogonal and can make those strategies cheaper to train.The related work includes methods that retain generated summaries or compact memory states.Source: §7 · p. 10
8 Conclusion
The paper concludes that KV-streams reduces experiment wall-time while matching or improving performance and allowing streamed KVs to carry information removed from context. Its broader impacts include both lower training costs and the need for safeguards when agents handle sensitive information.
8 Conclusion
- KV-streams is presented as a plug-and-play mechanism for any compaction strategy, reducing experiment wall-time by up to 3× while often improving performance and at worst not hindering it.The conclusion also describes streamed KVs as a pseudo-recurrent state and reports that this behavior does not require prior SFT.Source: §8 · p. 10
AI Use Statement
- The authors report using generative AI for manuscript drafting and editing, data analysis and plotting code, and creating and refining figures, diagrams, and illustrative icons.They state that they reviewed the AI-assisted work and take responsibility for the final content and artifacts.Source: AI Use Statement · p. 10
A Broader Impacts
- Lower training costs may also make harmful or unauthorized automation cheaper, so agents handling sensitive information still need access controls, monitoring, and secure data handling.The authors state that KV-streams does not replace underlying models’ and environments’ safety policies, permissions, or privacy protections.Source: A Broader Impacts · p. 13
B Training and Implementation Details
The appendix documents shared benchmark settings, rollout concurrency, implementation choices, and extended training and timing curves. It explains practical constraints including memory-driven concurrency, block padding, and the reduced context-length precision of turn-wise eviction.
B.1 Hyperparameters and Rollout Settings
- Hyperparameters are shared within each benchmark, while compaction strategy and rollout concurrency vary; re-prefill and KV-streams use the same concurrency settings, unlike full context.Text-based games use prime-RL with vLLM, while software-engineering tasks use Slime with SGLang; full context uses lower concurrency because long traces occupy more memory.Source: §B.1 · p. 14, Table 1 · p. 14, Table 2 · p. 14
B.2 Concurrency Sweep
- Full-context throughput falls above 12 concurrent rollouts per GPU, while sliding-window KV-streams sustains higher concurrency.The sweep uses vLLM on four A100 (80GB) GPUs; the figure reports rollouts per second and generated tokens per second, with hollow markers for KV-cache-full preemptions.Source: §B.2 · p. 15, Figure 9 · p. 15
B.3 vLLM Block Padding
- vLLM reuses only complete 16-token blocks, so turns are padded to multiples of 16 for KV reuse; SGLang uses block size 1 and needs no padding.Figure 4 token counts include the vLLM padding.Source: §B.3 · p. 15
B.4 Turn-Wise Compaction
- Turn-wise compaction preserves message boundaries, but gives less precise control over context length and per-trace memory use than token-level eviction.Token-level eviction produced degenerate text initially, while turn-wise eviction caused no such initial-response instability; the authors suspect removing boundary tokens disrupts learned format.Source: §B.4 · p. 15
C.1 Full Training Curves
- Summary and Sliding-Window re-prefill runs exceed the full-context wall-clock budget and have longer step times in the extended curves.The figures extend training and timing curves through each run’s last checkpoint, including runs stopped after exceeding that budget; Figure 12 plots reward and evaluation success by gradient step.Source: §C.1 · p. 16, Figure 10 · p. 16, Figure 11 · p. 17
C.2 SFT Warm-Start Training Loss
The SFT warm-start loss drops quickly before declining more slowly. The retrieval evaluations test whether information remains available in the KV cache after the assignment turn is evicted, including under changed prompts and targets.
C.2 SFT Warm-Start Training Loss
- Loss falls quickly over the first 50 steps, then more slowly during the 500-step SFT warm-start on 10k full-context trajectories.Figure 13 plots per-step loss and its centered 9-step rolling mean.Source: §C.2 · p. 18, Figure 13 · p. 18
D.1 Task Setup
- After assigning one of five fruits, models retain ascending-integer sequence KVs, evict the assignment and acknowledgment, then recall by exact match; correct recall therefore depends on the KVs.The retained sequence is checked to contain no candidate answer.Source: §D.1 · p. 18
D.2 Evaluation Conditions
- Prompt generalization changes the retrieval prompt without further training while keeping assignments, candidates, KV capacity, and eviction unchanged; the options are shuffled.The alternate prompt asks for a favourite object rather than asking which fruit was assigned.Source: §D.2 · p. 18
- The zero-support target evaluation trains on six objects, including television, although prompts list only the five fruits; it uses 240 trials balanced across objects.Correct recall must recover television even though it is not among the listed choices.Source: §D.2 · p. 19
- Actor-name generalization tests five names without further training, both with and without the names listed as candidates, using 200 trials per condition.Sequences containing a candidate answer are rejected, and recall is scored by exact match.Source: §D.2 · p. 19
D.3 Training Objectives
The appendix specifies training objectives and additional retrieval evaluations, then derives the repeated-prefill cost of compaction. KV-streams reuses retained KVs, eliminating that prefill cost in the stated analysis.
D.3 Training Objectives
- One model per objective and KV capacity is trained for 100 updates; self-distillation starts from the corresponding SFT checkpoint and runs for another 100 updates.Final checkpoints are evaluated; the additional figures select learning rates using the highest mean among configurations with three completed seeds.Source: §D.3 · p. 19, §D.4 · p. 20
- RL uses reward-weighted response-token loss with a leave-one-out baseline and no supervised target loss; RFT trains by negative log-likelihood on successful sampled trajectories.RFT excludes incorrect trajectories and trajectories whose decoded sequence contains a candidate answer.Source: §D.3 · p. 19
- SFT trains on demonstrations with an eviction mask, while self-distillation adds KL loss against a full-context pass of the current model as teacher.All four objectives use the same mask, allowing answer tokens to attend to sequence KVs but not the assignment turn.Source: §D.3 · p. 19
D.4 Additional Retrieval Results
- Figures 14 and 15 report mean accuracy over three seeds for in-distribution, prompt-generalization, zero-support-target, and actor-name conditions; cell labels give KV-cache capacity k.Figure 15 compares actor-name retrieval with and without listed options; Figure 14 includes the zero-support TV target.Source: Figure 14 · p. 20, Figure 15 · p. 20
E Prefill Token Cost
- For N generated tokens under budget B, each compaction retains r = o + s tokens; re-prefill compaction copies and prefills those retained tokens again.Here, o is overlap and s is generated summary length.Source: E Prefill Token Cost · p. 21
- ∆re-prefill = 0 with KV-streams, while re-prefill compaction costs N/3 when retaining only a B/4 summary; cost diverges as r approaches B.The re-prefill cost depends on total retained length, not its split between overlap and summary.Source: E Prefill Token Cost · p. 21