Papers
Topics
Authors
Recent
Search
2000 character limit reached

LongStraw: Execution Stack for Long-Context RL

Updated 18 July 2026
  • LongStraw is an architecture-aware execution stack for long-context RL post-training that decouples prompt capture from response gradients to reduce memory load.
  • It employs grouped GRPO updates with selective state retention and serial response replay, enabling scalable execution on million-token contexts.
  • The approach balances memory and compute trade-offs while contrasting with alignment methods that focus on output distribution consistency across varying context lengths.

Searching arXiv for the LongStraw paper and closely related long-context work to ground the article in current literature. LongStraw denotes an architecture-aware execution stack for long-context reinforcement-learning post-training under a fixed GPU budget, introduced for Group Relative Policy Optimization (GRPO) and designed to make prompts at and beyond 2 million tokens executable during RL updates without keeping the full prompt-plus-responses autograd graph live (Zhou et al., 16 Jul 2026). In adjacent long-context literature, the name has also been used informally for a distinct perspective centered on long-short alignment, namely the consistency of output distributions across sequences of varying lengths, measured by a Long-Short Misalignment metric and optimized through a training-time regularizer (Du et al., 13 Jun 2025). The primary usage in the 2026 paper is therefore infrastructural and systems-oriented, whereas the 2025 alignment work is methodological and output-distribution-oriented. Both usages address the widening gap between nominal inference context length and effective long-context behavior, but they operate at different levels of the stack.

1. LongStraw as a long-context RL execution stack

The 2026 formulation defines LongStraw as an execution strategy for RL post-training workloads in which the shared prompt is evaluated once without autograd, only model-specific prompt state needed by later tokens is retained, and short response branches are replayed one at a time under autograd (Zhou et al., 16 Jul 2026). The motivating problem is that inference systems are approaching million-token contexts, while RL post-training workloads often remain at 256K tokens or below and rely on length generalization at deployment. This gap is especially salient for AI agents, whose observations, tool outputs, documents, and prior decisions accumulate over long trajectories (Zhou et al., 16 Jul 2026).

Within a grouped GRPO update, the long prompt x1:Px_{1:P} is shared across multiple sampled responses {yi}i=1G\{y_i\}_{i=1}^G. Standard autograd would keep prompt and response activations live across the group, making GPU memory the critical bottleneck. LongStraw changes the autograd boundary. Its grouped update consists of prompt capture without gradients, pre-step scoring under old and reference policies, group-normalized advantage construction, policy replay on one short response suffix at a time, and a single optimizer step after all member-specific backwards have completed (Zhou et al., 16 Jul 2026).

The resulting memory-time trade-off is explicit. The live memory is approximated as

