HARTS: Efficient Agentic Reinforcement Learning for Hybrid-Attention Models over Arbitrary Rollout Trees

arXiv:2608.28158v1 [cs.LG] 28 Aug 2026

HARTS: Efficient Agentic Reinforcement Learning for Hybrid-Attention Models over Arbitrary Rollout Trees Boyuan Meng

Peihua Bao

Hong Liu

Ant Group China

Ant Group China

Ant Group China

Xiaowei Zhu

Chao Wang

Gen Li

Ant Group China

Ant Group China

Ant Group China

Zhenxuan Pan Ant Group China

Abstract

variation, and its reward trend is similar to the baseline over the first 120 steps of 𝜏 3 -Bench training.

Agentic reinforcement learning (RL) often produces irregular rollout trees whose trajectories share long interaction histories. Training each root-to-leaf trajectory independently recomputes these shared prefixes. Existing tree-structured training systems primarily target full-attention models and do not provide dense, differentiable hybrid-attention execution compatible with activation recomputation. We present HARTS (Hybrid-Attention RL over Tree Structures). HARTS jointly constructs microbatches, data-parallel (DP) replica assignments, and microbatch-slot schedules using the number of unique, non-replay compact token rows after prefix compression as its work metric. For chunkwise linear attention, a linear-time algorithm coordinates chunk-boundary state recovery and replay and produces the minimum possible number of sequential linear-attention calls under our packed chunkwise execution model. The resulting execution preserves the chunkwise state partitioning of conventional trajectory-wise training: it does not repeat projections, MLP/MoE computation, or final outputs, and performs only the bounded linear-attention state replay required for numerical alignment. At runtime, HARTS batches all branches in an execution round into one packed call, propagates gradients through differentiable state handoffs, supports activation recomputation, and restores per-token log-probabilities from compact outputs. For deterministic, no-token-drop top-𝑘 MoE routing, semantic multiplicities also restore the token weights and load statistics used by the MoE objectives. Existing RL objectives retain their interface. To our knowledge, HARTS is the first system to demonstrate arbitrary-rollout-tree prefix-sharing speedups on a real hybrid-attention model. On an Agentic RL workload generated from SWE-bench tasks, HARTS achieves 4.81– 4.87× forward/backward/gradient speedup with activation recomputation across multiple parallel configurations. Its numerical differences are comparable to baseline self-rerun

Keywords: agentic reinforcement learning, hybrid attention, rollout trees, prefix sharing, distributed training

1

Introduction

Agentic reinforcement learning (Agentic RL) trains an agent through multi-turn interactions, tool calls, parallel exploration, and retries after failed attempts. Sampling therefore need not produce independent linear sequences. It can instead produce a rollout tree with irregular depth, fan-out, and fork positions: each root-to-leaf path is a training trajectory, and multiple trajectories share the task description, environment state, and early interaction history (Figure 1(a)). Conventional trajectory-wise training expands and executes every trajectory independently, repeating the forward and backward computation of shared prefixes within the same training step. In our Agentic RL workload, generated by a Claude Code scaffold solving SWE-bench tasks [7], trajectory-wise expansion contains approximately 5.63× as many non-replay token rows as compact execution. Shared-prefix redundancy is therefore a major training cost for this workload. On a hybrid-attention model, a tree mask addresses only ancestor visibility in full attention; it does not satisfy the state-recovery and execution constraints of chunkwise linear attention (Figure 1(b)). In a full-attention layer, the tree primarily determines which ancestor keys and values each query can see. High-performance training algorithms for linear-attention variants such as RetNet, GLA, DeltaNet, Gated DeltaNet, and Kimi Delta Attention (KDA), however, use a chunkwise parallel/recurrent formulation. They compute within fixed-size chunks in parallel while propagating recurrent state across chunks [8, 15, 18–20]. An arbitrary fork rarely coincides with a fixed chunk boundary at which the operator can export a recurrent state for another branch. To preserve the state partitioning and floating-point path of

Contact: mengboyuan.mby@antgroup.com. 1

Boyuan Meng, Peihua Bao, Hong Liu, Xiaowei Zhu, Chao Wang, Gen Li, and Zhenxuan Pan

(a) One rollout tree and its trajectories Compact rollout tree

(b) Tree dependencies in full and linear attention

Root-to-leaf trajectories

A

T1

A

B

D

T2

A

B

E

T3

A

Full attention

QKV_D

Ancestor KV + current QKV output

→

ancestor KV [A, B] QKV_E

B

Full Attn D

Y_D

Full Attn E

Y_E

Linear Attn D

Y_D | S_D

Linear Attn E

Y_E | S_E

C T3

D

E

T1

T2

Linear attention

C

QKV_D

Recurrent state + current QKV output + next state

→

T1 and T2 share A and B, then use different next-token labels D and E.

shared state S_B QKV_E

Each box denotes one or more contiguous tokens.

(c) HARTS training pipeline for arbitrary rollout trees Placement is measured in physical compact tokens; execution follows the verified plan for the highlighted microbatch. 1 Balance DP-replica work: per-slot and current-step totals Bar length is schematic; microbatch 0* contains tree (a). Slot 1

Slot 2

DP replica 0

microbatch 0*

microbatch 2

DP replica 1

microbatch 1

microbatch 3

Balanced within each slot

Balanced within each slot

Balanced current-step totals

2 Plan linear-attention execution

3 Execute hybrid layers

Verified minimum: 2 sequential calls

on the linear-attention-aware ordering

4 Restore RL/MoE semantics One compact row, multiple semantic uses RL log-probabilities

Call 1

Planner-produced call layout

A

B

state dependency

D

Call 1: [A, B, D]

compact logit row zB

Call 2: [C] [E]

computed once T2: label E logsoftmax(zB)[E]

state dependency

MoE router training

Call 2: one variable-length batch C

T1: label D logsoftmax(zB)[D]

Full attention ancestor view

E

replica 0

Linear attention planned calls

MoE routing + experts

router result rB computed once

semantic multiplicity qB = 2

weights z-/aux losses and expert-bias counts

replica 1

Figure 1. HARTS carries one rollout tree through distributed hybrid-attention training. (a) The three rows are exactly the root-to-leaf trajectories represented by the tree; each box denotes one or more contiguous tokens. T1 and T2 share A and B but use different next-token labels D and E, while T3 shares only A before following C. (b) Both branch paths consume their current QKV. Full attention additionally reads the explicit ancestor KV of [𝐴, 𝐵], whereas linear attention consumes the shared recurrent state 𝑆 𝐵 and produces a branch-specific next state. (c) HARTS balances compact-token work across DP replicas at two levels: within each microbatch slot and in their current-step totals; bar lengths are schematic and do not report measured values. For the highlighted microbatch, the verified minimum plan executes [𝐴, 𝐵, 𝐷] in Call 1 and batches [𝐶] and [𝐸] as two independent sequences in Call 2. The resulting compact ordering is shared across layer types: full-attention layers recover ancestor visibility, chunkwise linear-attention layers execute the planned variable-length calls, and MoE layers perform compact routing and expert computation. One compact logit row 𝑧𝐵 yields the target-indexed log-probabilities for T1 and T2; the semantic multiplicity 𝑞𝐵 = 2 weights the router losses and the expert-load counts used by expert-bias update, without repeating routing or expert computation. trajectory-wise training, a side branch must resume from the nearest preceding chunk-boundary state and replay less than one chunk of ancestor state updates inside the linearattention core. Under the packed chunkwise operator interface that we target, a boundary state produced within a call cannot initialize a separate input sequence until that call returns, and gradients from all consuming branches must propagate through the shared state. Executing one rollout-tree node at a time, as in Tree Training’s linear-attention/GDN path [16], can traverse an arbitrary tree but fragments the work into many small sequential calls. When a fork falls inside a chunk, node boundaries also change the numerical path of conventional chunkwise training. Prefix sharing

on hybrid-attention models is thus simultaneously a staterecovery, call-planning, differentiability, and computationdensity problem. Existing tree-structured training systems do not solve this combination of problems. Tree Training [16] uses tree-node state routing in its linear-attention/GDN path. Its paper reports no real hybrid-attention training experiment that establishes the path’s efficiency or scalability. Its reported training performance, memory usage, and learning results use fullattention Qwen3-30B/32B models [17], and the system lacks effective support for parallel training [21]. Its strict zeroredundancy partitioning also requires retaining and reusing parent computation graphs and boundary states across partitions or microbatches. This lifetime conflicts with standard 2

HARTS: Hybrid-Attention RL over Tree Structures

This paper makes three contributions:

activation recomputation, which discards forward activations and rematerializes them during backward. Retaining the parent graph forfeits the corresponding memory savings; discarding it requires recomputing shared ancestors and is no longer strictly zero-redundancy. AReaL-DTA [21], in turn, targets standard full-attention Transformers, dynamically executes at tree-node granularity, and supports only data parallelism. It does not provide dense execution for models that interleave full- and linear-attention layers. Existing approaches therefore do not jointly support efficient execution on real hybrid-attention models, multidimensional parallelism, activation recomputation, and arbitrary rollout trees. We present HARTS (Hybrid-Attention RL over Tree Structures), a training system for Agentic RL over arbitrary rollout trees on hybrid-attention models (Figure 1(c)). HARTS carries prefix sharing through batch construction, full- and linear-attention execution, and dense or mixture-of-experts (MoE) feed-forward computation while reconstructing the original RL objective interface and the applicable MoE objective weights. It organizes and assigns training examples by the computation that remains after shared prefixes are removed. It then reconciles the distinct dependencies of full and chunkwise linear attention and coalesces irregular tree computation into a small number of packed calls. The same compact layout drives MoE routing and expert computation, so prefix sharing spans the entire hybrid model. This execution remains differentiable, supports activation recomputation, and composes with multiple forms of parallelism. At the objective interface, HARTS recovers a log-probability for every semantic token occurrence. Under deterministic, no-token-drop top-𝑘 routing, semantic multiplicity restores the token weights of MoE router losses and the expert-load counts used by expert-bias updates. Existing RL objectives remain unchanged. To our knowledge, HARTS is the first system to demonstrate arbitrary-rollout-tree prefix-sharing speedups on a real hybrid-attention model. We evaluate HARTS on Ling-3.0-tiny, a hybrid-attention MoE model with Multi-head Latent Attention (MLA) and KDA layers. The Agentic RL rollout-tree workload is generated by a Claude Code scaffold solving SWE-bench tasks. All performance experiments enable activation recomputation. Across multiple parallel configurations, 5.62–5.63× non-replay compact-row compression yields 4.81–4.87× forward/backward/gradient (F/B/Grad) speedup and 4.39–4.63× training-core speedup. F/B/Grad covers the complete measured training path after planning except for the optimizer update, whereas Core additionally includes HARTS planner latency. In separate numerical-fidelity experiments, with and without activation recomputation, the mean token-wise cosine similarity between HARTS and baseline full-vocabulary logits exceeds 0.9997 and is close to baseline self-rerun variation. Over the first 120 steps of online 𝜏 3 -Bench training [1], HARTS and the baseline also exhibit similar reward trends.

• Prefix-aware microbatch planning and schedule construction for arbitrary rollout trees. Subject to capacity constraints, HARTS generates candidate microbatch partitions together with their DP-replica assignments and microbatch-slot schedules. It measures cost by non-replay compact-token work and selects a balanced plan using total work, critical-path work in each slot, and cumulative work on each DP replica. • Minimum-call execution. When a fork falls inside a linear-attention chunk, each side branch resumes from the nearest preceding chunk-boundary state and replays less than one chunk of ancestor updates, preserving the chunkwise state partitioning of trajectorywise training. A linear-time algorithm selects one direct continuation at each fork and schedules every side branch in the earliest call where its initial state is available. Under the stated packed chunkwise execution model, the plan uses the minimum possible number of sequential linear-attention calls. It never repeats projections, MLP/MoE computation, or final outputs, and performs only the bounded state replay required for numerical alignment. At runtime, HARTS batches the branch sequences in each round into one packed call, selectively returns chunk-boundary states, and supports backward and full-layer activation recomputation through differentiable state handoffs. The resulting training path supports multiple parallel configurations. • RL and MoE semantics in compact coordinates, without expansion. Compact-to-semantic log-probability computation obtains current logprobabilities in the original token order directly from compact logits and accumulates gradients from reused compact logit rows back into compact coordinates. Existing policy objectives consume these log-probabilities and their original metadata unchanged. For deterministic, no-token-drop top-𝑘 routing, semantic multiplicity restores the token weights of router z-loss and the routing mass and assignment counts used by auxiliary losses and expert-bias updates. Routing, dispatch, and expert computation remain compact.

2

Background and Related Work

2.1

Rollout Trees in Agentic RL

An Agentic RL rollout need not be a single linear trajectory. For example, when a Claude Code scaffold attempts a SWE-bench task, several explorations can share the issue description, repository state, and early agent–environment interactions before trying different tool calls, code changes, 3

Boyuan Meng, Peihua Bao, Hong Liu, Xiaowei Zhu, Chao Wang, Gen Li, and Zhenxuan Pan

