TRL AsyncGRPOTrainer 树注意力支持线性注意力,提速最高 3.46x
SergioPaniego · x · 2026-10-27
TRL 的 AsyncGRPOTrainer tree attention 现已支持线性注意力(linear attention)模型,最高加速 3.46 倍。
技术实现:
- 基于 FLA 的 chunked Triton 实现改造了两个 state-scan 内核:前向时每个分支继承父节点的最终状态,反向时父节点对子节点的梯度求和。
- 共享前缀只处理一次,其最终 recurrent state 传给各 continuation 并行执行;反向时各 continuation 的梯度再回加到共享前缀。
- 同时修正了 causal conv1d 层的正确 token 处理。
配图基准(1x H100,BF16,真实 AsyncGRPO step):Qwen3.5-0.8B 在 grpo-g1-64k 场景提速 3.84x、agent-fork-long 4.36x;Qwen3.5-2B 最高 4.24x,且部分场景 baseline 直接 OOM,而 tree packing 方案仍可运行。
「Infra」频道最新
- 网友自制 Runpod 模板:MOSS-Image 搭 ComfyUI 一键开跑 — TheTimster666 · 2026-10-27
- 德国 AI 基础设施主权状况引争议,机器人被叫来整理数据 — BlackHC · 2026-10-27
- Prisma 从 v6 升到 v7:Rust 引擎换 JS,单核吞吐翻倍至约 290 req/s — DanielLockyer · 2026-10-27
- 美股 500 强利润增 30%,Nvidia 与 Micron 贡献其中三分之一 — LexSokolin · 2026-10-27
- WSJ:Broadcom CEO 周末斡旋促成 Anthropic 630 亿美元芯片租赁债务协议 — rohanpaul_ai · 2026-10-27
- 数字主权已成云架构决策:AI 管道让数据主权更难守住 — DavidLinthicum · 2026-10-27