Agentic reinforcement learning (RL) generates irregular rollout trees, but trajectory-by-trajectory training recomputes their shared interaction histories. Existing training-time prefix-reuse systems do not jointly address state recovery, dense execution, and distributed training for hybrid-attention models. We present HARTS (Hybrid-Attention RL over Tree Structures), to our knowledge the first system to demonstrate training acceleration from arbitrary rollout-tree prefix sharing on real hybrid-attention models. A two-level planner constructs compression-oriented microbatches, then jointly refines their membership, cross-microbatch sharing, placement, and execution order. Within each sharing window, gradients from dependent microbatches accumulate before one backward pass through the shared prefix computation. Differentiable boundary-state handoff and bounded replay preserve the per-trajectory chunkwise numerical path, while a linear-time planner attains the global minimum number of sequential linear-attention calls under our packed chunkwise-call model. The runtime uses a common compact layout across attention and dense or mixture-of-experts (MoE) layers, with support for distributed parallelism and activation recomputation. Recovering log-probabilities for each original trajectory position preserves RL objectives; for deterministic top-k routing without token dropping, occurrence counts restore MoE objective weights and expert-load statistics. On SWE-bench, HARTS delivers 4.77-10.45x the update throughput of trajectory-expanded training across two sampling schemes on 7.9B- and 124B-parameter hybrid-attention MoE models using 8-64 GPUs. On the 7.9B model, full-model logit differences are comparable to baseline self-rerun variation. The 124B model shows stable online learning trends through step 150 on SWE-bench and τ3-Bench.
Figures & tables
Figure 1. HARTS carries a rollout tree through distributed hybrid-attention training. (a) Three trajectories share token prefixes but may require different next-token labels. Boxes represent contiguous token ranges. (b) Full attention reads ancestor keys and values; linear attention consumes a recurrent state and produces a branch-specific next state. (c) L1 constructs compression-oriented microbatches; L2 refines membership, sharing, placement, and order. A slot aligns consumers across data-parallel replicas: MB0/MB1 form Slot 1 and MB2/MB3 form Slot 2. Producers P0/P1 compute shared prefixes for the consumers on their respective replicas. Bars show producer and per-slot residual work; stacked bars show current-step totals. All bar lengths are schematic. The call and layer examples use MB0’s complete tree in (a): Call 1 executes [A,B,D] , and Call 2 batches independent [C] and [E] . A common compact ordering supports both attention types and MoE execution. Within MB0, one logit row zB supplies T1/T2’s distinct target log-probabilities; occurrence count qB=2 restores MoE objective weights and expert-load counts without repeating routing or expert computation.
Figure 2. From tree forks to minimum-call linear-attention execution. (a) Boxes are fork-free segments; thick edges select direct continuations v1 and v3 . (b) At p2 , the state SB produced by the current call can initialize the independent side branch only after that call returns. The three boundary cases distinguish replay length from earliest legal start; this local view omits v2 . (c) Round 1 packs v0→v1→v3 with replay-started v2 ; Round 2 executes replay-started v4 . All sequences in each round form one packed call.
Figure 3. Overlapping sharing windows under consumer 1F1B. A contains consumers 1–2; B contains 3–4. PF/PB denote producer forward/backward; Fi/Bi denote consumer forward/backward. PB runs inside the window’s last consumer backward. Thin bars show stage-local retention from PF completion to release after PB; only stage 0 retains both windows concurrently. Dashed arrows show selected pipeline dependencies, not global barriers; event columns are not measured durations.
Property
Tiny
Flash
Total / active parameters (B)
7.9 / 1.3
124 / 5.1
Transformer layers
24
42
KDA:MLA per group
3:1
5:1
Dense FFN layers
1
1–2
MoE layers
2–24
3–42
Routed experts / top- k
128 / 8
512 / 8
Table 1. Evaluated hybrid-attention MoE models. Each repeated attention group contains the listed KDA:MLA ratio.
Figure 4. SWE-bench update throughput. Comparisons use matched inputs; speedups are relative to Baseline. Tiny additionally includes Node-wise LA.
Model / GPUs
Sampling
Compression
Tiny / 8
Tree-Step
3.91×
Flash / 32
Step-wise
7.58×
Flash / 32
Tree-Step
6.65×
Flash / 64
Step-wise
7.57×
Flash / 64
Tree-Step
6.56×
Table 3. Within-microbatch prefix compression on SWE-bench.
Planner
Compression
Balance
Time (s/step)
OR-Tools+FFD
5.535×
94.11%
3.1705
Harts L1
5.588×
98.62%
0.5097
Harts L1+L2
7.543×
97.59%
0.8891
Table 4. Two-level scheduling.
DP
Harts L1+L2 / OR-Tools
Microbatch slots
Balance P5 (%)
Step-total minimum (%)
2
599/614
96.65/87.91
98.40/96.35
4
301/319
94.40/79.62
98.20/95.08
8
152/166
93.61/72.91
95.96/91.45
Table 5. DP load-balance distributions.
Figure 5. Replay preserves chunkwise numerical fidelity. Solid lines: (a) symmetric KL and (b) mean absolute target-NLL difference, each relative to Baseline A. Orange dashed lines: No-replay’s excess over Baseline B’s error, (EN/EB−1)×100% . Token counts appear below handoff depths 0–5; lower is better.
PP schedule
Throughput (k tokens/s)
Non-overlapping windows
167.7
Overlap + backward-first
190.4
Harts (consumer 1F1B)
208.5
Table 6. PP scheduling ablation on Tiny/8 Tree-Step.
Parallel configuration
Baseline self-rerun
Harts w/o recomputation
Harts w/ recomputation
Absolute loss difference
DP8/TP1/PP1/EP8
0.99978662
0.99977463
0.99977816
9.66×10−5
DP4/TP2/SP/PP1/EP8
0.99975827
0.99974703
0.99974416
6.02×10−6
DP2/TP2/SP/PP2/EP4
0.99972408
0.99971610
0.99971732
1.33×10−4
Table 7. Full-model fidelity on Tiny. Cosine is the equally weighted mean of token-wise full-vocabulary logit cosines. Self-rerun compares baseline repeats; both Harts paths use the matched baseline. Absolute loss difference uses Harts without recomputation.
Figure 6. Online Flash training through step 150. Left pair: SWE-bench; right pair: τ3 -Bench. Each pair shows (a) reward and (b) the absolute difference between rollout/old and training-recomputed log-probabilities within each run.
Appendix figures & tables4 assets
Supplementary material from the paper’s appendix.
Appendix
Figure C.1. One compact physical ordering across hybrid-attention training. (a) Shared path P0–P1 branches into A, B, and C; A directly continues in Round 1, while B/C form independent sequences in Round 2. (b) Both side sequences start from the boundary state H0 after P0 and replay P1 inside the KDA core, without writing semantic outputs. (c) B’s MLA queries read ancestor KV P0/P1. Within B, diagonal blocks are causal; B1 reads earlier B0, not conversely. Outputs return to compact rows. (d) Compact routing/expert execution uses multiplicity-weighted router losses and expert-load counts; mapped log-probabilities restore original token order without expanded vocabulary-sized logits.
Figure D.1. Case I: three calls and R=373 . Circles show per-segment token counts. Thick solid edges form direct-continuation chains, while dashed edges start independent sequences and are labeled with their replay length under chunk size B=64 .
Figure D.2. Case II: two calls and R=813 . The three child subtrees of the single shared root R2 are unfolded vertically only to preserve label readability. Under chunk size B=64 , replay labels expose how many branches become ready together for the second variable-length call.
Figure D.3. Case III: two calls and R=780 . With chunk size B=64 , the first call follows the long G continuation chain while also executing three replay-started sequences; all side leaves become ready for one parallel second call.
Multi-step agentic reinforcement learning benefits from fine-grained credit assignment, yet existing approaches offer limited options: critic-free methods like GRPO assign a uniform advantage to every action in a trajectory, while learned value networks introduce notable overhead and can be fragile under sparse rewards. We observe that group rollouts targeting the same problem often traverse overlapping intermediate states, implicitly forming a tree whose branches diverge at successive decision points. Building on this insight, we introduce Rollout-Tree Monte Carlo (RTMC) advantage estimation, which aggregates return statistics across rollouts sharing a common state to produce per-step Q-values and advantages--without any learned critic. A state-action signature system compresses raw interaction histories into compact, comparable representations, making cross-rollout state matching tractable. On SWE-bench Verified, RTMC improves pass@1 by 3.2 percentage points over GRPO.
Agentic Reinforcement Learning (RL) enables LLMs to solve complex tasks by alternating between a data-collection rollout phase and a policy training phase. During rollout, the agent generates trajectories, i.e., multi-step interactions between LLMs and external tools. Yet, frequent tool calls induce long-tailed trajectory generation that bottlenecks rollouts. This stems from step-centric designs that ignore trajectory context, triggering three system problems for long-tail trajectory generation: queueing delays, interference overhead, and inflated per-token time. We propose Heddle, a trajectory-centric system to optimize the when, where, and how of agentic rollout execution. Heddle integrates three core mechanisms: trajectory-level scheduling using runtime prediction and progressive priority to minimize cumulative queueing; trajectory-aware placement via presorted dynamic programming and opportunistic migration during idle tool call intervals to minimize interference; and trajectory-adaptive resource manager that dynamically tunes model parallelism to accelerate the per-token time of long-tail trajectories while maintaining high throughput for short trajectories. Evaluations across diverse agentic RL workloads demonstrate that Heddle effectively neutralizes the long-tail bottleneck, achieving up to 2.5× higher end-to-end rollout throughput compared to state-of-the-art baselines.
Zili Zhang, Yinmin Zhong, Chengxu Yang +5
School of Computer Science, Peking University · Independent Researcher
Reinforcement learning with verifiable rewards (RLVR) is a promising approach for enhancing reasoning and agentic behavior in large language models. However, rollout-intensive policy optimization is often limited by insufficient reward contrast, arising when overly simple or complex prompts generate low-variance feedback and when outcome-only rewards assign the same terminal assessment to every decision in a multi-turn rollout. Past efforts have focused on allocating available rollout resources to promising prompts, yet they only leverage sample informativeness at the prompt level and neglect variation in prefix-level informativeness across turns within the same rollout. This work targets multi-turn agentic RL by modeling each ReAct-style thought-action-observation turn as a semantically distinct node, allowing budget allocation to extend from prompt roots to turn-level prefixes with further continuations, which naturally forms tree-structured rollouts. We introduce Tree Rollout Allocation for Contrastive Exploration (TRACE), a unified rollout allocation framework that enhances reward contrast within a fixed sampling budget. Technically, TRACE allocates rollout budget to both prompt roots and intermediate prefixes that are most likely to yield mixed terminal rewards. A shared generalizable predictor estimates conditional success probability at these anchors from prefix histories to guide this allocation. The resulting adaptive tree structure enriches outcome-only feedback and amplifies the policy-update signal. Empirically, TRACE achieves competitive performance and efficiency gains on typical agentic benchmarks, e.g., improving Qwen3-14B Multi-Hop QA average accuracy by 2.8 points over competitive baselines at equal sampling cost.