or recovery paths. A single task instance can therefore produce a rollout tree with irregular depth and fan-out, where every root-to-leaf path is a training trajectory. After conceptually expanding all trajectories, let 𝑠 = (𝑥, 𝑖) denote the 𝑖th token occurrence in trajectory 𝑥. We call 𝑠 a semantic position. Let S be the set of all semantic positions and 𝑆 = |S|. Prefix sharing executes the projection, MLP/MoE, and final output for a shared token only once. Let C denote the set of compact tokens that carry these unique outputs, and let 𝐶 = |C|. A compact token row is one row of a compact tensor that stores the hidden state, and later the logits, of one such token. Throughout the paper, compact-token work counts these unique, non-replay rows. It excludes bounded linear-attention replay, temporary MLA ancestor-KV view entries, and backend padding or alignment rows; these additional operator costs are reflected in measured execution time. The runtime assigns every compact token to a persistent, hardware-aligned physical ordering; its tensor location is a physical row. The backend may insert padding or alignment rows around these physical token rows. Such padding is part of the physical layout but not compact-token work. Temporary KDA replay/context rows and MLA ancestor-KV entries are operator-specific views rather than additional compact tokens. Sharing physical computation does not merge training targets. Two semantic positions from different trajectories can map to the same compact token row when they lie on a shared prefix, and thus reuse that row’s hidden state and logits. They may nevertheless have different next-token labels, advantages, loss masks, old or reference log-probabilities, or loss weights. Prefix sharing eliminates duplicate sharedprefix computation represented by compact rows, but preserves every semantic position and its training target.

2.2

𝐶:   𝐵 S [𝑐+1] = Diag 𝛾 [𝑐 ] S [𝑐 ]  ⊤
𝑖→𝐵 + Γ[𝑐 U [𝑐 ] − W [𝑐 ] S [𝑐 ] , ] ⊙ K [𝑐 ]

(1)

Here, S [𝑐 ] is the recurrent state at the beginning of chunk 𝑐, while K [𝑐 ] , U [𝑐 ] , W [𝑐 ] , and 𝚪 [𝑐 ] are the chunk’s dense updates and decay summaries. The exact recurrence can differ across linear-attention backends. HARTS assumes only that the training operator advances boundary states continuously across one or more execution chunks. We call this internal operator process chunkwise state propagation. A conventional training interface usually returns layer outputs and, optionally, a final state. It does not expose the internal boundary states that connect adjacent chunks for consumption by another tree branch. Define a fork position 𝑝 as the number of tokens in the shared prefix—the exclusive boundary immediately before a branch’s first new token. When 𝑝 lies inside a chunk, a side branch must start from the nearest preceding fixed chunk boundary j𝑝 k (2) 𝑎(𝑝) = 𝐵 𝐵 and concatenate the fewer-than-one-chunk ancestor interval [𝑎(𝑝), 𝑝) with the branch’s new tokens in the same chunkwise call. This construction preserves the state partitioning and floating-point path of conventional training on an independent trajectory. We call this fixed re-execution of ancestor state updates replay. Implementing this semantics requires selectively exposing legal boundary states from the training operator, passing them differentiably between branches, and organizing the fixed replay into dense execution. 2.3

Tree-Structured Training Systems

Tree Training. Tree Training [16] demonstrates the performance potential of arbitrary rollout-tree packing on fullattention Qwen3 models, but reports no training experiment on a real hybrid-attention model. Its full-attention path can organize a partition as a relatively dense computation using depth-first-search (DFS) serialization and a tree mask. Its GDN/linear-attention path, however, routes state and executes computation at tree-node granularity. Short, irregular node segments can therefore create many small sequential calls, and arbitrary node boundaries repartition the execution chunks used in standard trajectory training. Speedups on full-attention Qwen3 do not establish the performance of this hybrid-attention path. AReaL-DTA further reports that Tree Training’s static prefix packing and custom kernels lack effective support for parallel training [21]. Tree Training’s strict zero-redundancy partitioning also conflicts structurally with standard activation recomputation [2]. The former must retain parent computation graphs and boundary states across partitions or microbatches so that descendants can reuse them. The latter saves memory

Hybrid-Attention Training

We use hybrid-attention model to mean a model that interleaves full- and linear-attention layers in any order. This definition concerns only the attention layers; it permits either dense or MoE feed-forward networks (FFNs). For full attention, the tree primarily determines which ancestor keys and values are visible to each query, which an ancestor-aware mask or an equivalent key/value (KV) view can express. Linear attention additionally carries recurrent state along the token order. High-performance training kernels do not invoke this recurrence token by token. Instead, they divide a sequence into execution chunks of length 𝐵 and advance several chunks within one kernel using parallel or block scans. KDA is one example of this training path [8]. For illustration, we write KDA’s chunk-level state update as follows, renaming the chunk length in its original formulation to 𝐵 to avoid a conflict with the compact-token count 4

HARTS: Hybrid-Attention RL over Tree Structures

by discarding forward activations and rematerializing them during backward. Under this partition lifecycle, retaining the parent graph loses these memory savings; discarding it requires re-executing shared ancestors and is no longer strictly zero-redundancy. Activation recomputation is widely used to fit long contexts and large models, so this conflict limits the method’s practicality. Tree Training’s unique-token loss also cannot preserve the per-semantic-position training targets required by general RL objectives, because shared positions can have different labels and RL metadata.

communication also affect elapsed time. The number and membership of microbatches, their DP-replica assignments, and their microbatch-slot order must be planned jointly. Dense linear-attention execution. Even when the projection, MLP/MoE computation, and final output of every unique token execute only once, the linear-attention core must perform the bounded state replay required for numerical alignment. Calling that core once per tree node, as in Tree Training’s linear-attention/GDN path [16], would add many small kernels. Under the packed chunkwise-call interface that HARTS targets, an independent branch can start only after its initial chunk-boundary state is available; a state produced within a call cannot simultaneously initialize another independent input sequence in that call. The system must therefore minimize the sequential call depth induced by intercall state dependencies and coalesce all ready sequences in a round into dense packed execution. Shared execution must remain differentiable and recomputable. In full-attention layers, several descendants can read the same ancestor KV, and the standard gather/scatter backward path must accumulate their gradients into the shared compact activation. In linear-attention layers, each successor’s initial-state gradient must flow back into the earlier computation that produced the boundary state. Activation recomputation must also rematerialize the same compact layout and state dependencies without retaining a forward graph across microbatches. Compact computation must preserve RL and MoE semantics. Semantic positions mapped to the same compact logit row can have different next-token labels. Consequently, weighting a single shared loss by a coefficient, as in Tree Training’s unique-token loss, cannot recover their individual current log-probabilities [16]. Restoration must occur during log-probability computation so that policy objectives depending only on current log-probabilities and per-token metadata remain reusable. MoE router z-loss, load-balancing auxiliary loss, and expert-bias updates must likewise preserve the token weights of unshared training without duplicating expert execution.

AReaL-DTA. AReaL-DTA [21] targets standard fullattention Transformers and is neither designed for nor evaluated on hybrid-attention models. Its dynamic traversal at tree-node granularity reduces memory usage primarily by controlling activation residency, but splits model computation into many sequential calls and supports only data parallelism. Its tree-workload distribution and node-wise execution therefore do not provide dense execution for models that interleave full- and linear-attention layers. Prefix reuse and efficient kernels. Sequence packing [9] improves training utilization by removing padding, but does not reuse model computation for identical prefixes. Inference-time prefix caching reuses KV states, but does not address training backward or per-semantic-position RL objectives. FlashAttention [3] and chunkwise linearattention kernels optimize an individual operator call over a linear sequence or a conventional packed batch; they do not determine branch-state recovery, differentiable state transfer, or inter-call dependencies on an arbitrary rollout tree. HARTS plans shared tree execution above these kernels and extends the linear-attention training interface to return selected chunk-boundary states and propagate gradients through those states during backward. 2.4

Design Challenges

Prefix sharing and distributed load must be planned together. We use DP replica to mean the logical model copy that receives one local microbatch in a data-parallel schedule. With tensor parallelism (TP), sequence parallelism (SP), pipeline parallelism (PP), or expert parallelism (EP), the ranks in the corresponding parallel groups jointly execute a microbatch on one DP replica; these are logical groupings, and their rank sets can overlap. Partitioning work by raw trajectory-token counts or fixed DFS boundaries duplicates otherwise shared prefixes at microbatch boundaries. Each DP replica must execute the same number of nonempty local microbatches to preserve the call order of model operations and collectives. The 𝑗th local microbatch on every DP replica collectively forms microbatch slot 𝑗. With EP, token-dispatch and combine collectives introduce tight synchronization across DP replicas. The compact-work critical path of a slot is therefore determined by its heaviest microbatch, although replay, temporary attention views, and

3

Design: HARTS

3.1

Overview

HARTS accepts the same inputs as conventional training: complete trajectories from arbitrary rollout trees, next-token labels, and the advantages, rewards, masks, old or reference log-probabilities, loss weights, and trajectory boundaries required by the RL algorithm. From ingestion onward, the system maintains a mapping 𝑚 : S → C,

(3)

where 𝑚(𝑠) is the compact token row corresponding to semantic position 𝑠. This row, rather than a trajectory, stores the hidden state and later the logits of one unique compact 5

Boyuan Meng, Peihua Bao, Hong Liu, Xiaowei Zhu, Chao Wang, Gen Li, and Zhenxuan Pan

compact work, so 𝑤 (𝑀) = 𝐶 (𝑀). After stably sorting the trajectories in 𝑀 lexicographically as 𝑥 1, . . . , 𝑥𝑘 , this work is

token in a compact tensor. Only semantic positions with the same actual token prefix map to the same compact token row. For load balancing, trajectories from different rollout trees that share no actual prefix may still occupy the same microbatch and the same packed model call, but their computation follows separate paths. Figure 1(c) shows four consecutive stages. First, the prefixaware planner jointly chooses microbatch membership, DPreplica assignments, and a microbatch-slot schedule (Section 3.2). Second, the minimum-call planner converts each rollout tree into linear-attention input sequences organized by execution round and coalesces independent sequences in the same round (Section 3.3). This plan also defines a compact physical layout suited to linear-attention execution. The hybrid-attention runtime then uses that layout for fulland linear-attention layers and dense or MoE FFNs, while implementing differentiable boundary-state handoffs and recomputation (Section 3.4). Finally, the system uses 𝑚 to compute current log-probabilities in the original batch order directly from compact logits and restores the original token weights required by MoE router losses and expert-bias updates (Section 3.5).

𝐶 (𝑀) = |𝑥 1 | +

𝑘 ∑︁

(|𝑥𝑖 | − LCP(𝑥𝑖 −1, 𝑥𝑖 )) .

(5)

𝑖=2

When 𝑥𝑖 is inserted in lexicographic order, its immediate predecessor provides its longest prefix shared with any trajectory already inserted. Let 𝐿𝑖 = LCP(𝑥𝑖 −1, 𝑥𝑖 ), and let 𝑊 (𝑎, 𝑏) be the compact work of the contiguous interval 𝑥𝑎 , . . . , 𝑥𝑏 . Then 𝑊 (𝑎, 𝑡 − 1) + 𝑊 (𝑡, 𝑏) − 𝑊 (𝑎, 𝑏) = 𝐿𝑡 . (6) Equation (6) shows that splitting a contiguous microbatch at 𝑡 increases compact work by exactly 𝐿𝑡 , the LCP length of the boundary trajectories 𝑥𝑡 −1 and 𝑥𝑡 . A low adjacentLCP valley therefore has a low prefix-duplication cost. For any two trajectories in lexicographic order, the minimum adjacent LCP between them determines their shared-prefix length; a range-minimum-query (RMQ) index retrieves this value directly. This preprocessing supports constant-time queries for interval compact work and pairwise LCP. The number of unique token nodes in the compact trie is used directly for capacity checks and candidate comparison.

3.2 Prefix-Aware Microbatch Planning and Schedule Construction

(4) The three components minimize total compact work, aggregate slot-critical compact work, and maximum cumulative compact work on any replica, in that order. The intended effects of the latter two are to reduce synchronization-critical work within each slot and to balance work over the current training step.

Candidate range and natural partition. A low LCP between adjacent trajectories indicates a shallow fork and defines an adjacent-LCP valley. The planner recursively splits only contiguous intervals whose compact-token work exceeds 𝑇 , first choosing the lowest valley. Ties in the minimum adjacent LCP are broken by minimizing the larger compact work of the two sides, and then by minimizing the difference between their compact work. Let 𝐾 denote the number of microbatches after partitioning. We call the first capacityfeasible partition the natural partition P0 , and denote its number of microbatches by 𝐾0 . If merging all trajectories into one compact trie requires work 𝐺 = 𝐶 (𝑋 ), then    𝐺 𝐾min = 𝐷 max 1, . (7) 𝑇 ·𝐷 A value of 𝐾 divisible by 𝐷 is DP-compatible. The base search range is 𝐾 ∈ {𝐾min, 𝐾min + 𝐷, . . . , 𝐾max },       (8) 𝐾0 𝑁 𝐾max = max 𝐾min, min 𝐷 ,𝐷 . 𝐷 𝐷 𝐾min is a necessary lower bound, but need not admit a feasible grouping. Moreover, the compact-first objective is not monotone in 𝐾: fewer microbatches can still separate trajectory groups with long shared prefixes. The planner therefore compares multiple candidates rather than selecting the smallest feasible 𝐾.

