LongStraw: Long-Context RL Beyond 2M Tokens under a Fixed GPU Budget

Changhai Zhou, Kieran Liu, Yuhua Zhou, Qian Qiao, Jun Gao, Harry Zhang, Irvine Lu, Nolan Ho, Lucian Li, Andrew Lei, Cleon Cheng, Steven Chiang, Yihang Zeng, Di Zhang, Rio Yang, Kaijie Chen, Andrew Chen, Pony Ma, Weizhong Zhang, Cheng Jin

LongStraw enables million-token RL post-training by serializing response branches to fit within fixed GPU memory.

How can we perform RL post-training on 2M+ token contexts while keeping GPU hardware requirements fixed?

RL post-training for long-context agents is bottlenecked by memory: training requires backpropagating through multiple responses, which forces the full prompt and all response graphs to reside in GPU memory simultaneously. LongStraw decouples this by evaluating the prompt once without gradients, storing only the essential architecture-specific state, and then replaying response branches one at a time. This approach allows training on 2-million-token contexts using a fixed GPU budget, trading additional replay time for significantly lower peak memory usage.

Paper Primer

LongStraw implements a "capture-once, replay-suffix" execution stack. It evaluates the shared prompt in a no-gradient pass, retains only the minimal state required for future tokens (such as recurrent states or compressed attention pages), and then processes each response branch sequentially under autograd.

LongStraw enables RL post-training at 2.1 million tokens on a fixed 8-GPU budget.

The system successfully executed grouped scoring and response backward passes for Qwen 3.6-27B at 2.1M positions, with group sizes of 2 and 8. Increasing the group size from 2 to 8 added only 0.21 GB to the peak allocated memory.

The architecture supports extreme context scaling beyond 4 million tokens.

A stress test extended the execution envelope to 4.46 million positions on the same hardware configuration.

Why does this approach use serial replay instead of parallelizing across more GPUs?

The paper targets a "fixed-budget" constraint, asking how far a specific accelerator allocation can go. Serial replay avoids the memory spike of holding all response graphs simultaneously, allowing researchers to train on long contexts without needing massive, expensive GPU clusters.

Does this method guarantee a mathematically correct distributed update?

No. The authors report "execution capacity" rather than full training correctness. The current implementation detaches the prompt state and lacks full synchronization for certain gradient paths, meaning it establishes that the system can run the update graph without crashing, but does not yet prove full-gradient parity with conventional training.

The Long-Context RL Bottleneck

LongStraw tackles the memory gap between inference‑time and RL post‑training contexts.

Inference servers now support million‑token prompts, yet RL post‑training workloads remain limited to roughly 256 K tokens, relying on length generalization at deployment. This mismatch is especially acute for AI agents that must accumulate observations, tool outputs, and prior decisions across long trajectories.

During GRPO training the same prompt is recomputed for every response in a group, so the model repeatedly materialises identical forward states and backward graphs.

LongStraw addresses this redundancy by evaluating the shared prompt once without automatic differentiation, storing only the model‑specific state needed for later tokens, and then replaying each short response branch sequentially under autograd.

The critical gap is between the long contexts supported at inference and the much shorter contexts feasible during RL post‑training.

The GRPO Update Graph

We detail the GRPO update graph and its serial replay mechanism for long-context RL.

When a GRPO update processes a group of responses, the prompt is recomputed for every member, causing activation memory to grow linearly with the group size $G$.

Prompt capture is like taking a single photograph of a scene and reusing that image for every subsequent edit – the model records the prompt once, then replays only the response tokens.

Compute the importance ratio $\rho_{i,t}=0.25/0.20=1.25$ for every token (log‑ratio exponentiates to the same value).

Clip $\rho_{i,t}$ to $[0.8,1.2]$, yielding $\text{clip}(\rho_{i,t})=1.2$.

