Paper 2607.14952
LongStraw: Long-Context RL Beyond 2M Tokens under a Fixed GPU Budget
- Published
- Jul 2026
- Research lab
- Independent
- Citations
- 0
- GitHub
- 43 stars
01 In brief
Summary
LongStraw is a system for long-context GRPO post-training under fixed GPU budgets, addressing the memory bottleneck of multi-million-token prompts by virtualizing resident state, replaying responses, and executing distributed gradients.
It captures the shared prompt without autograd, retains only architecture-required state (e.g., recurrent state, KV pages, MLA/DSA pages), and replays one response at a time with autograd, accumulating gradients before a single optimizer step.
This bounds live graph memory by the response suffix while reusing prompt computation across the group.
The system is instantiated for two models: Qwen3.6-27B (hybrid GDN/full-attention) and GLM-5.2 (MLA/DSA/MoE).
On eight H20 GPUs, Qwen completes exact-attention response-only steps at 2,097,152 positions for G=2 and G=8, with group size increase adding only 0.208 GB peak memory.
At 4,456,448 positions, one resident prefix supports eight G=8 optimizer cycles (64 replays) at 83.894 GB per rank.
On 32 H20 GPUs, GLM-5.2 completes the full exact-2M online GRPO workflow, including vLLM rollout, reward computation, global DSA, two 78-layer backwards, gradient finalization, and one optimizer step.
These results demonstrate that practical context capacity is governed by resident-state lifetime, response replay, topology handoff, and distributed ownership, not just attention kernels.
02 From the paper
Abstract
Long-context RL post-training is constrained by the lifetime of state and gradients, not attention cost alone. In GRPO, one multi-million-token prompt must serve old-policy and reference scoring plus multiple policy responses, while conventional autograd keeps the prompt graph and all response graphs live alongside model weights, caches, and distributed communication buffers. We present LongStraw, an objective-aware, architecture-aware system for resident-state virtualization, response replay, and distributed-gradient execution. Its transaction captures the shared prompt without autograd, retains only the architecture-required state on explicitly owned pages, restores that state for each group member, scores old/reference branches without a graph, replays one policy response at a time with autograd, and accumulates the resulting gradients before one distributed finalization and optimizer step. This schedule bounds the live training graph by the response suffix while reusing the expensive prompt computation across the complete GRPO group. We instantiate this design for two incompatible model structures. Qwen3.6-27B combines 48 recurrent GDN layers with 16 full-attention layers; LongStraw keeps the compact recurrent state and physically CP8-sharded KV pages, composes global attention through cross-rank LSE/output merging, and performs blockwise response replay. GLM-5.2 combines a 78-layer MLA/DSA attention stack with a 256-expert, top-8 MoE tail. Its implementation keeps CP-sharded MLA latent pages and DSA indexer-key pages in CPU memory, stages one layer at a time, reconstructs IndexShare-aware global sparse selection over CP32, and dispatches routed response tokens over EP32. The two paths share one transaction contract while specializing the retained state, replay operator, and collective communication to the architecture...