PaperlayerView exact revision on arXiv

Source-linked AI summary

AReaL-TIK: Stateful Agentic Optimization of Unified RL Kernels through an Optimization IR

Ran Yan, Youhe Jiang, Jiayi Nie, Wenshuang Li, Yingqi Peng, Taiyi Wang, Tongkai Yang, Binhang Yuan

arXiv:2609.35140v1cs.DC

TL;DR

In synchronous RL post-training, numerical disagreement between rollout and policy-update kernels can perturb probability ratios, while recomputation adds a forward pass. AREAL-TIK uses an IR-guided agent to search for faster bitwise-consistent kernels. It reports 1.10× average throughput relative to AReaL with recomputation across 12 configurations, alongside a 2.52× attention-search speedup.

  • Problem

    Distinct rollout and policy-update kernels can disagree numerically and perturb the token-probability ratios used in RL training.

    Source: §1 · pp. 1, 3

  • Method

    AREAL-TIK starts from a bitwise-consistent implementation and uses an optimization IR linking source changes to numerical requirements, workload measurements, and derivation history.

    Source: §1 · p. 3

  • Results

    1.10× average throughput relative to AReaL with recomputation across 12 configurations; H20 unified-attention search achieves 2.52× speedup in summed workload latency over the shared starting implementation.

    Source: §1 · p. 3, §6.1 · p. 13

  • Takeaways & Limitations

    Numerical requirements, workload evidence, and verified alternatives support multistep exploration under bitwise-consistency constraints.

    Source: §7 · p. 17

    The agreement requirement compares the two paths within one implementation; coordinated changes may alter both paths’ outputs relative to a parent, and external-reference accuracy is checked separately.

    Source: §3.1 · p. 5

Abstract

from arXiv · show

Reinforcement learning (RL) post-training often uses distinct GPU kernels for rollout and policy update. In synchronous PPO and GRPO, numerical disagreement can perturb ratios between current token probabilities and those assigned during rollout. Recomputing rollout log-probabilities with the policy-update backend avoids this discrepancy but adds a forward pass. Bitwise-consistent unified kernels permit reuse when the policy snapshot and probability processing match the objective. Their optimization must preserve agreement across distinct execution regimes. We present KernelBraid, an agentic framework starting from a hand-tuned, bitwise-consistent implementation. Its optimization intermediate representation (IR) organizes source-code search by linking implementations and modifications to numerical requirements, workload measurements, and derivation history. The agent coordinates changes and retains verified intermediates for further exploration; promotion requires passing correctness checks and improving aggregate latency within per-workload limits. Across 12 end-to-end training configurations on H20, KernelBraid achieves 1.10x average throughput relative to AReaL with log-probability recomputation, and the mean training-reward ratio rounds to 1.00x. Isolated-layer profiling yields 1.40x average speedup in summed phase time across 15 model-GPU pairs. Operator-level evaluation covers correctness and performance for 10 operators on A100, H20, and H200, all passing the prescribed bitwise checks. Unified-attention search achieves 2.52x speedup in summed workload latency over the starting implementation using 7M LLM tokens; ablations assess the contributions of retained evidence and branch exploration to search efficiency and attained performance. Our code is open-sourced at https://github.com/areal-project/AReaL-TIK.

1 Introduction

Distinct rollout and policy-update kernels can disagree numerically, perturbing RL probability ratios and making recomputation costly. AREAL-TIK searches for faster bitwise-consistent kernels using an optimization IR that links implementations to numerical requirements, workload evidence, and derivation history.

Motivation

  • Numerical differences between rollout and policy-update kernels can perturb token-probability ratios, while recomputing the denominator with the update backend adds a forward pass.Reusing rollout log-probabilities is appropriate when the policy snapshot and probability processing match the objective.Source: §1 · pp. 1–2, Figure 1 · p. 2

