THD Packing & MTP Speculative Decoding Tame 64K Code Agent RL on Ascend NPUs

Overview

AReaL-Ascend v1.0.5 completes verification of Qwen3.6-27B in a Code Agent scenario with 64K context length on Ascend NPUs. The release focuses on two optimization axes: training-side token efficiency via THD Sequence Packing and rollout-side latency reduction via MTP Speculative Decoding. The full stack includes CANN 9.0.1, PyTorch 2.10.0, torch_npu 2.10.0.post2, vLLM 0.23.0, vLLM-Ascend, Megatron-Core, MindSpeed, and Megatron-Bridge.

Training Scenario

The validation uses Qwen3.6-27B with Mini-SWE-Agent v2 and R2E Gym at 64K sequence length. This workload is closer to real agent workloads than math/reasoning benchmarks because rollout token counts are large and generation dominates the step time.

THD Sequence Packing: Eliminating Padding Waste

Problem

In Agentic RL, response lengths vary widely across trajectories. Traditional padded batches align to the longest sequence in the microbatch, causing massive padding tokens that waste Attention and MLP compute.

Mechanism

THD Sequence Packing compacts multiple valid token sequences into a tighter token layout. It uses First-Fit Decreasing (FFD) bin-packing to fill a microbatch, records segment boundaries via cu_seqlens, and applies a block-diagonal attention mask so segments remain invisible to each other.

Concrete Example

Four trajectories with lengths 27, 7, 13, and 4 tokens (52 effective tokens total) would occupy 27×4 = 108 slots in a padded batch → ~48% utilization. After packing, all 52 tokens fit in 52 slots → 100% utilization.

Implementation Paths

Text models: wrapper-owned THD packing (existing)

Qwen3-VL Dense/MoE + Megatron-Bridge: model-owned THD packing (new in v1.0.5) — the model merges visual embeddings internally before packing, and AReaL remaps outputs back to the unified packed layout.

PR: https://github.com/areal-project/AReaL/pull/1608

MTP Speculative Decoding: Accelerating Agent Rollout

Problem

Standard autoregressive decoding emits one token per step; serial steps = generated tokens. For Code Agent workloads with huge rollout token counts, this dominates end-to-end RL step time.

Mechanism

MTP (Multi-Token Prediction) heads predict multiple candidate tokens per step (e.g., 3 tokens). The main model verifies them; when all are accepted, fewer serial steps are needed. In the illustrated case, 2 steps produce 4 tokens instead of 4 steps.

RL-Specific Observability

AReaL adds: vLLMConfig.speculative_config passthrough to vLLM

Support for Qwen3.5/Qwen3.6 MTP rollout

MTP head loaded frozen at rollout server start, excluded from policy weight updates

Automatic collection of acceptance_rate and mean_accepted_len into training statistics

Commit: https://github.com/HwVanICI/AReaL/commit/fd35cc8f28a8ba9775691f50b203e9791e66a2d6

As RL training progresses, the policy drifts from the frozen MTP head, so acceptance rate may change. Tracking these metrics lets practitioners monitor actual speculative decoding benefit over time.

Async Train-Rollout Consistency: Decoupled PPO Loss

Problem

Asynchronous RL overlaps training and rollout for throughput. The rollout engine uses stale weights (behavior policy π_behave), while the training engine updates the policy (π_θ). This off-policy bias inflates importance-ratio variance, distorts advantage estimates, and can slow convergence or crash training. Synchronous blocking sacrifices throughput.

Solution

Decoupled PPO Loss (enabled via use_decoupled_loss=true) splits the importance ratio:

Recompute log-probs with current weights on received trajectories → π_proximal

Decompose importance weight:

π_θ / π_behave = (π_θ / π_proximal) × (π_proximal / π_behave)

First term (trainable update) goes into PPO clip constraint

Second term behave_imp_weight = π_proximal / π_behave treated as constant correction factor

This preserves async throughput while explicitly accounting for off-policy drift. behave_imp_weight also serves as a monitoring metric; values near 1.0 indicate good consistency.

Experimental Result

In the 64K Code Agent run, behave_imp_weight stabilized around 0.9995, confirming minimal policy mismatch despite async execution.

Training Results & Performance Breakdown

Reward: 0.56 → 0.94 over ~120 steps, showing clear stable convergence

Rollout time per step: ~2000 s

Training time per step: ~1300 s

Weight update time per step: ~25 s

The combination of THD packing (training efficiency), MTP decoding (rollout speed), and Decoupled PPO (consistency) allows all three — efficiency, consistency, convergence — to be maintained simultaneously.

Getting Started

Ascend NPU install guide: https://areal-ai.io/AReaL/en/tutorial/installation_npu.html

Docker images: ghcr.io/hwvanici/areal_npu:v1.0.5-a2 (A2), ghcr.io/hwvanici/areal_npu:v1.0.5-a3 (A3)

Source:

git clone https://github.com/areal-project/AReaL && cd AReaL && git checkout ascend-v1.0.5