Per‑token surrogate term $\min(\rho_{i,t}A_i,\text{clip}(\rho_{i,t})A_i)=\min(1.25,1.2)=1.2$.

Sum over the $2$ tokens of a member: $2\times1.2=2.4$.

Average over the $3$ members: $L_{\text{policy}} = -\frac{1}{3}\frac{1}{2}(2.4+2.4+2.4)= -1.2$.

This tiny example shows how the clipped ratio caps the contribution of any single token, keeping the group loss stable even when the new policy deviates moderately from the old one.

How does GRPO differ from the standard PPO objective used for single‑response updates?

Standard PPO treats each trajectory independently, normalizing by the total number of tokens. GRPO first normalizes inside each response ($R_i$) and then across the whole group ($G$), and it adds a token‑wise KL penalty weighted by $\beta$ to keep the new policy close to a reference policy while reusing the prompt state.

The paper validates the replay pipeline with four increasingly strict evidence levels: (1) all forward/backward/optimizer events finish without NaNs; (2) the conditional response operator matches the model definition; (3) sharded gradient contributions reduce correctly to the owning parameters; (4) full‑gradient parity with a conventional full‑sequence run.

Architecture Anatomy and Bottlenecks

Understanding how dense vs. MoE FFNs and token‑mixing choices shape prompt‑state bottlenecks.

Retaining prompt‑state tensors across a long context inflates memory; the architecture decides which layers become the bottleneck.

Dense FFNs treat every token the same way, while MoE routes each token to a handful of specialized experts – like a single chef cooking every dish versus a kitchen staffed with many chefs each handling a subset of orders.

Dense path: each token $h$ is multiplied by $W_1$, gated, then multiplied by $W_2$, producing an 8‑dim output for all 4 tokens.

Router computes scores $g(h)$ and selects expert 1 for tokens 1‑2, expert 2 for tokens 3‑4.

Expert 1 processes its two tokens with its own weight set, expert 2 does the same; the two results are concatenated with the shared baseline.

Overall FLOPs: dense layer $3\times8\times8\times4 = 768$; MoE layer $3\times8\times8\times2 = 384$ (half the arithmetic) plus a small routing cost.

MoE cuts compute but adds a routing step and requires moving the selected token slices to the expert’s device, which can dominate runtime at scale.

Why does MoE not automatically give a speedup despite fewer FLOPs?

Because the selected token‑expert pairs must be transferred across GPU boundaries; the communication latency often outweighs the arithmetic savings, especially when the batch contains many tokens.

GDN carries a fixed‑size recurrent state forward, like a conveyor belt that moves the same cargo each step; full attention builds a growing index of all past tokens, like a library where each new query must scan an ever‑larger catalog. MLA/DSA compresses that catalog into a latent map, then reuses it across layers.

GDN: after processing the 4‑token prompt, the recurrent state remains a $2$‑dim vector.

Full attention: KV pages contain $4$ keys and $4$ values, total $8$ scalar entries.

During response replay (2 new tokens), GDN reuses the same $2$‑dim state; no new memory is allocated.

Full attention must allocate $2$ new keys and $2$ new values for the response, growing the KV storage to $12$ entries.

GDN’s memory stays flat regardless of prompt length, while full attention’s memory grows linearly, which is the core source of the bottleneck.

Why can MLA/DSA avoid the linear growth that full attention suffers?

MLA compresses the entire prompt into a fixed‑size latent matrix, and DSA selects only a sparse subset of positions per query; both steps decouple memory usage from $P$, unlike full attention which stores a separate key/value for every token.

**Figure 3.** Two independent architecture axes. Blue bands encode token mixing and retained prompt state; green bands encode FFN execution and routing. Qwen combines 48 recurrent GDN and 16 full-attention layers with dense FFNs. GLM combines 21 index-computing and 57 IndexShare layers with three dense and 75 MoE FFNs. The orange boxes preserve the semantic boundary of each receipt: Qwen establishes global forward partitioning but not coherent distributed adapter updates; the accepted-receipt GLM DSA fallback is local to each context-parallel shard.