$M_{\text{live} \;\approx\; M_{\text{fixed} + M_{\text{prefix-state}(P) + M_{\text{grad} + \max_i M_{\text{branch}(R_i) + M_{\text{score}\!\left(\sum_i R_i\right),$

while the update time is roughly

$T_{\text{update} \;\approx\; T_{\text{prefix}(P) + \sum_{i=1}^{G} T_{\text{score+replay}(R_i).$

This means group size primarily increases wall-time rather than peak memory, because only one short response branch remains live at a time (Zhou et al., 16 Jul 2026).

The stack is explicitly architecture-aware. It is instantiated for Qwen3.6-27B, which combines hybrid recurrent GDN layers with full self-attention, and for GLM-5.2, which combines MLA/DSA compressed sparse attention with MoE feed-forward blocks (Zhou et al., 16 Jul 2026). A plausible implication is that LongStraw is not a model-agnostic wrapper in the narrow sense; its retained prompt state, replay schedule, and device placement policy depend on architectural details.

2. GRPO under extreme context length

LongStraw is instantiated with GRPO, a PPO-style RL objective that groups several responses for the same prompt, normalizes rewards within the group, and uses a reference policy for KL regularization without a learned critic (Zhou et al., 16 Jul 2026). For response ii and timestep tt, the importance ratio is

$\rho_{i,t}(\theta)= \exp\!\left( \log\pi_\theta(y_{i,t}\mid x_{1:P},y_{i,<t})- \log\pi_{\mathrm{old}(y_{i,t}\mid x_{1:P},y_{i,<t}) \right),$

and the clipped policy loss is

$\mathcal{L}_{\mathrm{policy} = -\frac{1}{G}\sum_{i=1}^{G}\frac{1}{R_i}\sum_{t=1}^{R_i} \min\!\left( \rho_{i,t}A_i, \operatorname{clip}(\rho_{i,t},1-\epsilon,1+\epsilon)A_i \right),$

with an additional token-wise KL term to a reference policy with weight β\beta (Zhou et al., 16 Jul 2026).

At long context, the difficulty is not the GRPO objective itself but the execution of its conditioning. Old-policy, current-policy, and reference-policy log-probabilities all depend on the same very long prompt state. LongStraw therefore freezes the prompt-conditioned state during scoring, computes old and reference scores before any in-group optimizer step, and then replays each response suffix under autograd against stored prompt state zˉP\bar z_P (Zhou et al., 16 Jul 2026).

This does not alter GRPO semantics at the objective level. The paper states that LongStraw’s contribution is to restructure how log-probabilities and gradients are executed, without changing the objective itself (Zhou et al., 16 Jul 2026). The system-level novelty lies in making grouped RL scoring and response backward feasible at 2.1M positions for Qwen on eight H20 GPUs, with groups of 2 and 8, and validating the end-to-end execution path for a 2.1M-token prompt across all 78 layers of GLM-5.2 on 32 H20 GPUs (Zhou et al., 16 Jul 2026).

3. Architecture-specific state retention and replay

LongStraw’s defining systems idea is selective state retention. For Qwen3.6-27B, the retained prompt state consists of a compact recurrent state and short convolution tail for each GDN layer, whose size is fixed with respect to prompt length, together with sharded KV pages for each full-attention layer, whose storage grows linearly with prompt length {yi}i=1G\{y_i\}_{i=1}^G0 (Zhou et al., 16 Jul 2026). Everything else from the prompt forward is released, including FFN activations, attention scratch, and intermediate hidden states (Zhou et al., 16 Jul 2026).

For Qwen at prompt length {yi}i=1G\{y_i\}_{i=1}^G1, page size {yi}i=1G\{y_i\}_{i=1}^G2, CP world size {yi}i=1G\{y_i\}_{i=1}^G3, and 16 full-attention layers, the paper gives

{yi}i=1G\{y_i\}_{i=1}^G4

and estimates retained KV storage per rank as approximately {yi}i=1G\{y_i\}_{i=1}^G5 GiB (Zhou et al., 16 Jul 2026). The global dense-attention forward on sharded KV is then reconstructed through stable log-sum-exp composition: {yi}i=1G\{y_i\}_{i=1}^G6 followed by

{yi}i=1G\{y_i\}_{i=1}^G7

The paper states that this yields partition-correct global dense attention over the sharded KV, modulo BF16 merge (Zhou et al., 16 Jul 2026).

Qwen response replay is blockwise. The long prompt is processed in 510 chunks of 4,096 tokens without gradients, while each response is split into 4 blocks of 2,048 tokens. For each member, LongStraw first runs a no-grad suffix pass, recording GDN state at block boundaries and appending KV pages, then traverses response blocks in reverse order under whole-layer checkpointing, backpropagating block by block and popping block-specific KV after use (Zhou et al., 16 Jul 2026).

For GLM-5.2, the retained state is different. Each rank stores MLA latent pages for all 78 layers and DSA indexer-key pages for the 21 index-compute layers. With page size {yi}i=1G\{y_i\}_{i=1}^G8 and prompt length {yi}i=1G\{y_i\}_{i=1}^G9, the prompt state is stored on CPU; the paper computes per-rank CPU prompt state as

$M_{\text{live} \;\approx\; M_{\text{fixed} + M_{\text{prefix-state}(P) + M_{\text{grad} + \max_i M_{\text{branch}(R_i) + M_{\text{score}\!\left(\sum_i R_i\right),$0

or roughly $M_{\text{live} \;\approx\; M_{\text{fixed} + M_{\text{prefix-state}(P) + M_{\text{grad} + \max_i M_{\text{branch}(R_i) + M_{\text{score}\!\left(\sum_i R_i\right),$1 GiB across 32 ranks (Zhou et al., 16 Jul 2026). During replay, layer-specific prompt pages are staged from CPU to GPU, used to rebuild sparse attention and the short response path, and then released (Zhou et al., 16 Jul 2026).

GLM also exposes an MoE-specific activation bottleneck. On one CP shard of $M_{\text{live} \;\approx\; M_{\text{fixed} + M_{\text{prefix-state}(P) + M_{\text{grad} + \max_i M_{\text{branch}(R_i) + M_{\text{score}\!\left(\sum_i R_i\right),$2 tokens, top-8 routing expands to

$M_{\text{live} \;\approx\; M_{\text{fixed} + M_{\text{prefix-state}(P) + M_{\text{grad} + \max_i M_{\text{branch}(R_i) + M_{\text{score}\!\left(\sum_i R_i\right),$3

expert rows, and one BF16 hidden buffer for this routed set is estimated at approximately $M_{\text{live} \;\approx\; M_{\text{fixed} + M_{\text{prefix-state}(P) + M_{\text{grad} + \max_i M_{\text{branch}(R_i) + M_{\text{score}\!\left(\sum_i R_i\right),$4 GiB (Zhou et al., 16 Jul 2026). The paper presents this as evidence that sparsity alone does not resolve long-context autograd memory pressure; prompt detachment and whole-layer checkpointing over the short response are required (Zhou et al., 16 Jul 2026).

4. Reported execution capacity

The paper presents execution-capacity results rather than full training-validation results. For Qwen3.6-27B on eight H20 GPUs, grouped GRPO scoring and response backward complete at total context length $M_{\text{live} \;\approx\; M_{\text{fixed} + M_{\text{prefix-state}(P) + M_{\text{grad} + \max_i M_{\text{branch}(R_i) + M_{\text{score}\!\left(\sum_i R_i\right),$5 tokens for both $M_{\text{live} \;\approx\; M_{\text{fixed} + M_{\text{prefix-state}(P) + M_{\text{grad} + \max_i M_{\text{branch}(R_i) + M_{\text{score}\!\left(\sum_i R_i\right),$6 and $M_{\text{live} \;\approx\; M_{\text{fixed} + M_{\text{prefix-state}(P) + M_{\text{grad} + \max_i M_{\text{branch}(R_i) + M_{\text{score}\!\left(\sum_i R_i\right),$7 (Zhou et al., 16 Jul 2026). The reported measurements are as follows:

Configuration Wall time Peak allocated memory
Qwen, 2.1M context, $M_{\text{live} \;\approx\; M_{\text{fixed} + M_{\text{prefix-state}(P) + M_{\text{grad} + \max_i M_{\text{branch}(R_i) + M_{\text{score}\!\left(\sum_i R_i\right),$8 5198.780 s 97.503 GB
Qwen, 2.1M context, $M_{\text{live} \;\approx\; M_{\text{fixed} + M_{\text{prefix-state}(P) + M_{\text{grad} + \max_i M_{\text{branch}(R_i) + M_{\text{score}\!\left(\sum_i R_i\right),$9 6785.225 s 97.711 GB
Qwen, group-size increase $T_{\text{update} \;\approx\; T_{\text{prefix}(P) + \sum_{i=1}^{G} T_{\text{score+replay}(R_i).$0 +1586.445 s +0.208 GB

The increase of only $T_{\text{update} \;\approx\; T_{\text{prefix}(P) + \sum_{i=1}^{G} T_{\text{score+replay}(R_i).$1 GB, or $T_{\text{update} \;\approx\; T_{\text{prefix}(P) + \sum_{i=1}^{G} T_{\text{score+replay}(R_i).$2, when moving from group size $T_{\text{update} \;\approx\; T_{\text{prefix}(P) + \sum_{i=1}^{G} T_{\text{score+replay}(R_i).$3 to $T_{\text{update} \;\approx\; T_{\text{prefix}(P) + \sum_{i=1}^{G} T_{\text{score+replay}(R_i).$4 is used to support the claim that group size mainly scales time rather than peak memory under member-serial replay (Zhou et al., 16 Jul 2026). The same paper reports a separate stress test reaching $T_{\text{update} \;\approx\; T_{\text{prefix}(P) + \sum_{i=1}^{G} T_{\text{score+replay}(R_i).$5 positions on the same eight-H20 budget, and a prefix-frozen response-only run performing eight optimizer steps over the same $T_{\text{update} \;\approx\; T_{\text{prefix}(P) + \sum_{i=1}^{G} T_{\text{score+replay}(R_i).$6M context with peak $T_{\text{update} \;\approx\; T_{\text{prefix}(P) + \sum_{i=1}^{G} T_{\text{score+replay}(R_i).$7 GB (Zhou et al., 16 Jul 2026).

For GLM-5.2 on 32 H20 GPUs, the paper validates the end-to-end LongStraw execution path for a $T_{\text{update} \;\approx\; T_{\text{prefix}(P) + \sum_{i=1}^{G} T_{\text{score+replay}(R_i).$8M-token prompt across all 78 layers with $T_{\text{update} \;\approx\; T_{\text{prefix}(P) + \sum_{i=1}^{G} T_{\text{score+replay}(R_i).$9, though the main table reports peak GPU memory as not reported because the PyTorch counter was read between prefix capture and replay rather than across the whole transaction (Zhou et al., 16 Jul 2026). During prompt capture, rankwise max_memory_allocated ranged from ii0 GB to ii1 GB, indicating non-uniform allocation patterns (Zhou et al., 16 Jul 2026).

These results establish execution capacity rather than complete training correctness. The paper is explicit that the experiments show that one can capture ii2M+ prompts, run grouped GRPO scoring and response backwards, and issue optimizer calls without OOM, but do not yet demonstrate exact gradient equivalence to conventional full-sequence training (Zhou et al., 16 Jul 2026).

5. Correctness limits, detached prompts, and incomplete distributed gradients

A central limitation is that LongStraw detaches the captured prompt state. If ii3 denotes the full model state after the prompt, full-sequence training would require

ii4

LongStraw stores ii5 and computes only

ii6

which the paper describes as the response-conditional gradient (Zhou et al., 16 Jul 2026). The missing second term means the method is not equivalent to full-sequence training, and prompt state must be recaptured after each optimizer step because parameter updates make the stored state stale (Zhou et al., 16 Jul 2026).

The paper also identifies unfinished distributed gradient composition. In Qwen, forward CP8 global attention composition is correct, and query gradients are summed across ranks,

ii7

but K/V gradients remain local to each shard. Since LoRA K/V projection weights are replicated across ranks, the required reductions

ii8

are missing, so eight independent AdamW optimizers may produce inconsistent adapter updates (Zhou et al., 16 Jul 2026).

For GLM, the historical acceptance run used DSA sparse attention locally per CP shard rather than globally across shards, and the custom training loop bypassed finalize_model_grads(), so CP-replicated non-expert adapters were updated using unreduced local gradients (Zhou et al., 16 Jul 2026). The paper further notes that no full gradient parity against a short-context reference implementation was performed for either Qwen or GLM (Zhou et al., 16 Jul 2026).

These caveats are not incidental. They define the paper’s scope. LongStraw demonstrates that the execution path fits in memory and traverses all intended operations at million-token scale, but it does not yet establish optimizer-correct distributed RL training under the same execution schedule (Zhou et al., 16 Jul 2026).

6. Relation to long-short alignment and other long-context methods

The term “LongStraw” is also associated, in a different 2025 line of work, with long-short alignment: the idea that good long-context behavior is largely about making a model’s output distributions invariant to irrelevant changes in input length (Du et al., 13 Jun 2025). In that formulation, an autoregressive model ii9 should produce similar next-token distributions for suffixes tt0 and tt1 when the relevant contextual information is effectively the same (Du et al., 13 Jun 2025).

The corresponding discrepancy is measured by Symmetric Cross-Entropy,

tt2

and Long-Short Misalignment is defined as

tt3

The training objective adds this term to standard LM loss,

tt4

with best performance reported for tt5 and modest overhead of approximately tt6–tt7 (Du et al., 13 Jun 2025).

This alignment-based usage differs sharply from the 2026 LongStraw stack. The 2025 method is a training-time regularizer that does not modify positional encodings or architecture, whereas the 2026 method is an execution stack for RL post-training that also does not modify the underlying GRPO objective (Du et al., 13 Jun 2025, Zhou et al., 16 Jul 2026). The connection is conceptual rather than procedural: both respond to the inadequacy of relying solely on nominal context length. One addresses output-distribution consistency during training; the other addresses execution feasibility during RL updates at million-token scale.

Related long-context work also includes training-free length extrapolation methods such as Dimension-Wise Positional Embeddings Manipulation (DPE), which operates by measuring effective relative distance per RoPE dimension group and remapping only key dimensions at inference time (Lu et al., 26 Apr 2025). DPE is described as belonging to the “train short, test long” family and is evaluated on RULER, multi-needle NIAH, and InfiniteBench, which target retrieval, variable tracing, aggregation, and QA over long contexts (Lu et al., 26 Apr 2025). This suggests that the broader LongStraw landscape spans at least three levels: output-space regularization (Du et al., 13 Jun 2025), inference-time positional manipulation (Lu et al., 26 Apr 2025), and systems-level RL execution restructuring (Zhou et al., 16 Jul 2026).

7. Significance and open directions

LongStraw, in its 2026 sense, shifts attention from nominal context-window support to the systems question of what can actually be executed during RL post-training on a fixed hardware budget (Zhou et al., 16 Jul 2026). Its key contribution is to frame state lifetime and physical ownership as the controlling variables for million-token RL, rather than treating long-context feasibility as only a matter of attention FLOPs or sequence parallelism (Zhou et al., 16 Jul 2026). This is particularly relevant for agents, where prompts are long trajectories rather than single static documents.

At the same time, the paper’s own framing is deliberately cautious. The demonstrated result is execution capacity, not complete training correctness, downstream task improvement, or parity with conventional full-sequence gradients (Zhou et al., 16 Jul 2026). The workloads are synthetic, rewards are deterministic functions of group index, old and reference policies are identical at the first step so KL is zero and clipping is inactive, and no multi-step RL training curve or environment-integrated evaluation is reported (Zhou et al., 16 Jul 2026).

Open directions are therefore concrete. The paper identifies the need to restore missing distributed reductions, validate full-gradient parity at short context, and extend the execution stack from capacity receipts to correct large-scale RL optimization (Zhou et al., 16 Jul 2026). A plausible implication is that LongStraw’s long-term importance may depend less on the current detached-prefix implementation than on whether its prompt-capture and serial-replay design can be combined with exact distributed gradient composition.

In the broader literature, the coexistence of two “LongStraw” usages underscores a persistent theme in long-context research: effective long-context behavior is jointly limited by statistical generalization, positional or architectural representation, and systems-level training executability. Long-short alignment addresses the first of these by regularizing output consistency across lengths (Du et al., 13 Jun 2025). DPE addresses the second by dimension-wise RoPE manipulation at inference time (Lu et al., 26 Apr 2025). The LongStraw execution stack addresses the third by making grouped RL response backward feasible beyond 2M tokens under fixed GPU budgets (Zhou et al., 16 Jul 2026).

Topic to Video (Beta)

No one has generated a video about this topic yet.

Whiteboard

No one has generated a whiteboard explanation for this topic yet.

Follow Topic

Get notified by email when new papers are published related to LongStraw.