Exact prefix cost. Let LCP(𝑥, 𝑦) be the length of the longest common prefix of trajectories 𝑥 and 𝑦. Merging identical token-prefix nodes within a microbatch produces a compact trie; its number of unique token nodes is the

Exact search for small inputs. For each 𝐾 in Equation (8), exact-search mode finds the minimum-work partition whose microbatches are nonempty, capacity-feasible, and contiguous in lexicographic order. Let 𝐹 [𝑘, 𝑒] be the

Problem and objective. HARTS receives complete rootto-leaf trajectories drawn from arbitrary rollout trees, a data-parallel degree 𝐷, and a microbatch capacity 𝑇 . Let 𝑋 = {𝑥 1, . . . , 𝑥 𝑁 } be the indexed collection of trajectories. A trajectory is the smallest scheduling unit and cannot be split between tokens, while a microbatch may combine trajectories from different rollout trees. A valid plan provides exact coverage: each 𝑥𝑖 belongs to exactly one microbatch, and the plan assigns that microbatch to both a destination DP replica and a microbatch slot. Let P be a candidate set of microbatches, let 𝑀𝑑,𝑗 be the microbatch executed by replica 𝑑 in slot 𝑗, and let 𝑤 (𝑀) be its compact-token work. HARTS ranks fully instantiated candidate plans lexicographically by ! ∑︁ ∑︁ ∑︁ 𝐽 (P) = 𝑤 (𝑀), max 𝑤 (𝑀𝑑,𝑗 ), max 𝑤 (𝑀𝑑,𝑗 ) . 𝑀∈P

𝑑 𝑗

𝑑 𝑗

6

HARTS: Hybrid-Attention RL over Tree Structures

minimum compact-token work for partitioning the first 𝑒 trajectories into 𝑘 microbatches. Then 𝐹 [𝑘, 𝑒] =

min

{𝐹 [𝑘 − 1, 𝑠] + 𝑊 (𝑠 + 1, 𝑒)} .

For an existing microbatch, sort its members lexicographically by token sequence, and let 𝑝 and 𝑞 be the immediate predecessor and successor at the insertion position of trajectory 𝑥. Insertion replaces the adjacency represented by LCP(𝑝, 𝑞) with LCP(𝑝, 𝑥) and LCP(𝑥, 𝑞). The exact increase in compact work caused by 𝑥 is therefore

(9)

𝑠<𝑒 𝑊 (𝑠+1,𝑒 ) ≤𝑇

For state 𝐹 [𝑘, 𝑒], every split point 𝑠 defines a candidate partition: 𝑘 − 1 microbatches cover the first 𝑠 trajectories, and 𝑥𝑠+1, . . . , 𝑥𝑒 form the final microbatch. The dynamic program first compares total compact work—the sum of the preceding partition and the last interval in Equation (9). When several split points have equal total work, it chooses the partition with the smaller maximum single-microbatch work. This secondary criterion preserves compact-first semantics while avoiding an unnecessarily heavy microbatch when total work ties, giving subsequent replica and slot scheduling a better-balanced input. A split point 𝑠 is feasible only when the preceding state is feasible and the last interval satisfies 𝑊 (𝑠 + 1, 𝑒) ≤ 𝑇 . For fixed endpoint 𝑒, 𝑊 (𝑠 + 1, 𝑒) decreases monotonically as 𝑠 moves right. The planner can therefore skip every 𝑠 before the first capacity-feasible start without missing a valid contiguous partition. This mode is exact only over the contiguous-partition space and candidate 𝐾 range in Equation (8); it excludes microbatches whose members are noncontiguous in the global lexicographic order. Exact candidates remain in the final candidate set and are never replaced by later greedy candidates.

Δ𝐶 = |𝑥 | − LCP(𝑝, 𝑥) − LCP(𝑥, 𝑞) + LCP(𝑝, 𝑞).

(10)

At either end of the lexicographic order, Equation (10) omits the missing predecessor or successor term. To reduce sensitivity to processing order, the planner uses ascending token-lexicographic order, descending tokenlexicographic order, and descending trajectory length. The fourth order uses the descending unique-suffix score 𝑢 (𝑥) = |𝑥 | − max(LCP(𝑝𝑥 , 𝑥), LCP(𝑥, 𝑞𝑥 )), where 𝑝𝑥 and 𝑞𝑥 are the global lexicographic neighbors of 𝑥. This order prioritizes rare branches with fewer opportunities to share prefixes. For each of the four processing orders, the planner applies two placement policies. Among capacity-feasible destination microbatches, best-fit selects the one with the largest work after insertion to reduce capacity fragmentation, while max-shared selects the one with the largest shared amount |𝑥 | − Δ𝐶 to preserve prefix affinity. If no existing microbatch can accommodate 𝑥, the policy creates a new one. Applying both policies to all four orders yields eight plans. Each plan then attempts to eliminate small microbatches and aligns 𝐾 to 𝐷 using singleton splits with the smallest compactwork increase. The planner chooses one greedy candidate by minimizing microbatch count first and total compact work second. This preselection is a candidate-generation heuristic that injects a small-𝐾 alternative; it does not replace the compact-work-first final objective. The chosen greedy candidate remains in the final comparison under Equation (4), even if its 𝐾 is absent from the base candidate set.

Heuristic base search for large inputs. Let 𝑁 0 be the trajectory-count threshold between exact and scalable search; we set 𝑁 0 = 512. When 𝑁 > 𝑁 0 , the planner uses heuristic search with a sampled-𝐾 budget of 𝑄 max = 32. It retains every candidate 𝐾 when the interval fits within this budget; otherwise, it retains both endpoints, densely sampled small values starting at 𝐾min , and sparse values spanning the full interval. It then evolves the natural partition P0 . For a target 𝐾 < 𝐾0 , it repeatedly merges the capacity-feasible adjacent microbatches 𝐴 and 𝐵 with the largest compact-work saving, 𝐶 (𝐴) + 𝐶 (𝐵) − 𝐶 (𝐴 ∪ 𝐵). For a target 𝐾 > 𝐾0 , it repeatedly selects the contiguous split of microbatch 𝑀 into 𝐿 and 𝑅 with the smallest compact-work increase, 𝐶 (𝐿) + 𝐶 (𝑅) − 𝐶 (𝑀). Merge/split evolution stops when it reaches the target 𝐾; if no legal operation exists before then, construction of that candidate fails.

Adaptive candidate refinement. We call a postprocessing pass that changes a small number of memberships without changing 𝐾 local refinement. For 𝑁 ≤ 𝑁 0 , exact candidates remain unchanged and only the additional greedy candidate is refined. For 𝑁 > 𝑁 0 , candidates selected for refinement are anchors. All base plans remain in the comparison under Equation (4), but at most 𝐵 ref = 3 distinct anchors are refined: the candidate at the 𝐾 produced by greedy placement, another candidate with the lowest base compact work, and a remaining candidate whose 𝐾 is closest to the greedy value. For a lexicographically contiguous anchor, boundary refinement moves each neighboring cut within 𝐻 = 8 positions, first reducing the two affected microbatches’ total compact work and then their maximum individual work. It performs at most two sweeps and terminates early when a complete sweep changes no cut.

Scale-independent trie-aware greedy. Exact dynamic programming and sampled-𝐾 search construct base candidates over contiguous partitions and selected values of 𝐾, respectively, and thus do not cover arbitrary noncontiguous membership well. To expand this search space, the planner independently runs trie-aware greedy directly from the original trajectories for every input size, whether or not 𝑁 > 𝑁 0 . This path neither modifies an existing candidate nor fixes 𝐾 in advance. 7

Boyuan Meng, Peihua Bao, Hong Liu, Xiaowei Zhu, Chao Wang, Gen Li, and Zhenxuan Pan

Sparse repair then uses a small number of noncontiguous membership updates to further reduce compact work. Each iteration first considers moving one trajectory from a nonsingleton microbatch into a capacity-feasible destination that contains a trajectory within 𝐻 = 8 positions in the global lexicographic order. It accepts only the move with the largest strict decrease in the two affected microbatches’ total compact work. If no improving move exists, it considers swapping two trajectories whose global lexicographic positions differ by at most 𝐻 = 8 and likewise accepts only a strict improvement. Each iteration applies at most one best move or swap. Sparse repair performs at most four iterations and terminates early when neither operation improves the plan. Every update preserves 𝐾, exact trajectory coverage, nonempty microbatches, and the capacity constraint. After refinement, each candidate has a fixed 𝐾 and fixed microbatch membership. Subsequent stages only construct its DP-replica assignment and microbatch-slot schedule; they do not change the grouping.

preserves the fork topology of the rollout tree. Segment 𝑣 contains the tokens in the position interval [𝑠 𝑣 , 𝑒 𝑣 ). Rollout-tree forks determine segment boundaries, which generally differ from the fixed chunk boundaries of the linearattention backend. In Figure 2(a), the vertical lines at 0, 𝐵, 2𝐵 form the chunk grid, while 𝑝 1 and 𝑝 2 are fork positions. The planner later selects at most one direct continuation among the children of each segment. The thick 𝑣 0 → 𝑣 1 and 𝑣 1 → 𝑣 3 edges depict execution choices, not distinct edge types in the segment tree. One packed linear-attention call can process multiple input sequences with independent initial states, and each sequence can traverse multiple chunks within the call. Recurrent state produced within a sequence is immediately available to later tokens in that sequence. Under our packed chunkwise-call interface, however, a chunk-boundary state produced by the call can initialize another independent input sequence only after the call returns. At most one child of a fork can therefore follow the parent segment directly in the same input sequence; every other child starts a new sequence. In Figure 2(a), the planner chooses 𝑣 1 and 𝑣 3 as direct continuations at 𝑝 1 and 𝑝 2 , making 𝑣 0 → 𝑣 1 → 𝑣 3 one input sequence. Segments 𝑣 2 and 𝑣 4 start independent sequences, and Figure 2(c) schedules all three sequences across two rounds. We call each such packed call an execution round and minimize the number of sequential rounds required to complete a microbatch in each linear-attention layer. This objective captures the state-dependent sequential call depth under the targeted interface.

DP-replica assignment. For any candidate, the planner sorts its 𝐾 microbatches by work and groups each consecutive block of 𝐷 microbatches into one slot. Within a slot, it assigns heavier microbatches to replicas with less cumulative work so far. This produces an executable schedule and determines the slot-critical work and maximum replicacumulative work in Equation (4). Plan selection and guarantees. The planner compares all valid candidates lexicographically under Equation (4): it first minimizes total compact work and, when candidates tie on this primary objective, uses the two assignmentdependent objectives to select the final plan. It returns the best fully instantiated plan among the candidates it generates. Every returned plan guarantees exact coverage, nonempty microbatches within capacity, 𝐾 mod 𝐷 = 0, and exactly 𝐾/𝐷 microbatches per DP replica. These are feasibility guarantees. Membership search is exact only for the contiguous partitions and 𝐾 range stated above; sampled-𝐾 evolution, trie-aware greedy placement, and local refinement are heuristic. DP-replica assignment and slot construction are greedy in every search mode. 3.3

Fixed replay and earliest legal start. For segment 𝑣 = [𝑠 𝑣 , 𝑒 𝑣 ), let 𝑎 𝑣 = 𝑎(𝑠 𝑣 ) be the resume anchor at its entry and 𝑏 𝑣 = 𝑎(𝑒 𝑣 ) the resume anchor at its terminal position, which is a fork for an internal segment. By Equation (2), these are the nearest legal chunk boundaries no later than 𝑠 𝑣 and 𝑒 𝑣 , respectively. We call the corresponding S𝑎 𝑣 and S𝑏 𝑣 the anchor states. Every child except the direct continuation starts from S𝑏 𝑣 and, when 𝑏 𝑣 < 𝑒 𝑣 , replays [𝑏 𝑣 , 𝑒 𝑣 ) inside the linear-attention core. Replay positions produce no semantic outputs. Chunk size and fork position uniquely determine this replay interval. The planner schedules when the replay runs; replay length is not an optimization variable. The following boundary condition determines call depth. Suppose the state at 𝑎 𝑣 is available before the current round starts and

Minimum-Call Planning for Linear-Attention Execution

Execution model. We first consider one rooted rollout tree. A segment is a maximal contiguous token path with no internal fork. In Figure 2(a), 𝑣 0 forks into 𝑣 1 and 𝑣 2 at 𝑝 1 = 𝐵 − 1, and 𝑣 1 forks into 𝑣 3 and 𝑣 4 at 𝑝 2 = 𝐵 + 1. Contracting every segment into a vertex and connecting adjacent parent and child segments yields a rooted segment tree 𝐺 = (𝑉 , 𝐸). In this example, 𝑉 = {𝑣 0, . . . , 𝑣 4 }, 𝐸 = {(𝑣 0, 𝑣 1 ), (𝑣 0, 𝑣 2 ), (𝑣 1, 𝑣 3 ), (𝑣 1, 𝑣 4 )}, and 𝑣 0 is the root. The transformation contracts only fork-free token paths and

𝑒 𝑣 < 𝑎 𝑣 + 𝐵,

(11)