These architectural choices dictate where prompt‑state memory lives and therefore where LongStraw’s replay optimization must focus.

LongStraw Execution Design

LongStraw reshapes long‑context RL by keeping only prompt tensors that later responses need.

RL‑based long‑context training repeatedly re‑feeds the same prompt for every response in a group, inflating activation memory without adding new information.

Only the prompt tensors that a later response token will actually read are kept alive; everything else is released after the prompt pre‑fill.

How does LongStraw differ from a naïve KV‑caching scheme?

KV‑caching simply re‑uses the same key/value buffers across time steps; LongStraw goes further by discarding every activation that is not needed for any future response, turning the prompt into a read‑only snapshot and freeing the bulk of the $P$‑length activation graph.

Phase 1 – Prompt capture: run the prompt under $\theta_k$ with autograd disabled; save durable prompt tensors, release all transient activations.

Phase 2 – Score freeze: compute and store old/reference log‑probabilities for each group member; no model updates occur between members.

Phase 3 – Response replay: for each member $i$, rebuild the short response graph using the read‑only prompt state, back‑propagate its loss, then immediately discard the member’s graph.

Phase 4 – Gradient accumulation & optimizer step: sum the per‑member gradients locally and invoke a single optimizer call; the captured prompt state is now stale for the next update.

Phase 1: run tokens 1‑4 once, store KV pages for those four positions, discard all intermediate hidden states.

Phase 2: compute old log‑probs for both responses (no gradients), store two $4$‑element score vectors.

Phase 3‑a: replay response 1 using the saved KV pages, back‑propagate its 3‑token loss, then free its autograd graph.

Phase 3‑b: replay response 2 similarly with its 2‑token loss, then free its graph.

Phase 4: sum the two gradient contributions, apply one optimizer step, and mark the KV pages as stale for the next update.

Only the four prompt KV entries survive the whole update; all $P\!+\!R$ activations that would have existed in a naïve replay are eliminated, reducing peak memory roughly from $O(P\!+\!R)$ to $O(R)$.

**Figure 1** **LongStraw execution path.** The training runtime supplies a model snapshot and a group of responses. LongStraw evaluates the shared prompt once, stores model-specific state, then replays one response at a time and accumulates gradients. Qwen keeps recurrent state and sharded KV pages; GLM keeps CPU MLA/DSA pages and checkpoints its response layers. The orange boxes state the distributed reductions that the measured paths do not complete.

Table 2 enumerates durable versus transient prompt state. Durable rows (e.g., Qwen GDN recurrent state, GLM MLA latent pages) survive the no‑grad capture and are read‑only during response replay; transient rows (attention scratch, FFN activations) are released immediately after capture.

**Figure 4.** State lifetime across one grouped update. Durable Qwen state remains GPU-resident; GLM pages remain CPU-resident and only one layer is staged at a time. Prompt activations are released during no-grad capture, whereas only the current response graph is live under autograd. Group-indexed inputs and frozen scores may still grow with $G$. Worker-local gradients accumulate over $G$ serial backwards and each optimizer is called once. The optimizer changes $\theta_k$ to $\theta_{k+1}$, making every captured prompt state stale; a repeated loop would require recapture.

Device placement follows the durability contract: Qwen keeps its compact GDN and KV pages on GPU because eight‑way context parallelism fits the per‑rank memory budget; GLM keeps its large MLA and indexer pages on CPU and stages only the layer‑wise slices needed for a single response.

Whole‑layer checkpointing is applied only to the short response graph. For Qwen the KV storage per rank is $O(P/8)$; for GLM it is $O(P/32)$ per rank and $O(P/(32\times64))$ page records for IndexShare layers. The prompt itself is never checkpointed because it was already evaluated without autograd.

