PrefixFlow: Training-Time KV Caching for Schedule-Level Prefix Reuse in LLM RL Training
arXiv:2606.01143
Abstract
In reinforcement learning (RL) post-training, one prompt fans out into trajectories, yet standard trainers replay the shared prefix's forward and backward pass times; in long-context retrieval-augmented generation (RAG) and multimodal workloads, this repetition dominates training compute. Existing designs reuse the prefix within one computation graph (a packed sequence, fused kernel, or tree traversal); under the trainer's memory limits and parallelism, their compute savings do not reliably become shorter steps. We present PrefixFlow, which takes the prefix out of that graph and into the training schedule through a training-time KV cache: Phase A runs one prefix forward and caches its keys and values (); Phase B trains suffixes as ordinary microbatches that read only and write back only their gradients (); Phase C runs one delayed prefix backward on the accumulated . Memory and parallelism then become manageable: dormant prefix activations are offloaded during Phase B to free high-bandwidth memory (HBM) for suffixes, and context, pipeline, and data parallelism each schedule the phases through a local rule while tensor and expert parallelism need no change, so saved compute becomes shorter steps under any microbatch split. On dense models, and on mixture-of-experts (MoE) models under Rollout Routing Replay, PrefixFlow matches full-sequence training within bfloat16 (BF16) tolerance. Against an identically configured baseline, it speeds up end-to-end verl RL training on a long-prefix RAG workload by ( on the actor update). In controlled sweeps, actor-side speedup grows with prefix ratio and group size, reaching at an 83% prefix ratio, while suffix-dominated splits gain little; offloading dormant prefix activations lowers Phase-B peak HBM by up to 59% and extends trainable sequences by .
arXiv comment15 pages. Accepted to the 2026 ACM SIGOPS Annual Technical Conference (ATC 2026)