Partial Rollout Buffer in RL for LLMs

Updated 11 June 2026

Partial Rollout Buffer is a data structure that retains and manages incomplete trajectories in RL for LLMs, ensuring efficient use of computational resources.
It addresses challenges like long-tail rollout distributions and GPU idle times by decoupling sample collection from strict batch synchronization.
Empirical studies show that implementations such as APRIL, CBS, and WS-GRPO yield significant improvements in training throughput and model accuracy.

A partial rollout buffer is a data structure and scheduling mechanism used in reinforcement learning (RL) for LLMs that enables the retention, management, and efficient utilization of incomplete (partial) trajectories. Its central purpose is to address challenges posed by the long-tail distribution of rollout lengths, GPU idle time, and sample efficiency during RL training, particularly under large-scale RLHF or RLVR regimens. Recent practices have adopted partial rollout buffers to maximize hardware throughput, accelerate convergence, and, in some variants, transform sparse terminal rewards into dense, prefix-level guidance across RL formalisms including APRIL, Contextual Rollout Bandits, and WS-GRPO (Zhou et al., 23 Sep 2025, Lu et al., 9 Feb 2026, Mundada et al., 19 Feb 2026).

1. Motivation and Design Rationale

The long-tail distribution of LLM-generated rollout lengths results in batch-level synchronization inefficiencies, where a small fraction of extended rollouts can stall all batch members, causing substantial underutilization of GPU resources. This inefficiency is amplified as models and rollout sizes increase, driving RL practitioners to seek system-level remedies that avoid discarding partial work and that decouple sample collection from strict batch synchronization (Zhou et al., 23 Sep 2025). Furthermore, managing partial and historical rollouts in a buffer allows for improved sample reuse, adaptive rollout selection, and denser reward propagation strategies, particularly when only sparse correctness signals are available (Lu et al., 9 Feb 2026, Mundada et al., 19 Feb 2026).

2. Architectures and Operational Principles

Partial rollout buffer implementations differ according to their contextual use:

APRIL PartialRolloutBuffer: Buffers incomplete rollouts when over-provisioned requests are aborted upon achieving the target number of completed rollouts ( $N$ ). Each buffer entry retains a unique rollout ID, token prefix, model KV-cache, initiation step index, and policy version. At the subsequent RL step, entries are resumed from their saved state until rollouts are completed or again pre-empted, ensuring eventual completion and eliminating rollout waste (Zhou et al., 23 Sep 2025).
Replay Buffers with Truncation (CBS): Contextual Rollout Bandits store both complete and truncated rollouts (flagged via a binary feature). Each is annotated with multi-dimensional arm features (e.g., normalized length, usage count, sample age, policy-perception metrics). The scheduler adaptively selects high-value historical and partial rollouts for policy improvements. Feature tracking enables noise-aware, frequency-adaptive replay prioritization exposed to an online neural scheduler (Lu et al., 9 Feb 2026).
In-Batch Prefix Buffers (WS-GRPO): WS-GRPO constructs a per-optimization-step buffer consisting of all consecutive prefixes from sampled rollouts. There is no large replay memory, but each step collects and evaluates all prefixes in the minibatch. These prefixes are used to propagate dense pseudo-reward signals, mediating between sparse terminal correctness and per-token/local guidance (Mundada et al., 19 Feb 2026).

3. Algorithmic Procedures

A representative procedure from APRIL is summarized as follows (Zhou et al., 23 Sep 2025):

Over-provision rollout requests: Issue $N' = \alpha N$ requests, with $\alpha \in [1.5,2]$ .
Early termination: Upon $N$ completions, abort all remaining in-flight rollouts.
Partial rollout buffering: Extract states of aborted rollouts—including token sequence, KV-cache, and metadata—and enqueue in PartialRolloutBuffer.
Recycling: At each subsequent step, resume up to $M$ buffered partials before issuing new rollouts.
Policy update: Use only the $N$ completed rollouts for policy improvement.
Memory management: Typical buffers sustain up to $(\alpha-1)N$ entries per step.

For the neural rollout scheduler in CBS, each round updates the buffer, computes features (length, truncation), trains the scheduler, and selects rollouts to maximize a reward function tied to policy improvement. Truncated or incomplete trajectories are reweighted or downweighted according to their flagged status, enabling informed sample selection under a contextual bandit paradigm (Lu et al., 9 Feb 2026).

WS-GRPO, in turn, generates for each prompt a new buffer containing all prefix pairs, which are scored by a weakly-supervised preference model to compute prefix-level rewards. These guide policy optimization within each minibatch, sidestepping reliance on terminal-only rewards (Mundada et al., 19 Feb 2026).

4. Mathematical Models and Theoretical Guarantees

Table 1: Mathematical Forms and Empirical Gains