then 𝑏 𝑣 = 𝑎 𝑣 : the segment crosses no new fixed chunk boundary and produces no intermediate boundary state. A side branch can reuse the already available anchor state, replay through 𝑒 𝑣 , and start in the same round. If 𝑒 𝑣 ≥ 𝑎 𝑣 + 𝐵, the required S𝑏 𝑣 is produced only inside the current call, so the side branch cannot start before the next round. The inequality is strict. At 𝑒 𝑣 = 𝑎 𝑣 + 𝐵, replay length is zero, but the boundary 8

HARTS: Hybrid-Attention RL over Tree Structures

(a) Rollout tree

→ segments

0

B

(c) Optimal plan: Tpre(root) = 2

(b) Replay length and earliest side-branch call

boxes: no-fork token paths bold edges: selected continuations

Local view at p2 = B+1; v2 omitted

2B

v3 v1

S0

Round 1 · one variable-length linear-attention call

SB

Current call

v0

v1

sequence 0

p2

v3 (direct)

S0

v0

v1

v3 requested

SB becomes usable by a side branch only after the current call returns

v0 (root)

sequence 1 (starts in the same call)

Next call · starts v4

p2 = B+1

SB

v4

replay [B,p2)

S0

v4

replay [0,p1)

v2

SB after return

state dependency

Boundary cases (pre-ready anchor S0, a=0)

p1 = B−1

v2

p = B−1 replay = B−1 S0 pre-ready

p = B (key case) replay = 0 SB produced now

p = B+1 replay = 1 SB produced now

side: same call

side: next call

side: next call

Round 2 · one variable-length linear-attention call

SB

schematic chunk grid · not to scale

replay [B,p2)

v4

planner changes when fixed replay runs, not its volume direct continuation

replay

state

Figure 2. From rollout-tree forks to minimum-call linear-attention execution. (a) Each box is a segment with no internal fork. Diamonds mark two rollout-tree forks; thick edges show the planner choosing 𝑣 1 and 𝑣 3 as direct continuations of the current input sequence. (b) A local expansion of the second fork, 𝑝 2 : S𝐵 , produced by the current call, can initialize the independent side branch 𝑣 4 only after the call returns. The three boundary cases below distinguish replay length from the earliest call in which a side branch can start. This local view omits 𝑣 2 . (c) The complete plan executes 𝑣 0 → 𝑣 1 → 𝑣 3 together with replay-prefixed 𝑣 2 in Round 1, then replay-prefixed 𝑣 4 in Round 2. All input sequences in one round form a single packed linear-attention call. • 𝑇in (𝑣) is the minimum when 𝑣 has entered the current call as its parent’s direct continuation and its entry state propagates within the current input sequence.

state is still an output of the current call and cannot serve as another sequence’s initial state within that call. Figure 2(b) expands the second fork in Figure 2(a), at 𝑝 2 = 𝐵 + 1. The current call follows 𝑣 0 → 𝑣 1 → 𝑣 3 and produces S𝐵 , but this state cannot initialize the independent side branch until the call returns. Segment 𝑣 4 therefore starts from S𝐵 in the next call, replays [𝐵, 𝑝 2 ), and then processes its own tokens. This local view omits 𝑣 2 , the side branch of the first fork at 𝑝 1 ; Figure 2(c) shows the complete plan. The three 𝐵 − 1, 𝐵, and 𝐵 + 1 cases at the bottom of Figure 2(b) emphasize that fork position determines replay length, whereas state readiness determines the earliest call in which the side branch can start.

Figure 2 gives a complete example. Each leaf segment, 𝑣 2 , 𝑣 3 , or 𝑣 4 , needs one round. At 𝑝 2 = 𝐵 + 1, only one of 𝑣 3 and 𝑣 4 can directly continue 𝑣 1 . The figure keeps 𝑣 3 in Round 1; 𝑣 4 waits for S𝐵 and enters Round 2, so the subtree rooted at 𝑣 1 needs two rounds. At 𝑝 1 = 𝐵 − 1, S0 is available before Round 1. Side branch 𝑣 2 can therefore replay [0, 𝑝 1 ) and execute in the same round as 𝑣 0 → 𝑣 1 → 𝑣 3 . The entire tree requires two sequential calls. The recurrence below generalizes this choice to any segment tree. A single call can scan the complete input of a leaf segment even when it crosses multiple chunks. Thus,

Minimum-call recurrence under two state-readiness cases. At each fork, the planner must select one direct child and decide when the remaining side branches start. If the anchor state required by a side branch is available before the current call, that branch may execute in the same round as the direct child. If the call itself produces the state, the side branch must wait until the next round. Different side branches in one round can be coalesced into the same packed call. The planner minimizes call depth over all choices of direct child. To represent the two entry conditions, we define two minimum call counts for the complete subtree rooted at each segment 𝑣. Both counts include the current round in which 𝑣 executes:

𝑇pre (𝑣) = 𝑇in (𝑣) = 1,

ch(𝑣) = ∅.

(12)

For a non-leaf segment, let 𝑜 ∈ ch(𝑣) be the child that directly continues the current input sequence, and define the maximum of an empty set as zero. When the entry state arrives within the current call, the direct-continuation subtree does not wait, while every side branch is delayed by one round. Therefore,  𝑇in (𝑣) = min max 𝑇in (𝑜), 1 + 𝑜 ∈ch(𝑣)

• 𝑇pre (𝑣) is the minimum number of rounds when S𝑎 𝑣 is available before the current call starts; and

 max

𝑇pre (𝑢) . (13)

𝑢 ∈ch(𝑣)\{𝑜 }

When the entry anchor is available before the round, Equation (11) determines whether side branches can start in that 9

Boyuan Meng, Peihua Bao, Hong Liu, Xiaowei Zhu, Chao Wang, Gen Li, and Zhenxuan Pan

Algorithm 1 Minimum-call planning and plan recovery

round: 𝑇pre (𝑣) =

   max 𝑇pre (𝑢), 𝑒 𝑣 < 𝑎 𝑣 + 𝐵,  𝑢 ∈ch(𝑣)  𝑇in (𝑣), 𝑒 𝑣 ≥ 𝑎 𝑣 + 𝐵. 

Require: rooted segment tree 𝐺, chunk size 𝐵 Ensure: ordered rounds of packed sequences and requested states 1: for each segment 𝑣 in postorder do 2: if 𝑣 is a leaf then 3: 𝑇pre [𝑣] ← 𝑇in [𝑣] ← 1 4: else 5: obtain the largest and second-largest child 𝑇pre 6: for each child 𝑜 do 7: ℎ𝑜 ← max𝑢≠𝑜 𝑇pre [𝑢]

(14)

In the first case, 𝑣 and all of its children can start in the current round from the same ready anchor, so no additional wait is necessary. The second case reduces to the directchild choice and side-branch delay in Equation (13). The root segment’s zero state is available before Round 1, making the minimum number of calls for the tree 𝑇pre (root). For multiple independent rollout trees in one microbatch, the planner solves each tree separately and merges sequences with the same round number. Because every root starts from an independent zero state, the microbatch call count is the maximum across its trees. The first pass of Algorithm 1 computes the minimum call count. Maintaining the largest and second-largest child values of 𝑇pre yields all ℎ𝑜 values with a linear total number of child visits. Recovery then propagates the root’s round budget down the tree. Here, 𝐾side is the remaining budget for side branches, and 𝑓 counts children that exceed it. In a fixed child order, recovery selects as the direct continuation the first child that both fits its own continuation budget and covers every child that would otherwise exceed the sidebranch budget. Counting and selection visit each edge only a constant number of times, so recovery is also linear and can use slack in noncritical subtrees. The dynamic program and recovery together take 𝑂 (|𝑉 | + |𝐸|) time and 𝑂 (|𝑉 |) space. The minimum-call result is scoped to the execution model above: replay starts at the fixed nearest anchor in Equation (2); every non-replay output is produced exactly once; a state produced inside a call can initialize a separate sequence only after that call returns; and the planner neither uses an earlier anchor with multi-chunk replay nor duplicates segment computation to reduce call depth. Under these constraints, Equations (13)–(14) choose the minimum possible sequential call count, and recovery constructs a plan that attains it.

8: 9:

𝑇in [𝑣] ← min𝑜 max(𝑇in [𝑜], 1 + ℎ𝑜 ) 𝑇pre [𝑣] ← max𝑢 𝑇pre [𝑢] if 𝑒 𝑣 < 𝑎 𝑣 + 𝐵, else 𝑇in [𝑣]

10: 𝐾round ← 𝑇pre [root]; create zero-state root sequence 𝑞 0

in round 1 11: procedure Recover(𝑣, ready, 𝐾round, 𝑟, 𝑞) 12: append 𝑣 to 𝑞; if 𝑣 is a leaf then return 13: 𝑠𝑎𝑚𝑒 ← ready ∧ (𝑒 𝑣 < 𝑎 𝑣 + 𝐵)

𝛿 ← 0 if 𝑠𝑎𝑚𝑒, else 1; 𝐾side ← 𝐾round − 𝛿 for each Í child 𝑜: 𝐻 [𝑜] ← 𝑇pre [𝑜] if 𝑠𝑎𝑚𝑒, else 𝑇in [𝑜] 𝑓 ← 𝑢 I[𝑇pre [𝑢] > 𝐾side ] choose the first child 𝑜 such that 𝐻 [𝑜] ≤ 𝐾round and 𝑓 − I[𝑇pre [𝑜] > 𝐾side ] = 0 18: Recover(𝑜, 𝑠𝑎𝑚𝑒, 𝐾round , 𝑟 , 𝑞) 19: for each side child 𝑢 ≠ 𝑜 do 20: if 𝑠𝑎𝑚𝑒 then 21: launch 𝑞𝑢 in round 𝑟 from S𝑎 𝑣 with fixed replay 22: else 23: request S𝑏 𝑣 ; launch 𝑞𝑢 in round 𝑟 + 1 with fixed replay 24: Recover(𝑢, true, 𝐾side , 𝑟 + 𝛿, 𝑞𝑢 ) 25: Recover(root, true, 𝐾round , 1, 𝑞 0 ) 26: batch all sequences assigned to each round into one varlen call; return rounds 14: 15: 16: 17:

independent side branch and its fixed replay interval enter the earliest round permitted by state dependencies. All sequences in the same round form one packed batch and explicitly record their initial-state sources, replay intervals, and requested boundary states. The runtime consumes this deterministic plan without repeating tree planning. Appendix A expands the recurrence and recovery on nested forks inside chunks and on a high-fan-out tree.

Replay cost and packed output. If an internal segment 𝑣 has | ch(𝑣)| children, each of its | ch(𝑣)| − 1 side branches requires replay of length 𝑒 𝑣 − 𝑎(𝑒 𝑣 ) = 𝑒 𝑣 mod 𝐵. The total mandatory replay is therefore ∑︁ (| ch(𝑣)| − 1) (𝑒 𝑣 mod 𝐵). 𝑅= (15) 𝑣:ch(𝑣)≠∅

Here, 𝑅 counts repeated state-update positions inside the linear-attention core; it is not part of compact-token count 𝐶 or the non-replay compact-row compression metric. Its value is independent of which child becomes the direct continuation. The planner optimizes sequential call depth while holding mandatory replay fixed. After recovery, each directcontinuation chain forms one input sequence, while each

3.4

Prefix-Sharing Execution across Hybrid-Attention Layers

Section 3.3 fixes the input sequences in every round, the source of each sequence’s initial state, and its replay interval. Given this plan, the runtime derives separate execution views for full- and linear-attention layers from the same compact 10

HARTS: Hybrid-Attention RL over Tree Structures

Full attention on compact tokens. Full-attention layers require no recurrent-state interface. HARTS organizes compact-token queries by branch into packed Q and gathers the complete ancestor key/value prefix of each branch into a corresponding packed KV view. It then writes attention outputs back to their original compact positions. When multiple KV views read the same shared ancestor, the gather operation’s backward pass accumulates all gradients into that ancestor’s unique compact activation, reusing the standard attention backward path. Only the projected ancestor KV views needed by attention are expanded temporarily; queries and token-wise operations such as QKV projection, MLP, and MoE computation remain compact.

hidden states, while keeping the computation differentiable and recomputable. Differentiable boundary states. A conventional chunkwise linear-attention training interface often returns only outputs and an optional final state. HARTS needs a selective training contract that materializes only the chunk-boundary states requested by the execution plan and lets other input sequences consume them as differentiable initial states. Suppose input sequence 𝑖 in a packed call crosses 𝑛𝑖 execution chunks, and let I𝑖 be the chunk-boundary indices that the plan requests from this sequence. HARTS extends the persequence contract to   (𝑖 ) (𝑖 ) (𝑖 ) ChunkLA(Q𝑖 , K𝑖 , V𝑖 , S [0] ; I𝑖 ) → O𝑖 , S [𝑛 , {S } 𝑐 ∈ I 𝑖 . [𝑐 ] 𝑖] (16) All sequences in one round execute in a single packed call. This contract applies to a linear-attention training operator that advances recoverable recurrent state over chunks or blocks and accepts an initial state; it does not depend on the KDA-specific recurrence in Equation (1). KDA is the only backend that we fully implement and evaluate. The operator already produces intermediate boundary states to connect adjacent chunks during ordinary chunkwise state propagation. HARTS writes out only the small subset indexed by each I𝑖 and keeps these states live across the later execution rounds of the same layer; it does not add a second state scan. The additional costs are materializing and retaining the requested states for their consumers and for (𝑖 ) backward. When a later sequence uses S [𝑐 as its initial state, ] backward accumulates that initial-state gradient into the earlier computation that produced the boundary. If several side branches share one boundary tensor, autograd sums all consumer gradients at that tensor. Packed sequences with no explicit state sharing keep their states and gradients isolated.

