SearcharxivSearch

arXiv subjects

Peihua Bao

Publications and source records attributed to Peihua Bao.

2 recordsLinked to original sources

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

Agentic reinforcement learning (RL) often produces irregular rollout trees with shared histories. Training root-to-leaf trajectories independently recomputes these shared prefixes. Existing systems primarily target full-attention models and lack dense, differentiable hybrid-attention execution compatible with activation recomputation. We present HARTS (Hybrid-Attention RL over Tree Structures). HARTS jointly plans microbatches, data-parallel (DP) replica assignments, and microbatch-slot schedules using non-replay compact-token work after prefix compression. For chunkwise linear attention, a linear-time algorithm coordinates chunk-boundary state recovery and replay and produces the minimum number of sequential linear-attention calls under our packed execution model. HARTS preserves the chunkwise state partitioning of trajectory-wise training: it does not repeat projections, MLP/MoE computation, or final outputs, and performs only bounded state replay for numerical alignment. Per round, HARTS batches all branches into one packed call, propagates gradients through differentiable state handoffs, supports activation recomputation, and restores per-token log-probabilities. For deterministic, no-token-drop top-$k$ MoE routing, semantic multiplicities restore MoE-objective token weights and load statistics. 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\times$ forward/backward/gradient speedup with activation recomputation across multiple parallel configurations. Its numerical differences are comparable to baseline self-rerun variation, and its reward trend is similar to the baseline over the first 120 steps of $τ^3$-Bench training.

cs.LG

Stencil-Lifting: Hierarchical Recursive Lifting System for Extracting Summary of Stencil Kernel in Legacy Codes

We introduce Stencil-Lifting, a novel system for automatically converting stencil kernels written in low-level languages in legacy code into semantically equivalent Domain-Specific Language (DSL) implementations. Targeting the efficiency bottlenecks of existing verified lifting systems, Stencil-Lifting achieves scalable stencil kernel abstraction through two key innovations. First, we propose a hierarchical recursive lifting theory that represents stencil kernels, structured as nested loops, using invariant subgraphs, which are customized data dependency graphs that capture loop-carried computation and structural invariants. Each vertex in the invariant subgraph is associated with a predicate-based summary, encoding its computational semantics. By enforcing self-consistency across these summaries, Stencil-Lifting ensures the derivation of correct loop invariants and postconditions for nested loops, eliminating the need for external verification. Second, we develop a hierarchical recursive lifting algorithm that guarantees termination through a convergent recursive process, avoiding the inefficiencies of search-based synthesis. The algorithm efficiently derives the valid summaries of stencil kernels, and its completeness is formally proven. We evaluate Stencil-Lifting on diverse stencil benchmarks from two different suites and on four real-world applications. Experimental results demonstrate that Stencil-Lifting achieves 31.62$\times$ and 5.8$\times$ speedups compared to the state-of-the-art verified lifting systems STNG and Dexter, respectively, while maintaining full semantic equivalence. Our work significantly enhances the translation efficiency of low-level stencil kernels to DSL implementations, effectively bridging the gap between legacy optimization techniques and modern DSL-based paradigms.

cs.SE