Optimization challenges

  • Full-sequence and decode paths favor different GPU schedules, but changes to tiling, reduction order, or precision can alter outputs and break their required numerical agreement.Coupled source and layout choices can help one path while slowing the other, and useful improvements may require retaining intermediates that do not immediately improve aggregate performance.Source: §1 · p. 3

Approach

  • AREAL-TIK starts from a hand-tuned, bitwise-consistent implementation and uses an optimization IR linking source versions and modifications to numerical requirements, workload measurements, and parent–child history.The agent consults shared-source and numerical records when proposing path-specific or joint modifications; evaluation checks bitwise agreement and accuracy against external references.Source: §1 · p. 3

Search and evaluation

  • Verified alternatives can be retained for follow-up changes, while promotion requires reducing weighted workload latency and respecting every workload’s slowdown limit.The workflow covers attention, normalization, position encoding, routing, activation, and KV-cache kernels.Source: §1 · p. 3

Contributions and results

  • 1.10× average throughput relative to AReaL with recomputation across 12 configurations; H20 unified-attention search reaches 2.52× speedup in summed workload latency over the shared starting implementation.The attention search uses 7M LLM tokens.Source: §1 · p. 3

2 Related Work

Prior work spans RL training infrastructure, training–inference mismatch, and automated kernel optimization. AREAL-TIK focuses on jointly optimizing coupled rollout and policy-update paths while preserving bitwise consistency.

RL training infrastructure

  • RL frameworks coordinate distributed rollout and policy updates; AREAL-TIK complements them by jointly optimizing full-sequence and token-wise decode kernels while allowing distinct GPU schedules.AReaL overlaps phases asynchronously, while veRL supports flexible RL dataflows and model resharding [7] [8].Source: §2 · p. 4

Training–inference mismatch

  • Training–inference mismatch studies examine BF16 rounding, mismatch growth, and errors on low-probability tokens; proposed mitigations include importance sampling, adaptive learning rates, and deterministic inference [20] [21] [22] [35].AREAL-TIK instead targets performance optimization of already bitwise-consistent rollout and policy-update kernels.Source: §2 · p. 4

Agentic GPU kernel development

  • Kernel-generation systems automate GPU code generation and refinement, while AREAL-TIK evaluates coupled entry-point changes against bitwise consistency and a joint objective with per-workload slowdown limits [23] [25] [26].This joint evaluation addresses full-sequence and token-wise decode paths.Source: §2 · p. 4

Persistent agent memory

  • Prior agentic and evolutionary systems retain interaction history or program feedback [27] [44] [46]; AREAL-TIK’s IR links numerical requirements and workload measurements to implementations, follow-up opportunities, and derivation history.The linked records guide modifications and selection of verified parents, including intermediates, without training another optimization policy.Source: §2 · p. 4

3 Motivating Coupled Kernel Optimization

The section frames unified-kernel optimization as a coupled search: both execution paths must preserve numerical agreement, yet favor different schedules and can affect each other’s performance. Retaining verified intermediate implementations can make later improvements reachable.

3.1 Numerical Inconsistency

  • At a fixed policy snapshot, the two paths’ corresponding logits should agree bit for bit; this isolates numerical disagreement from genuine policy change.The comparison is within an implementation, and coordinated changes can alter both paths’ outputs relative to a parent; numerical accuracy against external references is checked separately.Source: §3.1 · pp. 4–5