Parallel layout (Table 3) separates prompt‑parallelism (CP) from expert‑parallelism (EP). Qwen’s CP8 distributes full‑attention KV pages; GLM’s CP32 distributes MLA pages while EP32 distributes routed experts. The two dimensions are orthogonal, so response‑time collectives differ: Qwen performs a global LSE/normalizer merge, GLM performs local EP dispatch/combine without a cross‑CP attention merge.

Adapter scope is fixed: GLM uses rank‑8 LoRA on attention, FFNs, and expert modules; Qwen uses NF4 QLoRA on its dense targets. All base weights, embeddings, and routing parameters remain frozen, so LongStraw’s memory savings are orthogonal to these parameter‑efficiency tricks.

Qwen: Dense Hybrid Replay

Qwen processes over two million tokens within an eight‑GPU H20 budget, scaling modestly with group size.

Increasing the response group size from $G=2$ to $G=8$ adds 1,586 seconds of runtime but only 0.208 GB of peak memory.

Measured end‑to‑end times of 5,198.78 s (G=2) and 6,785.23 s (G=8) on eight H20 GPUs.

The runtime grows because each additional response member incurs a full backward replay, while the memory footprint barely changes since the prompt KV pages are already resident.

The prompt’s KV tensors are retained on each rank and reused across multiple response generations, so only the response tokens are recomputed.

How does Qwen Hybrid Replay differ from standard KV‑caching?

Standard KV‑caching stores a single prompt KV set and reuses it for one response; Qwen Hybrid Replay retains the KV pages across multiple responses and interleaves gradient‑tracked replay of response blocks, enabling many responses without re‑encoding the prompt.

**Figure 5** Qwen CP8 forward fidelity and backward synchronization gap. The full-attention forward partitions physical KV pages and recomposes the global softmax over all shards, subject to the BF16 numerator reduction. Backward all-reduces the query gradient, but K/V-derived parameter and hidden-gradient contributions remain rank local before eight local optimizer calls. The receipt therefore establishes global conditional-forward and update-shaped execution capacity, not a coherent distributed model update.

Qwen can handle over two million tokens within an eight‑GPU H20 budget.

GLM: Sparse Attention Replay

GLM replays retained prompt-state tensors to answer grouped responses without re‑processing the whole prompt.

The LongStraw premise is that RL‑driven grouped responses waste memory by re‑encoding the prompt for every response; GLM solves this by keeping the prompt‑state tensors alive across the prompt boundary and replaying only the response tokens.

GLM splits attention into a dense latent‑space component (MLA) and a sparse index‑based component (DSA), each of which can be materialized once per prompt and reused across multiple responses.

How does GLM’s dynamic sparse attention differ from ordinary dense attention?

Dense attention computes a full P × P similarity matrix for every layer, scaling quadratically with prompt length. GLM’s DSA replaces that matrix with a top‑k lookup that touches only 2,048 keys per response, while the dense MLA component is computed once and cached, so the overall cost grows linearly with the number of responses, not with the prompt size.

MLA projects each token into a 576‑dim latent vector and stores it as a page (8 × 576 values).

DSA builds an index for the first response: each of the 2 heads scores the 8 keys and selects the top‑2 positions, yielding an index tensor of shape [2, 2].

The second response reuses the stored MLA pages and recomputes only its own DSA index (again top‑2).

Both responses attend over the same 8 × 576 MLA pages, so the expensive projection is performed once.

Retaining the MLA pages eliminates the O(P·H) projection cost for every response; only the cheap DSA index (O(k·H)) is recomputed.

The core trick is to materialize the prompt‑side MLA and DSA pages once, keep them resident on CPU, and during each response forward only materialize the tiny response‑side tensors.

Why not simply recompute the prompt attention for each response instead of retaining pages?

Recomputing would require projecting the full prompt (≈2 M tokens) through the MLA projection on every response, which would explode GPU memory and runtime (quadratic in prompt length). Retaining the pages moves the heavy projection to CPU once, keeping the per‑response GPU footprint constant and enabling the 32‑GPU budget to be respected.