Full-layer activation recomputation. All boundary states are produced and consumed within the forward pass of one linear-attention layer in one microbatch; no parent graph persists across microbatches or layers. Without activation recomputation, the requested states and their graph connections remain live until the layer’s backward pass. With standard full-layer activation recomputation, the backward pass re-executes the layer from the fixed microbatch plan and rematerializes the same packed inputs, rounds, replay intervals, boundary states, and graph connections without replanning the tree. HARTS therefore requires no forward graph to remain resident across microbatches.

3.5

Preserving RL and MoE Semantics without Expansion

Except for the bounded replay required for numerical alignment, Sections 3.3–3.4 produce each shared token’s model output only once. The RL objective, however, remains defined over per-trajectory positions in the original batch, and MoE router losses and expert-bias updates must still weight every token occurrence in that batch. A direct restoration would copy 𝐶 compact logit rows into 𝑆 rows and repeat the vocabulary-wide log-softmax on all 𝑆 rows. Doing so would create 𝑆 × 𝑉 logits, gradients, and memory traffic, forfeiting much of the benefit of prefix sharing. HARTS instead keeps vocabulary-dimensional and expert computation in compact coordinates and restores per-trajectory differences only in scalar objectives and expert-load counts.

Round-wise execution. Each execution round issues exactly one packed linear-attention call. Root sequences start from zero states; all other sequences start from boundary states returned by earlier rounds. Sequences in one round share a call but retain independent lengths, initial states, and recurrent-state evolution. When a fork lies inside a chunk, a side branch starts from the nearest chunk-boundary state and replays the ancestor updates from that boundary to the fork inside the linear-attention core. Replay positions update only recurrent state and neither produce nor write back semantic outputs. The corresponding ancestor projections and causal-convolution results are still materialized only once, and the final layer output is written only once. Replay does not repeat projection, causal convolution, MLP, or MoE computation. If a linear-attention block contains a bounded local operator such as causal convolution, its context must be constructed separately from the true ancestor tokens; core replay and adjacent packed sequences are not semantic history for that operator.

Compact-to-semantic log-probability computation. Let V = {1, . . . , 𝑉 } index the vocabulary. For compact token row 𝑐 ∈ C, let 𝑧𝑐 = (𝑧𝑐,𝑣 )𝑣 ∈ V ∈ R𝑉 be its logit row after the language-model (LM) head. Let P ⊆ S contain all nonterminal semantic positions in the original batch. For 𝑠 ∈ P, let 𝑦𝑠 ∈ V be the next token in the same trajectory, and let 𝑚(𝑠) identify the compact logit row for position 𝑠. Finally, let A ⊆ P contain the positions selected by the policy objective’s action mask. The current log-probability of each 𝑠 ∈ P 11

Boyuan Meng, Peihua Bao, Hong Liu, Xiaowei Zhu, Chao Wang, Gen Li, and Zhenxuan Pan

MoE semantics without expert expansion. For compact token 𝑐, define

is 𝜆𝑠 = log softmax(𝑧𝑚 (𝑠 ) ) 𝑦𝑠 = 𝑧𝑚 (𝑠 ),𝑦𝑠 − log

∑︁

exp 𝑧𝑚 (𝑠 ),𝑣 .

𝑞𝑐 = |{𝑠 ∈ S | 𝑚(𝑠) = 𝑐}| ,

𝑣∈V

(17) The implementation uses mapped cross-entropy to compute the negative of Equation (17), then negates the result to obtain 𝜆𝑠 . At a rollout-tree fork, one compact logit row can correspond to several semantic positions with different targets 𝑦𝑠 . HARTS computes the vocabulary-wide row maximum and sum-exp once per compact logit row, then reads each target logit in original batch order to obtain 𝜆𝑠 for every 𝑠 ∈ P. The original action mask selects 𝝀 = (𝜆𝑠 )𝑠 ∈ A . Vocabularydimensional work decreases from 𝑂 (𝑆𝑉 ) to 𝑂 (𝐶𝑉 ), while mapping, target lookup, and log-probability writes require 𝑂 (𝑆); the computation never constructs expanded 𝑆 × 𝑉 logits. With vocabulary parallelism, TP ranks reduce only the scalars needed for each compact row’s maximum and sumexp, plus one target-logit scalar per nonterminal semantic position. They never all-gather complete vocabulary logit rows.

where | · | denotes set cardinality. Thus, 𝑞𝑐 is the number of semantic-token occurrences Í in the original batch that map to compact token 𝑐, and 𝑐 ∈ C 𝑞𝑐 = 𝑆. With deterministic, no-token-drop top-𝑘 routing, a shared compact hidden state requires one router decision, dispatch, and expert computation. For per-token router z-loss, 𝑞𝑐 is the loss weight of compact token 𝑐. For load-balancing auxiliary loss, 𝑞𝑐 weights the accumulated routing-probability mass and top-𝑘 assignment count. If expert-bias load balancing is enabled, the same weighted assignment count updates each expert’s routing bias. Multiplicity changes neither routing decisions nor shared-prefix expert computation; it restores the training weight represented by each compact row in the original batch. In summary, vocabulary-sized tensors and their gradients remain organized over 𝐶 compact rows. Only linear-size objects scale with the original number of semantic positions 𝑆: labels, log-probabilities, loss metadata, multiplicities, and expert-load counts. MoE communication and expert execution use hardware-aligned physical buffers containing the 𝐶 compact token rows. These buffers may include alignment padding, but padding represents no semantic position and never restores the 𝑆 duplicate semantic rows. Under the stated policy-objective interface and deterministic no-tokendrop routing condition, HARTS preserves the original policy objective, router-loss weighting, and expert-bias update semantics without restoring duplicate model computation.

Reusing the original policy objective. Let 𝝃 denote algorithm-specific metadata aligned with the positions in A or their trajectories, such as advantages, rewards, masks, old or reference log-probabilities, loss weights, trajectory boundaries, and reduction denominators. If the original policy objective receives model outputs only through current log-probabilities for chosen tokens in original order, together with this metadata, then we can write it as Lpolicy = 𝐹 (𝝀, 𝝃 ).

(18)

The 𝝀 produced by HARTS has the same order and semantics as trajectory-wise computation and remains attached to the model’s computation graph. The metadata 𝝃 likewise retains the original action and trajectory structure. Consequently, 𝐹 is unaware of prefix sharing, and its cross-token or crosstrajectory coupling and reduction rules remain unchanged.

4

∑︁ 𝜕Lpolicy ∇𝑧𝑐 𝜆𝑠 . 𝜕𝜆𝑠 𝑠∈A

Implementation

We integrate HARTS with the AReaL Agentic RL framework [4] and the Megatron training stack [14], and evaluate the resulting system with Ling-3.0-tiny [5]. This hybrid-attention MoE model interleaves several KDA linearattention layers with one MLA full-attention layer in each repeated attention group, while its Transformer blocks use either dense or MoE FFNs. The same prefix-sharing layout must therefore satisfy MLA ancestor visibility, KDA recurrent-state dependencies, the local context of causal convolution, and MoE training statistics. Section 5.1 gives the exact layer counts, ratios, and expert configuration; the design of HARTS does not depend on these constants.

Gradient accumulation for reused compact logits. A compact logit row 𝑧𝑐 can supply log-probabilities to several semantic positions. By the chain rule, backward accumulates their gradient contributions into that same row: ∇𝑧𝑐 Lpolicy =

(20)

(19)

𝑚 (𝑠 )=𝑐

4.1

Thus, forward computes a shared logit row once, while backward preserves the contribution from every semantic position. Because positions can have different labels, advantages, masks, and loss weights, this operation sums independently formed gradients; it does not merely multiply one gradient by a reuse count. Backward produces 𝑑Logits only for compact rows and never creates an 𝑆 × 𝑉 gradient. Both forward and backward therefore cost 𝑂 (𝐶𝑉 + 𝑆) asymptotically.

System Integration

RL framework. The RL framework constructs rollout trees, retains per-token labels and RL metadata, and implements the two planners from Sections 3.2 and 3.3. The first planner chooses microbatch membership, DP replicas, and microbatch slots. The second produces per-round linearattention sequences, replay intervals, and state sources for every rollout tree in a microbatch, then merges same-round 12

HARTS: Hybrid-Attention RL over Tree Structures

inputs from independent trees. The model-training stack consumes only the fixed plans; it neither retraverses rollout trees nor solves a scheduling problem during execution.

For a branch with ℎ ancestor tokens and ℓ current tokens, its query row contains only the ℓ current tokens that require outputs, while its key/value row contains both the ℎ ancestors and the current tokens:

Megatron and backend. Megatron executes MLA, KDA, dense or MoE FFNs, and the LM head on the compact token layout. It also implements the differentiable state transfer from Section 3.4, the compact-to-semantic log-probability computation from Section 3.5, multidimensional parallelism, and activation recomputation. We extend the KDA training backend’s chunkwise state-propagation interface to write requested chunk-boundary states during the same recurrence and to propagate later sequences’ initial-state gradients through those states in backward. KDA is the only linear-attention backend implemented and evaluated in this work. Our Agentic RL instance uses Group Relative Policy Optimization (GRPO) [13]. Megatron returns current logprobabilities in original token order, while the RL framework retains its existing group-relative advantages, old or reference log-probabilities, masks, and reduction rules. The same interface also supports objectives such as Proximal Policy Optimization (PPO) [12] when their model-facing inputs consist of per-token current log-probabilities and the original metadata; their loss definitions remain unchanged. 4.2

|𝑄 | = ℓ,

|𝐾𝑉 | = ℎ + ℓ.

(21)

HARTS generates separate packed-sequence boundaries for Q and KV and uses bottom-right causal alignment from Transformer Engine [10]. Within one temporary attention row, the current branch has ℓ queries and ℎ ancestor KV entries. Let 𝑖 ∈ {0, . . . , ℓ − 1} index a query within the branch, and let 𝑗 ∈ {0, . . . , ℎ + ℓ − 1} index the concatenated KV row. Query 𝑖 can access exactly the entries satisfying 𝑗 ≤ ℎ + 𝑖,

(22)

which covers the complete ancestor prefix and all earlier tokens in the branch. Conventional top-left causal alignment would incorrectly align the shorter Q row with the left edge of the KV row. An indexed scatter writes attention outputs back to their unique compact positions. Projection, RoPE, and final outputs therefore operate on compact tokens, while the attention core temporarily expands only the projected ancestor KV views that it needs. HARTS does not eliminate this attention-core expansion across branch views. This mapping reuses a high-performance packed-attention backend and requires no arbitrary-tree custom kernel for MLA. The construction does not depend on MLA-specific projections and therefore also applies to other full-attention variants, including multi-head attention (MHA) and grouped-query attention (GQA), provided that the layer supports the asymmetric packed Q/KV views described above.

Compact Execution of MLA, KDA, and MoE

Figure 3 summarizes this implementation path. The call plan first fixes a compact physical ordering suited to linearattention execution. This ordering persists across model layers. Within a layer, MLA and KDA may derive temporary ancestor-KV, replay, context, or padded views, but MLA, KDA, and MoE all write their semantic outputs back to the persistent compact rows.

KDA state reuse and causal convolution. KDA directly consumes the physical ordering arranged by execution round: each round issues one packed call rather than one call per rollout-tree node. The extended chunkwise state propagation returns boundary states on demand, later sequences start from those states, and backward propagates their initial-state gradients into the corresponding earlier computation. When a fork lies inside a chunk, a side branch reuses compact activations whose projection and causal convolution have already executed; only the KDA core performs the replay defined in Section 3.3. KDA’s causal convolution must not treat an adjacent packed sequence as history. For convolution width 𝑤, HARTS temporarily prepends the most recent 𝑤 − 1 projected activations from the true ancestors of each independent input sequence. Segment identifiers prevent a convolution window from crossing sequence boundaries, and the runtime retains one post-convolution result for each compact token belonging to the sequence. Replay gathers these existing results rather than re-executing convolution. When several replay sequences read the same compact activation, indexed-gather backward accumulates all consumer gradients at its unique projection/convolution producer.

Compact physical ordering. Section 3.3 chooses the direct continuation at each fork and joins consecutively executed segments into input sequences. The implementation uses these choices to define one physical ordering for the entire model. Tokens in an input sequence are contiguous; all sequences in an execution round can directly form one packed KDA call; and the planned replay sources and requested boundary states are directly addressable. The model does not repeatedly restore trajectory-wise order between MLA or KDA layers and MoE FFNs. KDA consumes this ordering directly, MLA constructs temporary Q/KV views from it, MoE performs compact-token routing on it, and every layer writes its output back to the same physical positions. MLA execution with asymmetric packed Q/KV views. MLA projections and rotary position embeddings (RoPE) execute once per compact token, with RoPE using that token’s position ID in the original trajectory. HARTS then constructs attention views from the projected compact rows. 13

Boyuan Meng, Peihua Bao, Hong Liu, Xiaowei Zhu, Chao Wang, Gen Li, and Zhenxuan Pan