3.2 Key Observations for Coupled Optimization

  • The two paths can favor different GPU strategies, but numerical dependencies and shared source code couple their optimization; an edit in one path may require a corresponding change in the other.Changing the KV tile from 128 to 64 reduced shared memory per block from approximately 56 KiB to 34 KiB and raised residency from four to six CTAs per SM, while requiring decode to process twice as many sequential KV blocks and updates.Source: §3.2 · p. 5
  • A shared layout change can affect both entry points: changing CuTe tiling order reduced transfers for a 32 KiB transposed-V tile from 32 operations of 1 KiB to two of 16 KiB.Both paths construct TMA descriptors and WGMMA operands from the shared layout; resource changes alone do not establish a speedup, so both paths’ correctness and performance must be measured.Source: §3.2 · p. 5
  • A shared CuTe tiling change reduced static TMA instructions in each kernel from 34 to 4, passed all 12 bitwise-consistency tests, and improved decode by up to 1.27× and full-sequence execution by up to 1.09×.Average gains were 1.15× for decode and 1.06× for full-sequence execution.Source: §3.2 · p. 5
  • Verified intermediate implementations can enable later improvements even when they do not immediately replace the current best.The H200 attention example retains a buffer-removal change whose revised layout enables later TMA loading for transposed V and WGMMA computation of PV.Source: §3.2 · pp. 5–6
  • The search must link numerical requirements, measurements for both paths, and implementation history so edits can respect correctness while preserving follow-up opportunities.The IR connects these records to guide exploration and promotion.Source: §3.2 · p. 6

4 Optimization IR

The optimization IR links implementation versions to numerical constraints, workload measurements, and derivation history. It defines the joint correctness and latency objective, records search evidence, and preserves alternatives for further exploration.

4.1 Optimization Problem

  • The consistency constraint compares policy-update logits with rollout-prefill logits at prompt positions and rollout-decode logits at response positions.The sequence-wide comparison includes final response-position logits, although scoring the recorded response does not require them.Source: §4.1 · pp. 6–7
  • The joint objective minimizes weighted workload latency subject to output consistency, per-workload latency limits, and prescribed correctness requirements.The fixed evaluation specification sets workload and objective criteria, correctness tests, baseline latencies, and maximum permitted slowdowns; promotion requires a verified, admissible implementation with strictly lower T(I).Source: §4.1 · p. 7

4.2 Records and Relationships

  • The IR represents search through linked implementation, opportunity, action, audit, and lineage-management records.Together, these records connect source changes and their motivations to correctness outcomes, measurements, branch status, and evidence-supported follow-up opportunities.Source: §4.2 · p. 7

4.3 Persistent Search State

  • The persistent search state retains a derivation graph, active verified implementations, archived alternatives, pending opportunities, and the current best.The graph preserves linked evidence and decisions, including failed offspring and previously explored opportunities.Source: §4.3 · p. 8
  • The stopping counter κt tracks consecutive audited offspring that fail to improve the current best; its threshold limits unproductive search rather than certifying convergence.The threshold is normally configured between 20 and 50, and records remain in derivation history after implementations leave the frontier.Source: §4.3 · p. 8

5 IR-Guided Kernel Search

AREAL-TIK organizes unified-kernel tuning as an evidence-linked search over implementations, modifications, and audit results. Its workflow separates verification, promotion, and branch exploration so verified intermediates can support later improvements.

5.1 Workflow and Initialization

  • Initialization starts search only after the bitwise-consistent root compiles, passes mandatory checks, is fully measured, and meets per-workload latency limits.The root and its audit evidence seed the initial search state and opportunities.Source: §5.1 · p. 8

5.2 Candidate Generation

  • Actions select a pending opportunity and verified parent, then make a full-sequence-only, decode-only, or joint source-code change.The agent records the change’s scope and rationale before producing a complete offspring.Source: §5.2 · pp. 9–10

5.3 Verification and Promotion

  • Audit checks bitwise agreement, external numerical accuracy, and required workloads; promotion requires full verification, every workload limit, and lower aggregate latency.Failed offspring remain recorded but cannot enter the verified frontier or serve as parents.Source: §5.3 · p. 10

5.4 Branch Exploration

  • Branch decisions are separate from promotion: the agent can continue from verified implementations, record compatible changes for a later merge, archive alternatives, or abandon falsified directions.Failed offspring cannot serve as parents, while archived verified implementations can be reactivated when new evidence supports further work.Source: §5.4 · p. 10