Capture MLA pages for the 4 tokens (4 × 576 values) and store them on CPU.

First response token queries the DSA indexer, which selects the top‑2 keys from the 4‑token pool, producing an index tensor [1, 2, 2].

The response forward loads the 4 × 576 MLA pages, applies the cached projection, and uses the DSA index to attend only to the selected keys.

Second response repeats the DSA indexing step; the MLA pages are reused without any additional projection.

The per‑response GPU memory never exceeds the size of the two response hidden states plus the tiny DSA index, regardless of how many prompt tokens were originally captured.

**Figure 6.** GLM CP32 ownership and the two missing global operations. Megatron zigzag partitioning gives each rank two mirrored 512-page chunks. The accepted-receipt DSA fallback selects top-2,048 independently inside each local 65,536-token shard. A faithful full-context operator needs a cross-rank candidate merge, selected-value exchange, and one composed response hidden state. Separately, the accepted-receipt backward writes local DDP gradient buffers and calls the distributed optimizer without Megatron's gradient-finalization reduction. Parameter all-gather may restore replica equality, but it cannot reconstruct the omitted CP gradient sum.

**Figure 7.** GLM resident layer replay and whole-layer checkpointing. CPU pages retain only the CP-local prompt state. A compute layer stages MLA and DSA-key pages and publishes a local top-2,048 tensor; an IndexShare layer stages MLA pages and consumes the matching per-forward selection. The live attention projection and dense or EP32 MoE tail execute inside one checkpoint boundary. Backward traverses layers in reverse, restages and recomputes one layer, emits the audited tensor shapes, and releases its workspace. Page shapes and byte counts are contract-derived; they are not whole-step peak measurements.

During the initial forward pass, each rank captures MLA pages ([1, 64, 1, 576]) and, for index‑computing layers, DSA key pages ([1, 64, 128]).

Captured pages are written to the page manager, which maps logical prompt positions to fixed‑size CPU pages.

When generating a response token, the rank loads its local prompt pages, reuses the cached MLA values, and runs the DSA indexer to produce a top‑2,048 index tensor.

IndexShare layers read the previously published top‑k tensor instead of recomputing the index.

The layer’s sparse attention, output projection, and MoE tail are executed, producing a short response hidden state ([2, 1, 6,144]).

After the response forward, the short hidden state is kept for the next layer; the large prompt pages remain on CPU until the final backward.

During backward, each layer restages its CPU pages, recomputes sparse attention and MoE tails, and releases workspace, while gradients flow only through the retained response tensors.

GLM Scaling and Constraints

Ablation study isolates each scaling component, quantifying its memory impact and functional necessity.

This section systematically removes each scaling ingredient to expose its quantitative contribution.

Removing prompt‑state retention forces a fully expanded routed hidden buffer that consumes 6 GiB, exceeding the per‑rank H20 budget.

Stage I observed a 6 GiB buffer at a 65 K‑token shard.

Omitting autograd while retaining the full 2,097,152‑token prompt succeeds in a prefill‑only run, proving that prompt‑state capacity alone is not the memory limiter.

Stage II prefix‑only path processed the entire 2.1 M‑token prompt.

At 1 M tokens the GPU‑resident state and MoE work become prohibitive, indicating that beyond‑1 M scaling requires CPU‑resident paging.

Stage III layer‑0 replay showed costly GPU memory at the 1 M‑token mark.

Running all 78 layers at 64 K tokens uncovers index‑share lifetime and view‑contract bugs, confirming that full‑layer execution is fragile without the fixes introduced later.

Stage IV all‑layer runs at 32 K and 64 K exposed these failures.

The parallel ownership topology relies on exactly 32 ranks (TP1/CP32/EP32/ETP1/PP1); deviating from this rank count would break the budgeted placement.

