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.