(a) Tree relation and KDA-friendly compact ordering Each box is a compact output span; rows in the same round form one variable-length call. Tree: A continues; B/C fork after P1

P0

P1

A0

A1

B0

B1

C0

C1

Round 1 · one sequence seq 0

P0

P1

A0

A1

seq 0

B0

B1

seq 1

C0

C1

(b) KDA replay for a fork inside a chunk

(c) MLA view for branch B

A, B, and C fork after P1; H0 is the reusable chunk-boundary state after P0.

Gathered from the compact rows in (a).

seq 0

zero

P0

P1

A0

A1

KV

P0

P1

B0

ancestor KV

returned H0 after P0

• round / sequence boundaries • state sources / replay intervals • semantic map / multiplicity

MoE C compact hidden rows

top-k routing + selected experts

compact outputs

B1

current branch

Q

B and C both fork after P1

(d) Compact compute, restored semantics

Projection + RoPE once per compact token

Round 1 · one variable-length KDA call

Plan metadata

Round 2 · two side sequences

B0

B1

semantic multiplicity qc

router z-/aux losses expert-load counts for bias update

bottom-right aligned

Round 2

one KDA call · two sequences

seq 0 shared H0 initial state for B and C

seq 1

Bottom-right causal mask (Q rows × KV cols)

replay P1 state update only no output

B0

B1

replay P1 state update only no output

C0

C1

P0

P1

B0

B1

B0

✓

✓

✓

—

B1

✓

✓

✓

✓

all ancestors

Replay boundary. The fork after P1 lies inside a KDA chunk; H0 after P0 is the nearest reusable boundary state. Round 2. One call packs two H0-initialized members, [replay P1, B] and [replay P1, C]. Replay positions update state but produce no output.

causal lower triangle

RL log-probabilities

C compact logit rows + row map + action labels

Mapped logprobability computation

actionordered log-probs

P0/P1: all ancestor keys. B0/B1: causal lower triangle. No repeated vocabulary-sized logits.

write MLA outputs to compact B0 / B1 rows

Figure 3. One compact physical ordering across hybrid-attention training. (a) The rollout tree’s shared path 𝑃0–𝑃1 forks into 𝐴, 𝐵, and 𝐶 after 𝑃1. The call planner chooses 𝐴 as the direct continuation, so Round 1 executes 𝑃0–𝑃1–𝐴, while Round 2 packs side branches 𝐵 and 𝐶 as two input sequences in one call. (b) The fork lies inside a KDA chunk. The nearest reusable boundary state, 𝐻 0 , is produced at the chunk boundary after 𝑃0. Round 1 continues along 𝐴 and returns 𝐻 0 ; Round 2 packs [replay(𝑃1), 𝐵] and [replay(𝑃1), 𝐶], both initialized by 𝐻 0 , into one KDA call. Replay positions update only KDA state and write no semantic outputs. (c) For side branch 𝐵, MLA gathers 𝑃0 and 𝑃1 from compact token rows as ancestor KV. The diagonal blocks for 𝐵0 and 𝐵1 each use a causal lower-triangular mask. 𝐵1 sees the complete earlier 𝐵0 block, while 𝐵0 cannot see future 𝐵1 tokens. Outputs return to their compact token rows. (d) MoE routing and selected-expert computation operate directly on compact hidden states. Semantic multiplicity 𝑞𝑐 restores router z-loss weights, routing mass and assignment counts for auxiliary loss, and expert-load counts for expert-bias updates. Compact-to-semantic log-probability computation combines compact logits, row mappings, and token labels into current log-probabilities in original token order, without expanding duplicate vocabulary-sized logits. 4.3

This bounded temporary input view supplies the ancestor context required by causal convolution. Recurrent-state handoffs carry KDA state across rounds, with core replay restoring it when necessary.

Parallel Execution and Recomputation

HARTS composes with existing DP, TP/SP, PP, and EP training parallelism and also supports the context-parallel (CP) scheme in DeepSpeed-Ulysses [6]. The planner assigns microbatches to DP replicas and execution slots to balance compact-token work both within each slot and over the training step. TP partitions model parameters as usual. With SP, devices in a TP group divide the hardware-aligned compact physical rows evenly. All PP stages use the same microbatch order and compact layout. EP routes and computes experts only for compact tokens. DeepSpeed-Ulysses uses all-to-all communication to transform a sequence-dimension partition into partitions over independent attention heads. Because this transformation depends on the presence of multiple

MoE execution and objectives. The MoE router receives the 𝐶 compact hidden states directly. Under deterministic top-𝑘 routing with no token dropping, each compact token row 𝑐 undergoes routing, EP communication, selected-expert computation, and combination once, and the result returns to that row. A shared prefix is never expanded according to its multiplicity 𝑞𝑐 for execution. The implementation weights per-token router z-loss by 𝑞𝑐 , computes load-balancing auxiliary loss from 𝑞𝑐 -weighted routing-probability mass and assignment counts, and feeds the same weighted assignment counts to expert-bias updates. These weights change neither routing decisions, dispatch, nor expert computation. 14

HARTS: Hybrid-Attention RL over Tree Structures

Table 1. Ling-3.0-tiny model used in evaluation. Item

Configuration

Model Transformer layers Attention layout Hidden size / heads Dense FFN MoE FFN Expert intermediate KDA convolution width KDA chunk size Vocabulary / precision Hardware

Decoder-only hybrid-attention MoE 24 6 × (3 KDA + 1 MLA) 1,536 / 16 Layer 1, intermediate 4,608 Layers 2–24, 128 experts, top-8 512 4 64 157,184 / BF16 8× H20-3e GPUs

5.2 Prefix-Aware Microbatch Planning and Schedule Construction On the recorded rollout-tree workload with 𝑇 = 65,536, we compare HARTS against a planning baseline built from the OR-Tools partitioner [11] used by Tree Training [16]. Tree Training uses OR-Tools to divide precollected rollout trees into capacity-constrained partitions, but does not determine DP-replica placement or slot schedules. Our OR-Tools planning baseline first assigns rollout trees to DP replicas using shard-level first-fit decreasing (FFD), runs the Tree Training OR-Tools partitioner independently within each replica, and performs additional splits until every replica has the same number of microbatches. Because grouping occurs independently after data distribution, this construction cannot jointly optimize replica work within a slot or over a training step. This comparison evaluates partitioning and scheduling, not an end-to-end Tree Training runtime. HARTS instead considers all trajectories in the current step together, jointly determines microbatch membership, DP-replica placement, and slot order, first minimizes compact-token work, and then reduces synchronization-critical and cumulative replica work.

attention heads rather than the internal attention implementation, HARTS’s modified attention execution does not affect CP. With full-layer activation recomputation, backward reuses the plan fixed during forward and re-executes the same compact layout, linear-attention calls, replay, and state handoffs, without retaining activations or a computation graph across microbatches.

5

Planning quality and planning-time speedup. Table 2 reports final compact-token work and planning-time speedup for ten rollout-tree steps. Both ratios are ORTools/HARTS: Compact OR/HARTS above one means that HARTS produces less compact work, and Planning time OR/HARTS above one means that HARTS plans faster. HARTS produces less compact work in every step. Summing the ten table rows, it reduces compact tokens from 80,213,867 to 77,581,958, a reduction of 2,631,909 tokens or 3.28% relative to OR-Tools. Per-step compact ratios range from 1.007421× to 1.079273×, and planning-time speedups range from 7.236× to 9.732×, covering different workload sizes and tree structures.

Evaluation

We evaluate HARTS by asking four questions. (1) Does prefixaware planning retain shared work and balance DP replicas? (2) Does compact-row reduction translate into speedup across parallel configurations? (3) Do replay and compact hybrid-model execution avoid a measurable shift in the training signal? (4) Does complete online Agentic RL retain a similar learning trend? 5.1

DP-replica work balance. We next compare static DP load across all ten rollout-tree batches. HARTS targets two complementary properties: within-slot balance, which reduces waiting caused by the heaviest replica in a synchronized slot, and current-step-total balance, which balances each replica’s cumulative work over all slots in the current optimization step. Let 𝑑 ∈ {1, . . . , 𝐷 } index a DP replica, where 𝐷 is the DP degree, and let 𝑤𝑑 be the compact-token work of replica 𝑑 in the scope being measured. Both properties use the same imbalance metric:

Experimental Setup

We use Ling-3.0-tiny, a hybrid-attention MoE model composed of MLA and KDA, and run the performance and training-engine fidelity experiments on eight H20-3e GPUs. Table 1 summarizes the model. Each of its six attention groups contains three KDA layers followed by one MLA layer, so prefix sharing repeatedly alternates between recurrent-state and ancestor-KV paths within one Transformer forward pass. The performance experiments use an Agentic RL rollouttree batch generated by a Claude Code scaffold executing SWE-bench tasks [7]. Each DP replica has microbatch capacity 𝑇 = 65,536. The baseline expands the trees into root-toleaf trajectories and executes them conventionally, whereas HARTS applies the compact planning and execution path described above. Both use the same stated training configuration and enable full-layer activation recomputation.

max 𝑤𝑑 1≤𝑑 ≤𝐷 𝐼 = 1 Í𝐷 − 1. 𝑑=1 𝑤𝑑 𝐷

(23)

This metric is the excess work of the heaviest DP replica relative to the mean; lower values indicate better balance. For within-slot balance, each synchronized slot is one observation with 𝑤𝑑 = 𝑤 (𝑀𝑑,𝑗 ), and the table reports the 15

Boyuan Meng, Peihua Bao, Hong Liu, Xiaowei Zhu, Chao Wang, Gen Li, and Zhenxuan Pan

Table 2. Prefix-aware planning quality and planning-time speedup across ten rollout-tree steps.

Step Paths 1 2 3 4 5 6 7 8 9 10

1,551 1,693 2,368 1,283 1,613 1,752 1,457 1,411 1,738 1,304

Raw tokens

Compact Planning time Harts OR-Tools compact tokens compact tokens OR/Harts OR/Harts

42,329,859 44,799,214 58,835,220 34,603,796 43,931,290 48,612,718 38,966,280 38,541,212 48,352,605 36,062,624

7,780,586 7,670,675 7,747,595 5,659,215 8,266,355 9,671,964 7,015,681 7,504,249 9,188,374 7,077,264

8,058,974 7,915,111 8,361,767 5,924,124 8,556,470 9,774,737 7,233,955 7,697,111 9,256,559 7,435,059

P50 and P95 imbalance over all slots. For current-step-total balance, Í each optimization step is one observation with 𝑤𝑑 = 𝑗 𝑤 (𝑀𝑑,𝑗 ), and the table reports P50 and maximum imbalance across the ten rollout-tree batches. Table 3 shows that HARTS improves both forms of DPreplica work balance. Within a slot, both P50 and P95 imbalance decrease substantially, so compact-token work is closer across replicas in typical and high-percentile synchronized rounds. For current-step totals, both P50 and maximum imbalance remain lower, indicating stable balance in cumulative work across slots. HARTS also uses fewer scheduled slots, so these improvements do not come from adding microbatches or duplicating more shared prefixes. These results follow from modeling compact work together with both DP-replica scheduling objectives. Slotcritical work is nonseparable across replicas: the heaviest microbatch determines the slot’s compact-work critical path. Current-step totals instead accumulate each replica’s work across slots. HARTS compares total compact work, aggregate slot-critical work, and maximum replica-cumulative work. It then greedily constructs the slot composition and replica placement for each retained grouping. This combination reduces synchronization-critical work within slots and avoids persistently concentrating the current step’s work on one replica. Overall, HARTS lowers final compact work and planning time while improving both within-slot and current-step-total balance. 5.3

1.035780× 1.031866× 1.079273× 1.046810× 1.035096× 1.010626× 1.031112× 1.025700× 1.007421× 1.050556×

9.732× 7.899× 7.236× 7.403× 8.002× 8.993× 7.571× 8.164× 8.845× 9.213×

backward, parameter-gradient computation, and the communication and synchronization on this path. The interval base and 𝑇 H denote ends before the optimizer update. Let 𝑇train train this interval for the baseline and HARTS, respectively. Let H denote HARTS planning time, including prefix-aware 𝑇plan microbatch/slot planning and per-microbatch minimumcall planning; the trajectory-wise baseline invokes neither HARTS planner. The two reported speedups are 𝑆 F/B/Grad =

base 𝑇train

, H 𝑇train

𝑆 Core =

base 𝑇train . H H 𝑇plan + 𝑇train

Thus, Core speedup charges all HARTS planner latency against the same complete training interval; neither metric includes the optimizer update. Table 4 shows 5.62–5.63× non-replay compact-row compression across the three parallel configurations, translating into 4.81–4.87× F/B/Grad speedup and 4.39–4.63× Core speedup. DP2/TP1/PP4/EP2 has the highest F/B/Grad speedup, 4.8650×, while DP4/TP2/SP/PP1/EP8 has the highest Core speedup, 4.6281×. Even when TP/SP, PP, and EP are combined, F/B/Grad and Core speedups remain at least 4.8053× and 4.3886×, respectively. Non-replay compact-row compression does not translate one-for-one into F/B/Grad speedup because several costs do not shrink in proportion to 𝐶: KDA replay at forks inside a chunk, compact-runtime materialization, materialization and gradient traffic for requested boundary states, packed-call launch overhead, TP/PP/EP communication and synchronization, and the expanded projected ancestor-KV views consumed by MLA attention. Core speedup is lower still because the Core interval additionally includes all HARTS planner time. These two gaps are consistent with the lower conversion in the configuration that combines TP/SP with PP.