Stage V describes the 32‑rank allocation.

A single‑member (G = 1) canary run succeeds at the full 2,097,152‑token prompt within the fixed H20 budget, showing that group accumulation, not prompt length, is the remaining bottleneck.

Stage VI reports the successful canary execution.

Even with balanced placement, per‑rank memory usage varies by 32.577 GB (74.7–96.3 % of device capacity), highlighting load‑balance inefficiencies.

Stage VII measured allocation spread from 112.571 GB to 145.148 GB per rank.

Execution Receipts and Trace Evidence

Execution receipts expose the true cost and consistency of fixed‑budget runs.

Qwen’s per‑response wall‑time drops from ≈ 2 600 s to ≈ 850 s when the response group grows four‑fold (G = 2 → 8), i.e. about a three‑fold speedup.

Wall times per response are 2 599.390 s (G = 2) versus 848.153 s (G = 8) on identical eight‑H20 hardware.

These two metrics together demonstrate that fixed‑budget execution can scale response groups without exploding compute or memory, but the wall‑time comparison must be read with the hardware‑difference disclaimer in mind.

They are terminal logs that certify each worker has reached the required scoring, backward, and optimizer‑call boundaries, yet they do not guarantee a globally coherent distributed update.

How do Execution Receipts differ from a full‑gradient‑parity verification?

Receipts stop at the local AdamW call on each rank; they do not include the all‑reduce of K/V adapter gradients or a synchronized parameter update, which is required for full‑gradient parity.

**Table 5.** Audited fixed-budget execution receipts. Qwen combines a 2,088,960-position prompt with 8,192 response inputs for an exact 2,097,152 context positions; GLM's prompt alone contains 2,097,152 positions. The fixed allocations are eight H20s for Qwen and 32 H20s for GLM. Terminal evidence records worker-local events, not a coherent distributed parameter update. Hardware and suffix workloads differ, so wall times are not comparable.

**Table 8.** Audited evidence matrix. Execution capacity records that the requested program reached its terminal optimizer calls; it does not imply a faithful distributed forward, a synchronized update, or conventional full-gradient parity.

Fixed-Budget Systems Lessons

Limits of fixed‑budget execution across Qwen and GLM systems.

Recall that LongStraw keeps prompt‑state tensors alive across the prompt boundary, replaying only response tokens to stay inside a fixed accelerator budget.

9.1 Capacity gains stem from letting prompt‑related intermediates die after the prompt is captured; the live autograd set then scales with response length, not with the full prompt.

9.2 The algorithm only benefits when the logical page table matches physical memory ownership; copying pages into correctly sized tensors moves the 2 M‑position prefix peak into the feasible range.

9.3 In Qwen the dense FFN dominates peak storage, while MoE routes shift the peak to expert‑selection structures; both paths still require every prompt token to traverse every decoder layer.

9.4 Context parallelism (splitting token history) and expert parallelism (splitting FFN parameters) address orthogonal resources; their collectives cannot substitute for one another in forward or backward passes.

9.5 Qwen’s forward pass is partition‑correct but its shard‑local adapter gradients are not reduced; GLM’s DSA path keeps gradients local as well, so neither matches a full‑sequence gradient without an extra reduction step.

9.6 Increasing the GRPO group size from 2 to 8 barely raises peak memory (+0.208 GB) while multiplying post‑prefix work; the wall‑time per additional member drops to ~264 s, showing that group cardinality is a scheduling knob rather than a memory limiter.

**Figure 9.** Qwen group accounting within eight H20s. The $G = 2$ and $G = 8$ bars are serial-loop anchors, not group-size limits. Four times as many members add 1,586.445 s but only 0.208 GB (0.213%) at peak because the shared prefix dominates residency; the final segment is a phase-sum residual, not a pure optimizer timer.

Read the original paper

Open the simplified reader on Paperglide

Browse all simplified papers