发表机构
HKUST; University of Cambridge; Ant Group; Reflection AI（香港科技大学; 剑桥大学; 蚂蚁集团; Reflection AI）
机构由 AI 辅助整理，请以论文原文为准。AI 中文总结
针对强化学习后训练中不同内核的数值不一致问题，提出智能体框架KernelBraid，通过优化IR组织位级一致统一内核的搜索，在多硬件配置下实现显著性能提升且保证数值正确性。
AI 中文摘要
强化学习（RL）后训练阶段通常为 rollout（采样推演）和策略更新使用不同的GPU内核。在同步PPO和GRPO中，数值不一致会扰动当前token概率与rollout阶段分配概率之间的比值。使用策略更新后端重新计算rollout对数概率可避免这种差异，但会增加一次前向传播。当策略快照和概率处理与目标匹配时，位级一致的统一内核支持计算复用，其优化必须保证不同执行模式下的一致性。我们提出KernelBraid——一个从手工调优的位级一致实现起步的智能体框架。它的优化中间表示（IR）通过将实现与修改关联到数值需求、工作负载测量和推导历史，来组织源代码搜索。智能体协调修改并保留经验证的中间结果用于进一步探索；只有通过正确性检查且在单工作负载限制内提升总延迟的版本才会被保留。在H20上的12种端到端训练配置中，相对于带对数概率重计算的AReaL，KernelBraid实现了1.10倍的平均吞吐量，且训练奖励的平均比值约为1.00倍。单层剖析显示，在15组模型-GPU配对中，阶段总耗时的平均加速比为1.40倍。算子级评估覆盖了A100、H20和H200上的10个算子的正确性与性能，全部通过规定的位级检查。统一注意力搜索在700万大语言模型token的总工作负载延迟上，相比初始实现实现了2.52倍加速；消融实验评估了保留证据和分支探索对搜索效率与最终性能的贡献。我们的代码已在https://github.com/areal-project/AReaL-TIK开源。
英文摘要
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.