5.5 Attention Case Study

  • Removing I0’s temporary K buffer passed correctness checks and raised the residency limit from two to three CTAs per SM, but I1 was not promoted.The subsequent cache-hint variant passed all 12 bitwise tests yet reduced aggregate performance by 0.08% and was archived.Source: §5.5 · p. 11
  • 1,024-token prefill latency fell from approximately 0.21 to 0.10 ms when verified I1 enabled a follow-up optimization in I2.The follow-up used I1’s layout for TMA loading of transposed V and WGMMA computation of PV; I2 met promotion criteria.Source: §5.5 · p. 11

6 Evaluation

Evaluation covers end-to-end PPO and GRPO training, isolated-layer and operator performance, kernel-search comparisons, and ablations. Results measure throughput, reward, latency, and search efficiency under the specified workloads and correctness constraints.

6 Evaluation

  • The evaluation tests two models with PPO and GRPO across staleness values s ∈{0, 2, 4}, using 300 updates and three seeds.It also evaluates isolated layers across five models and three GPUs, and compares search procedures on unified attention.Source: §6.1 · p. 11, §6 · p. 11, §6.2 · p. 13

6.1 Q1: End-to-End PPO and GRPO Training

  • 1.10× average throughput versus AReaL with recomputation, improving 11 of 12 configurations; the average reward ratio rounds to 1.00×.The Qwen2.5-1.5B PPO configuration at s = 2 has a 0.955× reward ratio.Source: §6.1 · pp. 12–13, Figure 8 · p. 12, Figure 5 · p. 12

6.2 Q2: Isolated-Layer Performance Across Models

  • 1.40× average speedup across 15 GPU–model pairs, with gains in 13; the five H20 models average 1.15×.The summed phase time includes AReaL’s recomputation forward pass and is not wall-clock training throughput.Source: §6.2 · p. 13, Figure 9 · p. 13

6.3 Q3: Per-Operator Performance

  • Attention’s average prefill and decode speedups remain below one on all three GPUs, while Hidden RMSNorm and NeoX RoPE average 10.30× and 18.03× prefill speedups.NeoX RoPE averages 2.90× for decode; the operator checks cover their tested domains, not arbitrary shapes.Source: §6.3 · p. 14, Table 1 · p. 14, Figure 11 · p. 15

6.4 Q4: Kernel-Search Comparison

  • 2.52× speedup over the shared starting implementation; AREAL-TIK reaches 2.52×, versus 2.00× for SkyDiscover and 1.50× for OpenEvolve.It uses 7M tokens, compared with 8.5M for SkyDiscover and 10M for OpenEvolve.Source: §6.4 · pp. 15–16, Figure 12 · p. 15

6.5 Q5: Ablation Study

  • 2.00× versus 2.52× for the complete system when lineage management is removed; every ablation uses more than the complete system’s 7M tokens.Withholding numerical requirements uses the most tokens, 16M, despite unchanged correctness enforcement.Source: §6.5 · p. 16, Figure 13 · p. 16

7 Conclusion

AREAL-TIK uses an optimization IR to guide joint source-code search across rollout and policy-update kernels. The evaluation reports faster throughput and isolated-layer phase times, with comparable average reward and prescribed bitwise checks passed.

7 Conclusion

  • AREAL-TIK uses an optimization IR as persistent memory for joint source-code search across rollout and policy-update kernels.Numerical requirements, workload evidence, and verified alternatives support exploration under bitwise-consistency constraints.Source: §7 · p. 17

A Training Hyperparameters

The table lists training hyperparameters for Qwen2.5-1.5B and Qwen3-4B on GSM8K; unless noted, settings apply to both systems, algorithms, and staleness values.

A Training Hyperparameters

  • Table 2 presents training hyperparameters for Qwen2.5-1.5B and Qwen3-4B on GSM8K.Unless noted, settings apply to both systems, both algorithms, and all data-staleness values.Source: Table 2 · p. 20
Loading 2609.35140v1…