Buffer Variant	Key Metric	Empirical Impact
APRIL	Throughput $T$	$\Delta T\sim$ 24–44%
	Convergence step reduction	15–25% fewer steps
	Final accuracy gain	+2–8 percentage points
CBS	Regret $R_T$ (bandit formalism)	$N' = \alpha N$ 0
	Performance bound increase with buffer size	$N' = \alpha N$ 1
WS-GRPO	Generalization bound vs buffer/complexity	$N' = \alpha N$ 2

APRIL models the expected improvement in rollout throughput as $N' = \alpha N$ 3 and shows empirical gains of 24–44% on representative language reasoning tasks. GPU idle fraction, modeled as $N' = \alpha N$ 4, decreases in proportion to throughput gains. Convergence steps to reach a fixed evaluation target are reduced by 15–25%, with a +2–8% absolute gain in final accuracy (Zhou et al., 23 Sep 2025).

CBS guarantees sublinear regret under standard neural-bandit assumptions. The achievable performance upper bound scales monotonically with buffer size: larger buffers (and thus longer rollout histories, including partials) improve optimization opportunities (Lu et al., 9 Feb 2026).

WS-GRPO provides theoretical consistency bounds for prefix-level reward propagation, robustness to preference model error, and generalization guarantees matching $N' = \alpha N$ 5 rates, with low additional variance from prefix rewards (Mundada et al., 19 Feb 2026).

5. System-Level and Implementation Considerations

APRIL’s buffer requires memory management tuned to store up to $N' = \alpha N$ 6 in-flight partials per RL step. Each entry contains $N' = \alpha N$ 7 bytes for token and KV-cache states. Lock-free queues or ring buffers support high-throughput operation. Strict on-policy constraints necessitate discarding or pruning entries older than $N' = \alpha N$ 8 steps (Zhou et al., 23 Sep 2025).

Hardware-agnostic integration requires only the ability to abort GPU-side inference and extract KV-cache state. On NVIDIA, CUDA streams and vLLM abort APIs suffice; on AMD, similar functionality is achieved via HIP streams and custom memory snapshotting (see torch_memory_saver PR #43). Minimal changes are needed for frameworks like slime, with a PartialRolloutBuffer class handling all buffer manipulations.

CBS and related bandit-scheduling buffers employ a FIFO design with capacity $N' = \alpha N$ 9 and per-rollout arm features. Periodic featurization incurs $\alpha \in [1.5,2]$ 0 extra training cost. WS-GRPO’s buffer is in-batch and transient, rebuilt each optimization step and does not require persistent replay memory; it scales with the aggregate number of token- or sentence-level prefixes per batch (Mundada et al., 19 Feb 2026).

6. Empirical Results and Comparative Performance

On math reasoning benchmarks with Qwen3-4B and 8 AMD MI300 GPUs, APRIL achieves the following improvements over strong RLHF baselines (Zhou et al., 23 Sep 2025):

Algorithm	Throughput (Kt/s) Baseline	APRIL	Accuracy Baseline (%)	APRIL (%)	Throughput Gain	Accuracy Gain
GRPO (D17k)	50	67	62	68	+34%	+6
GRPO (DS)	48	63	59	65	+31%	+6
DAPO (D17k)	52	56	61	65	+8%	+4
GSPO (D17k)	51	70	60	66	+37%	+6

WS-GRPO, on ARC/GSM8K/CommonsenseQA/DeepMath, achieves 50–90% reduction in average rollout lengths (tokens or reasoning steps)—with at most 2–3 percentage points of accuracy lost, and sometimes no loss at all (e.g., GSM8K steady at 92.4%) (Mundada et al., 19 Feb 2026).

CBS, across six reasoning benchmarks, improves performance and efficiency by continually reusing (and adaptively deprioritizing) partial and out-of-sample rollouts, with monotonic gain in attainable policy performance as buffer size increases (Lu et al., 9 Feb 2026).

7. Applications and Implications

Partial rollout buffers are now standard in system-level RL optimizations for LLMs, permitting scalable, efficient, and accurate training at industrial scale (Zhou et al., 23 Sep 2025). Their use in contemporary frameworks (e.g., slime) is trivial to port across CUDA/HIP devices. Adaptive sampling schedulers (CBS-style) further boost efficiency by leveraging historical and incomplete rollouts rather than relying solely on freshly generated samples (Lu et al., 9 Feb 2026). When paired with reward propagation via weak supervision (WS-GRPO), partial buffers enable dense feedback even under sparse, delayed objectives, thus reducing redundancy and overthinking without significant loss in accuracy (Mundada et al., 19 Feb 2026).

The underlying mechanisms—over-provisioning, early abort and retention, feature-rich buffer storage, adaptive scheduler utilization, and prefix-level reward propagation—now form the backbone of RL pipelines for reasoning, coding, and alignment tasks in LLMs. They enable substantial improvements in both training dynamics and sample efficiency, while maintaining compatibility with diverse hardware and policy optimization schemes.