Skip to main content
Back to timeline
arXivSource publication:

Long-context RL training stops waiting on autoregressive rollouts: sparse self-drafting plus full-attention verification lifts throughput 1.39x

Synopsis

Long-context RLVR rollout generation now drafts tokens with sparse attention and verifies them with full attention, raising end-to-end training throughput to 1.39x and 1.35x of autoregressive on Qwen3-1.7B and Qwen3-8B while keeping comparable task performance.

Source-provided article image: S^3R: Fast and Faithful RL Rollouts via Self-Sparse Speculative Decoding
Figure 2 · arXiv

Interpretation

Long-context reinforcement learning with verifiable rewards no longer has to run full attention for every rollout token. S3R confines sparse attention to proposing candidate tokens, while the final responses come from full-attention verification and exact speculative sampling, so the response groups used for policy optimization follow the dense autoregressive distribution. Earlier attempts to cut rollout cost sampled directly from a sparse policy and compensated on the training side with filtering or importance reweighting; in group-relative optimization such as GRPO, each response's advantage depends on the rewards of other group members, so exact response-wise reweighting does not generally recover the dense-policy expected update. In a matched ablation at the same aggressive sparse drafting budget, S3R trains stably while direct sparse sampling with training-side correction shows larger gradient-norm fluctuations, a sharp spike around step 70, and reward deterioration.

Drafting and verification share the same policy parameters, so no separate drafter has to be trained or updated. Drafting selects a small set of historical KV entries using attention scores, verification recomputes the proposed block against the complete historical KV, and exact speculative sampling either accepts candidates or draws from the residual distribution. Existing speculative decoding schemes rely on retrieved historical responses, an external drafter network, or layer skipping; an external drafter goes stale as the policy keeps updating and loses acceptance, while retrieval depends on past trajectories covering the evolving policy's continuations. On Qwen3-1.7B and Qwen3-8B trained on DAPO-MATH-17k with veRL on a single node of eight H100 GPUs, S3R reaches block efficiencies of 5.59 and 5.47, above 1.31 and 1.29 for N-gram Lookup, 2.03 and 2.04 for History Retrieval, and 1.89 and 1.86 for a frozen EAGLE3 auxiliary proposer.

Dense verification still needs the complete historical KV cache, and that state can exceed GPU capacity as responses grow, so a KV-aware AIMD scheduler adapts request admission to limit repeated KV offloading and restoration while sustaining rollout parallelism. A permanently low concurrency limit reduces thrashing but sacrifices batching when memory is available, while recomputing evicted KV states adds generation cost. Over a ten-step Qwen3-1.7B window, the AIMD scheduler raises generation throughput from 3989.0 tokens/s/GPU for autoregressive to 6362.1, a further 5.8% over offloading at a fixed concurrency of 256, and cuts transferred pages per 1,000 generated tokens from 18.0 to 1.3.

Perspective

The method targets long-context reinforcement learning with verifiable rewards under group-relative optimization, trained on Qwen3-1.7B and Qwen3-8B with DAPO-MATH-17k on a single node of eight H100 GPUs in veRL, and evaluated with AIME26 and ArXiv-Math pass@1. It requires retaining the complete historical KV cache for dense verification, so it suits settings where memory pressure can be managed by KV-aware scheduling. Future directions include adaptive sparse budgets and context selection, and extensions to MoE and linear-attention models.

The sparse budget trades accepted progress against execution cost: raising it from 0.05 to 0.50 lifts acceptance from 80.4% to 96.7% while speedup falls from 1.370 to 1.058, and lowering it to 0.03 reduces both acceptance and speedup, so the default 0.05 is an empirical balance in this workload rather than a unique optimum. Acceptance declines across draft positions, and on Qwen3-8B the sixth position declines more than the first. Periodic online drafter adaptation did not produce a monotonic increase in acceptance in the reported configuration.