Efficiency across Parallel Configurations

We report three metrics. Non-replay compact-row compression is the ratio of token rows in the unshared trajectory baseline to unique, non-replay compact token rows in HARTS. It measures the model-level row reduction from prefix deduplication, not total operator work: KDA replay, temporary MLA ancestor-KV views, and backend padding are excluded from the count and appear only in elapsed time. F/B/Grad speedup compares the complete measured training path once an execution plan is available, including compact-runtime materialization, model forward, mapped cross-entropy and policy-loss computation,

5.4

Training Fidelity and RL Effectiveness

We evaluate training fidelity at three levels. First, we isolate replay using a fixed checkpoint and input. Next, we compare full-model logits across five parallel configurations. Finally, we use online reward on 𝜏 3 -Bench [1] to observe end-to-end Agentic RL learning behavior. 16

HARTS: Hybrid-Attention RL over Tree Structures

Table 3. Balancing DP-replica compact-token work within each slot and in current-step totals. Results cover all ten rollout-tree batches; lower imbalance is better. Scheduled slots Harts/OR

Topology

DP4/TP2/SP/PP1/EP8 DP2/TP1/PP4/EP2 DP2/TP2/SP/PP2/EP4

301/319 599/614 599/614

Balanced within each slot

Balanced current-step totals

P50 imbalance (%) Harts/OR

P95 imbalance (%) Harts/OR

P50 imbalance (%) Harts/OR

Max imbalance (%) Harts/OR

0.0247/0.2628 0.0069/0.0851 0.0069/0.0851

2.1379/34.6410 0.4963/16.1598 0.4963/16.1598

0.2607/3.2079 0.0343/0.7799 0.0343/0.7799

1.1050/5.6756 0.4667/3.8462 0.4667/3.8462

Table 4. Training efficiency on rollout-tree workload generated by running SWE-bench tasks with a Claude Code scaffold, using one eight-GPU node. Baseline and Harts both enable full-layer activation recomputation with microbatch capacity 𝑇 = 65,536.

Parallel configuration

Non-replay compact-row compression

F/B/Grad Speedup

Core Speedup

DP4/TP2/SP/PP1/EP8 DP2/TP1/PP4/EP2 DP2/TP2/SP/PP2/EP4

5.6219× 5.6331× 5.6330×

4.8352× 4.8650× 4.8053×

4.6281× 4.4611× 4.3886×

Replay ablation. This experiment uses the same Ling-3.0tiny checkpoint, rollout-tree input, and microbatch boundaries to compare four Megatron/KDA execution paths: Baseline A, the common reference; Baseline B, an ordinary rerun used to estimate repeat-execution noise; HARTS Replay; and No-replay, which passes the parent’s final state directly without restoring the original chunk phase. Every error is measured relative to Baseline A. Let 𝑋 ∈ {𝐵, 𝑅, 𝑁 } identify 𝑋 be the logit the latter three paths, respectively, and let 𝑧𝑡,𝑣 from path 𝑋 at token 𝑡 and vocabulary entry 𝑣. Define 𝑝𝑡𝑋 (𝑣) = Í

When a fork lies inside a chunk rather than on its boundary, the child starts from a legal boundary state and performs replay. Crossing one parent–child segment boundary counts as one handoff. Thus, handoff depth is a logical segment-tree depth, not a count of materialized state tensors. Tokens in the root segment have 𝑑 = 0; tokens in a child segment at depth 𝑑 have inherited state 𝑑 times. We report 𝑑 = 0, . . . , 5. Figure 4 shows that HARTS Replay remains close to baseline-rerun variation at every handoff depth. No-replay develops a noticeably larger error after the first handoff, and the gap generally grows on deeper branches. At 𝑑 = 4–5, No-replay exceeds baseline-rerun error by 32.4%–36.1% in SymKL and by 17.0%–18.7% in mean |ΔNLL|. Resuming at the nearest chunk boundary and replaying less than one chunk is therefore necessary to align shared execution with the chunkwise state partitioning of conventional training under this backend.

𝑋 exp 𝑧𝑡,𝑣

.

(24)

𝑋 𝑢 ∈ V exp 𝑧𝑡,𝑢

For the set T𝑑 of tokens with state-handoff depth 𝑑, the left panel of Figure 4 reports mean per-token symmetric KL divergence: 1 ∑︁ 1  SymKL𝑑 (𝑋, 𝐴) = 𝐷 KL (𝑝𝑡𝐴 ∥𝑝𝑡𝑋 ) |T𝑑 | 𝑡 ∈ T 2 (25) 𝑑  𝐴 𝑋 + 𝐷 KL (𝑝𝑡 ∥𝑝𝑡 ) , Í where 𝐷 KL (𝑝 ∥𝑞) = 𝑣 𝑝 (𝑣) log[𝑝 (𝑣)/𝑞(𝑣)]. This metric compares the complete vocabulary distribution at each token; lower values indicate closer agreement with the reference. Let 𝑦𝑡 be the next-token target at position 𝑡 and NLL𝑡𝑋 = − log 𝑝𝑡𝑋 (𝑦𝑡 ). The right panel reports 1 ∑︁ MeanAbsDeltaNLL𝑑 (𝑋, 𝐴) = NLL𝑡𝑋 − NLL𝑡𝐴 . |T𝑑 | 𝑡 ∈ T

Full-model logit similarity. We next compare complete model outputs. Let T be the set of aligned token positions in two executions, and let 𝑧𝑡𝐴 , 𝑧𝑡𝑋 ∈ R𝑉 be the full-vocabulary logits from the baseline and comparison path at position 𝑡. We first compute logit cosine similarity for each token and then average tokens with equal weight: MeanCos(𝐴, 𝑋 ) =

1 ∑︁ ⟨𝑧𝑡𝐴 , 𝑧𝑡𝑋 ⟩ . |T | 𝑡 ∈ T ∥𝑧𝑡𝐴 ∥ 2 ∥𝑧𝑡𝑋 ∥ 2

(27)

𝑑

(26) Mean target-NLL deviation directly measures the change in the target token’s current log-probability, independently of advantages, masks, and loss reduction. We call the dependency in which a child segment uses recurrent state produced by its parent segment a state handoff.

Table 5 shows that HARTS MeanCos exceeds 0.9997 both with and without full-layer activation recomputation and remains close to baseline self-rerun variation. Enabling recomputation produces no visible additional gap, consistent with rematerializing the same per-round plan, replay, and boundary-state connections. The last column reports the 17

Boyuan Meng, Peihua Bao, Hong Liu, Xiaowei Zhu, Chao Wang, Gen Li, and Zhenxuan Pan

Replay vs A

(a) Full-token SymKL

3.9

No-replay vs A

40%

3.6

30%

3.3

20%

3.0

10%

2.7

0%

d=0

d=1

d=2

d=3

d=4

d=5

n=433.8k

n=169.6k

n=174.4k

n=42.8k

n=27.9k

n=24.7k

No-replay / baseline repeat noise (%)

(b) Full-token mean |ΔNLL| Mean |ΔNLL| (nats per token)

Mean symmetric KL per token (×10−3)

Baseline B vs A

40%

0.059

30%

0.056

20%

0.053

10%

0.050

0%

d=0

d=1

d=2

d=3

d=4

d=5

n=433.8k

n=169.6k

n=174.4k

n=42.8k

n=27.9k

n=24.7k

State-handoff depth d (discrete)

State-handoff depth d (discrete)

Figure 4. Replay preserves chunkwise numerical fidelity across state-handoff depth. The three solid lines on the left axis report the absolute errors in Equations (25) and (26) for Baseline B, Replay, and No-replay relative to the common reference, Baseline A. The dashed orange line on the right axis reports how much No-replay exceeds baseline-rerun error, (𝐸 𝑁 /𝐸𝐵 − 1) × 100%. Token counts for each depth appear below the horizontal axis; lower is better for every metric. Table 5. Mean token-wise full-vocabulary logit cosine and loss difference. Baseline self-rerun compares two ordinary executions; the two Harts cosine columns compare against the matched baseline under the stated recomputation setting. Absolute loss diff reports the matched Baseline–Harts scalar-loss difference for the Harts run without activation recomputation.

Parallel configuration

Baseline self-rerun

Harts w/o recomputation

Harts w/ recomputation

Absolute loss diff

DP8/TP1/PP1/EP8 DP4/TP2/SP/PP1/EP8 DP4/TP1/PP2/EP4 DP2/TP2/SP/PP2/EP4 DP2/TP1/PP4/EP2

0.99978662 0.99975827 0.99975816 0.99972408 0.99972486

0.99977463 0.99974703 0.99974010 0.99971610 0.99971866

0.99977816 0.99974416 0.99974332 0.99971732 0.99971735

1.99 × 10−9 7.45 × 10−10 7.45 × 10−10 9.19 × 10−9 9.19 × 10−9

6

absolute scalar-loss difference between the matched baseline and HARTS run without activation recomputation. It is below 10−8 in all five parallel configurations, providing additional evidence that compact execution does not introduce a measurable loss shift in these experiments.

Conclusion

We presented HARTS, a hybrid-attention training system for Agentic RL over arbitrary rollout trees. HARTS assigns prefix-aware microbatches to DP replicas and slots. A linear-time planner coordinates reuse and replay of chunkboundary states. It repeats only the core state updates required for numerical alignment; projections, MLP/MoE computation, and final outputs execute once. Under the stated packed chunkwise execution model, the schedule minimizes sequential linear-attention calls. A single compact layout feeds every model layer. State handoffs are differentiable and compatible with activation recomputation. Mapped log-probabilities restore the original RL interface. Under deterministic top-𝑘 routing with no token dropping, semantic multiplicities restore MoE objective weights. We evaluate HARTS on an Agentic RL rollout-tree workload generated from SWE-bench tasks, with full-layer activation recomputation and several parallel configurations.

Online 𝜏 3 -Bench reward. Finally, we start two complete online Agentic RL runs on 𝜏 3 -Bench from the same checkpoint and training configuration and compare their first 120 common training steps. In Figure 5, faint lines show raw perstep reward and thick solid lines show a seven-step centered moving average. Over these 120 steps, the two reward curves have similar overall trends and variation. Together with the replay ablation and full-model cosine results, this experiment shows that HARTS maintains a stable learning trend in the complete online Agentic RL pipeline. 18

HARTS: Hybrid-Attention RL over Tree Structures

[2] Tianqi Chen, Bing Xu, Chiyuan Zhang, and Carlos Guestrin. 2016. Training Deep Nets with Sublinear Memory Cost. In Advances in Neural Information Processing Systems. Curran Associates, Inc., Barcelona, Spain, 2116–2125. [3] Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. 2022. FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. In Advances in Neural Information Processing Systems. Curran Associates, Inc., New Orleans, LA, 16344–16359. [4] Wei Fu, Jiaxuan Gao, Xujie Shen, Chen Zhu, Zhiyu Mei, Chuyi He, Shusheng Xu, Guo Wei, Jun Mei, Jiashu Wang, Tongkai Yang, Binhang Yuan, and Yi Wu. 2025. AReaL: A Large-Scale Asynchronous Reinforcement Learning System for Language Reasoning. In Advances in Neural Information Processing Systems, Vol. 38. 36256–36282. doi:10.52202/085713-1218 [5] InclusionAI. 2026. Ling-3.0-tiny. Hugging Face model card. https: //huggingface.co/inclusionAI/Ling-3.0-tiny [6] Sam Ade Jacobs, Masahiro Tanaka, Chengming Zhang, Minjia Zhang, Shuaiwen Leon Song, Samyam Rajbhandari, and Yuxiong He. 2023. DeepSpeed Ulysses: System Optimizations for Enabling Training of Extreme Long Sequence Transformer Models. arXiv:2309.14509 [cs.LG] [7] Carlos E. Jimenez, John Yang, Alexander Wettig, Shunyu Yao, Kexin Pei, Ofir Press, and Karthik R. Narasimhan. 2024. SWE-bench: Can Language Models Resolve Real-world Github Issues?. In The Twelfth International Conference on Learning Representations. OpenReview.net, Vienna, Austria, 21 pages. https://openreview.net/forum? id=VTF8yNQM66 [8] Kimi Team. 2025. Kimi Linear: An Expressive, Efficient Attention Architecture. arXiv:2510.26692 [cs.LG] [9] Mario Michael Krell, Matej Kosec, Sergio P. Perez, and Andrew Fitzgibbon. 2021. Efficient Sequence Packing without Cross-Contamination: Accelerating Large Language Models without Impacting Performance. arXiv:2107.02027 [cs.LG] [10] NVIDIA Corporation. 2026. NVIDIA Transformer Engine. Software library. https://github.com/NVIDIA/TransformerEngine [11] Laurent Perron and Vincent Furnon. 2025. OR-Tools. https:// developers.google.com/optimization/ Version 9.12. [12] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. 2017. Proximal Policy Optimization Algorithms. arXiv:1707.06347 [cs.LG] [13] Zhihong Shao, Peiyi Wang, Qihao Zhu, Runxin Xu, Junxiao Song, Xiao Bi, Haowei Zhang, Mingchuan Zhang, Y. K. Li, Y. Wu, and Daya Guo. 2024. DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models. arXiv:2402.03300 [cs.CL] [14] Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro. 2019. Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism. arXiv:1909.08053 [cs.CL] [15] Yutao Sun, Li Dong, Shaohan Huang, Shuming Ma, Yuqing Xia, Jilong Xue, Jianyong Wang, and Furu Wei. 2023. Retentive Network: A Successor to Transformer for Large Language Models. arXiv:2307.08621 [cs.CL] [16] Jinghui Wang, Shaojie Wang, Yinghan Cui, Xuxing Chen, Chao Wang, Liang Huang, Can Tang, Xiaojiang Zhang, Junyi Peng, Li Wan, Haotian Zhang, and Bin Chen. 2026. Tree Training: Accelerating Agentic LLMs Training via Shared Prefix Reuse. arXiv:2511.00413 [cs.LG] [17] An Yang, Anfeng Li, Baosong Yang, et al. 2025. Qwen3 Technical Report. arXiv:2505.09388 [cs.CL] [18] Songlin Yang, Jan Kautz, and Ali Hatamizadeh. 2025. Gated Delta Networks: Improving Mamba2 with Delta Rule. In The Thirteenth International Conference on Learning Representations. OpenReview.net, Singapore, 18 pages. [19] Songlin Yang, Bailin Wang, Yikang Shen, Rameswar Panda, and Yoon Kim. 2024. Gated Linear Attention Transformers with HardwareEfficient Training. In Proceedings of the 41st International Conference

