AREAL-DTA: Dynamic Tree Attention for Efficient Reinforcement Learning of Large Language Models

TL;DR

Reinforcement learning post-training of large language models is hampered by repeated computation of long shared prefixes across rollout sequences. AReaL-DTA introduces a DFS-based dynamic tree attention that traverses a rollout prefix tree, reusing shared computations and materializing only a single root-to-leaf path at a time, with memory proportional to the longest sequence length. It also uses a load-balanced distributed batching strategy to partition and process multiple prefix trees across GPUs. The approach yields substantial throughput improvements (up to $8.31\times$ on a single worker and $6.20\times$ cluster-wide) and memory reductions, enabling scalable RL training without heavy activation checkpointing. Overall, the work demonstrates a practical path to efficient RL post-training of large LLMs by exploiting prefix sharing and dynamic scheduling.

Abstract

Reinforcement learning (RL) based post-training for large language models (LLMs) is computationally expensive, as it generates many rollout sequences that could frequently share long token prefixes. Existing RL frameworks usually process these sequences independently, repeatedly recomputing identical prefixes during forward and backward passes during policy model training, leading to substantial inefficiencies in computation and memory usage. Although prefix sharing naturally induces a tree structure over rollouts, prior tree-attention-based solutions rely on fully materialized attention masks and scale poorly in RL settings. In this paper, we introduce AREAL-DTA to efficiently exploit prefix sharing in RL training. AREAL-DTA employs a depth-first-search (DFS)-based execution strategy that dynamically traverses the rollout prefix tree during both forward and backward computation, materializing only a single root-to-leaf path at a time. To further improve scalability, AREAL-DTA incorporates a load-balanced distributed batching mechanism that dynamically constructs and processes prefix trees across multiple GPUs. Across the popular RL post-training workload, AREAL-DTA achieves up to $8.31\times$ in $τ^2$-bench higher training throughput.

Figures (8)

• Figure 1: A visual illustration of AReaL-DTA. AReaL-DTA orders sequences ($\mathbf{s}_1$–$\mathbf{s}_4$) as a prefix tree and maintains a stack of the active root-to-leaf prefix (tokens + KV cache) while pushing new segments and reusing shared prefixes (e.g., $\text{seg}_1$ and $\text{seg}_2$). At each leaf node, AReaL-DTA immediately computes the loss, backpropagates, and pops the branch --- discarding leaf-only activations (e.g., $\text{seg}_3$) while keeping shared-prefix states for the next sibling. Note that each node represents a set of tokens in sequence segments.
• Figure 2: Across training step, we present the reward comparison between AReaL-DTA and AReaL on the TauBench dataset and the 1.7B/8B models.
• Figure 3: Across accumulate RL training time, we present the reward comparison between AReaL-DTA and AReaL on the TauBench dataset and the 1.7B/8B models.
• Figure 4: Ablation study for the training throughput of AReaL-DTA on the $\tau^2$-Bench dataset and the 1.7B/4B/8B/14B models.
• Figure 5: Ablation study for the backward pass memory utilization of AReaL-DTA on the TauBench dataset and the 1.7B/4B/8B/14B models.