HARTS

τ 3-Bench Reward

Baseline

0

20

40

60

80

100

Training Step

Figure 5. 𝜏 3 -Bench reward over the first 120 online training steps. HARTS and the baseline show similar overall learning trends and ranges of variation. Faint lines are raw per-step rewards; thick lines are seven-step centered moving averages.

Non-replay compact-row compression is 5.62–5.63×. This reduction yields 4.81–4.87× F/B/Grad speedup and 4.39–4.63× Core speedup. Replay and full-model experiments show numerical differences close to baseline self-rerun variation, with mean token-wise full-vocabulary logit cosine above 0.9997. The first 120 steps of online 𝜏 3 -Bench training also exhibit a reward trend similar to the baseline. Several residual costs separate compact-row compression from measured speedup. Discrete tree partitioning leaves some within-slot imbalance. KDA incurs bounded replay, while requested states add materialization and gradient traffic. Packed calls add launch overhead. Parallel execution adds communication and synchronization, and MLA expands projected ancestor-KV views. A synchronized slot runs at the pace of its slowest DP replica, so residual imbalance reduces elapsed-time savings even when current-step totals are close. Future work can incorporate these costs into placement and optimize state materialization, kernel launches, and collective overlap. The benefits of HARTS are expected to increase with model scale. As layer count, hidden dimension, and expert computation grow, sharing one prefix token saves more model computation, while control costs such as planning, layout materialization, and kernel launch grow more slowly. Better amortization is therefore expected to bring speedup on larger models closer to the ideal indicated by non-replay compact-row compression.

References [1] Victor Barres, Honghua Dong, Soham Ray, Xujie Si, and Karthik Narasimhan. 2025. 𝜏 2 -Bench: Evaluating Conversational Agents in a Dual-Control Environment. arXiv:2506.07982 [cs.AI] 19

Boyuan Meng, Peihua Bao, Hong Liu, Xiaowei Zhu, Chao Wang, Gen Li, and Zhenxuan Pan

on Machine Learning. PMLR, Vienna, Austria, 56501–56523. [20] Songlin Yang, Bailin Wang, Yu Zhang, Yikang Shen, and Yoon Kim. 2024. Parallelizing Linear Transformers with the Delta Rule over Sequence Length. In Advances in Neural Information Processing Systems. Curran Associates, Inc., New Orleans, LA, 16 pages. [21] Jiarui Zhang, Yuchen Yang, Ran Yan, Zhiyu Mei, Liyuan Zhang, Daifeng Li, Wei Fu, Jiaxuan Gao, Shusheng Xu, Yi Wu, and Binhang Yuan. 2026. AReaL-DTA: Dynamic Tree Attention for Efficient Reinforcement Learning of Large Language Models. arXiv:2602.00482 [cs.LG]

20

HARTS: Hybrid-Attention RL over Tree Structures

A

Linear-Attention Execution Planning Cases

Figures 6–8 apply Algorithm 1 with 𝐵 = 64; 𝑅 denotes the mandatory replay total from Equation (15), and each caption reports packed members and call assignments.

R1 3 tok Call 1

r=3

A1 2 tok Call 1

B1 632 tok Call 1

r=5

A2 403 tok Call 1

A3 10 tok Call 1

A5 21119 tok Call 1

r=24

A6 1355 tok Call 1

A7 2 tok Call 2

B2 1 tok Call 1

A9 3588 tok Call 2

r=2

B3 15 tok Call 2

C2 1 tok Call 1

B6 13 tok Call 1

B8 13 tok Call 2

C3 45 tok Call 2

r=60

B4 10 tok Call 1

r=3

B5 14 tok Call 2

C4 10 tok Call 1

r=6

r=26

A8 1 tok Call 2

C1 255 tok Call 1

r=59

r=15

A4 7753 tok Call 1

r=3

r=13

B7 6 tok Call 2

C6 110 tok Call 1

r=12

r=59

B9 11 tok Call 2

C8 630 tok Call 1

r=30

C7 6 tok Call 2

C14 189 tok Call 2

C15 187 tok Call 2

r=19

C9 630 tok Call 2

r=27

A10 29628 tok Call 2

C5 27 tok Call 2

C10 116 tok Call 2

C11 553 tok Call 2

r=7

A11 3604 tok Call 2

C12 435 tok Call 2

C13 427 tok Call 3

Figure 6. Case I: three calls and 𝑅 = 373. Each circle represents one segment and shows its token count. Thick solid edges form direct-continuation chains, while dashed edges start independent members and are labeled with their replay length under chunk_size= 64. Call 1 (five members; replay 26). 𝑅1 → 𝐴1 → 𝐴2 (𝑟 = 0); 𝐴3 → 𝐴4 → 𝐴6 (5); 𝐴5 (15); 𝐵1 → 𝐵2 → 𝐵4 → 𝐵6 (3); and 𝐶1 → 𝐶2 → 𝐶4 → 𝐶6 → 𝐶8 (3).

Calls 2–3. Call 2 has thirteen members: 𝐴7 → 𝐴8 → 𝐴10 (24), 𝐴11 (27), 𝐴9 (26), 𝐵7 → 𝐵8 (6), 𝐵9 (12), 𝐵5 (60), 𝐵3 (59), 𝐶9 (59), 𝐶7 → 𝐶10 → 𝐶12 (13), 𝐶11 (19), 𝐶5 → 𝐶14 (3), 𝐶15 (30), and 𝐶3 (2), totaling 340 replay tokens. Call 3 contains only 𝐶13 (𝑟 = 7).

21

Boyuan Meng, Peihua Bao, Hong Liu, Xiaowei Zhu, Chao Wang, Gen Li, and Zhenxuan Pan D subtree R2 3 tok Call 1

D1 2 tok Call 1

r=5

D2 399 tok Call 1

D3 16884 tok Call 1

r=57

D4 59 tok Call 2

D5 10346 tok Call 1

r=35

D6 133 tok Call 2

D7 6555 tok Call 1

r=62

D8 1902 tok Call 1

D9 70 tok Call 2

E subtree r=3

E1 633 tok Call 1

r=60

E2 1 tok Call 1

r=3

E3 1 tok Call 2

r=61

E6 2 tok Call 1

r=60

r=61

E7 15 tok Call 2

E10 15 tok Call 2

r=61

E11 15 tok Call 2

E12 16 tok Call 2

r=60

E4 15 tok Call 2

r=61

r=61

E13 15 tok Call 2

E5 15 tok Call 2

r=61

E14 15 tok Call 2

E15 14 tok Call 2

r=63

E8 13 tok Call 1

E9 14 tok Call 2

F subtree F1 256 tok Call 1

r=3

r=3

F2 1 tok Call 2

F3 1 tok Call 2

r=4

F6 2 tok Call 2

F10 346 tok Call 2

r=3

F11 1469 tok Call 2

r=4

F12 390 tok Call 2

r=4

r=4

F13 168 tok Call 2

F4 170 tok Call 2

r=4

F5 445 tok Call 1

r=4

F14 335 tok Call 2

F15 226 tok Call 2

F9 155 tok Call 2

F16 10 tok Call 1

F17 5 tok Call 2

r=6

F7 220 tok Call 2

F8 212 tok Call 2

Figure 7. Case II: two calls and 𝑅 = 813. The three child subtrees of the single shared root 𝑅2 are unfolded vertically only to preserve label readability. Under chunk_size= 64, replay labels expose how many branches become ready together for the second variable-length call. Call 1 (four members; replay 11). 𝑅2 → 𝐷1 → 𝐷2 (𝑟 = 0); 𝐷3 → 𝐷5 → 𝐷7 → 𝐷8 (5); 𝐸1 → 𝐸2 → 𝐸6 → 𝐸8 (3); and 𝐹 1 → 𝐹 5 → 𝐹 16 (3).

Call 2 (24 members; replay 802). 𝐷4, 𝐷6, 𝐷9 replay 57, 35, 62. On the 𝐸 side: 𝐸9 (63), 𝐸7 (61), 𝐸3 → 𝐸10 (60), 𝐸11–𝐸15 (61 each), and 𝐸4, 𝐸5 (60 each). On the 𝐹 side: 𝐹 2 → 𝐹 6 → 𝐹 7 (3), 𝐹 8 (6), 𝐹 3 → 𝐹 9 (3), 𝐹 10–𝐹 15 (4 each), 𝐹 4 (3), and 𝐹 17 (0).

22

HARTS: Hybrid-Attention RL over Tree Structures

R3 3 tok Call 1 r=3

G1 2 tok Call 1 r=5

G2 397 tok Call 1

G3 13257 tok Call 1

G4 585 tok Call 1

G5 251 tok Call 2

r=3

H1 633 tok Call 1 r=60

H2 1 tok Call 1

r=14

H3 15 tok Call 2

I1 256 tok Call 1 r=60

r=3

H4 15 tok Call 2

r=61

H5 14 tok Call 1

I2 1 tok Call 1

r=3

I3 943 tok Call 2

I4 282 tok Call 2

r=4

H6 15 tok Call 2

I5 1243 tok Call 1

I6 410 tok Call 2

r=23

G6 480 tok Call 1

G7 225 tok Call 2

r=55

G8 644 tok Call 1

G9 319 tok Call 2

r=59

G10 999 tok Call 1

G11 49 tok Call 2

r=34

G12 772 tok Call 1

G13 485 tok Call 2

r=38

G14 1527 tok Call 1

G15 38 tok Call 2

r=29

G16 116 tok Call 1

G17 63 tok Call 2

r=17

G18 39239 tok Call 1

G19 76 tok Call 2

r=24

G20 477 tok Call 1

G21 222 tok Call 2

r=53

G22 3247 tok Call 1

G23 765 tok Call 2

r=36

G24 838 tok Call 1

G25 550 tok Call 2

r=42

G26 909 tok Call 1

G27 799 tok Call 2

r=55

G28 2111 tok Call 1

G29 867 tok Call 2

r=54

G30 874 tok Call 1

G31 395 tok Call 2

r=32

G32 1773 tok Call 1

G33 962 tok Call 2

r=13

G34 1592 tok Call 1

G35 583 tok Call 2

Figure 8. Case III: two calls and 𝑅 = 780. With chunk_size= 64, the first call follows the long 𝐺 continuation chain while also executing three replay-started members; all side leaves become ready for one parallel second call. Call 1 (four members; replay 11). 𝑅3 → 𝐺1 → 𝐺2 (𝑟 = 0); 𝐺3 → 𝐺4 → 𝐺6 → · · · → 𝐺32 → 𝐺34 (5); 𝐻 1 → 𝐻 2 → 𝐻 5 (3); and 𝐼 1 → 𝐼 2 → 𝐼 5 (3).

Call 2 (22 singletons; replay 769). 𝐺35, 𝐺33, 𝐺31, 𝐺29 replay 13, 32, 54, 55; 𝐺27, 𝐺25, 𝐺23, 𝐺21 replay 42, 36, 53, 24; 𝐺19, 𝐺17, 𝐺15, 𝐺13 replay 17, 29, 38, 34; and 𝐺11, 𝐺9, 𝐺7, 𝐺5 replay 59, 55, 23, 14. The remaining 𝐻 6, 𝐻 3, 𝐻 4, 𝐼 6, 𝐼 3, 𝐼 4 replay 61, 60, 60, 4, 3, 3.

23

Related documents

• Adaptive Anisotropic Attention for Axis-Structured Signals#668051
• SOLAR: SVD-Optimized Lifelong Attention for Recommendation#1142634
• Divide-and-Remember: Recursive Action-Relevant Memory for Long-Horizon VLA Policies#1155917
• Remember When It Matters: Proactive Memory Agent for Long-Horizon Agents#353091
• PF-RL: Progress Field Reinforcement Learning via Goal-Conditioned Value Geometry for Vision-Language-Action Models#1116965