ICML（International Conference on Machine Learning，国际机器学习大会）是机器学习领域最具影响力和学术声誉的国际顶级会议之一，也是人工智能研究领域的重要学术交流平台。会议长期聚焦机器学习的基础理论、核心算法、系统架构与实际应用，涵盖深度学习、强化学习、机器学习系统以及跨学科应用等前沿方向。PKU-DAIR实验室的论文《DARTS: Distribution-Aware Active Rollout Trajectory Shaping for Accelerating LLM Reinforcement Learning》、《EchoAttention: Exploiting Token-Pair Redundancy and Frame-Block Similarity for Efficient Long Video Generation》和《SALE:Low-bit Estimation for Efficient Sparse Attention in Long-context LLM Prefilling》 被 ICML 2026 录用。

DARTS: Distribution-Aware Active Rollout Trajectory Shaping for Accelerating LLM Reinforcement Learning

作者：Yujie Wang, Siwei Chen, Longzan Luo, Xinyi Liu, Xupeng Miao, Fangcheng Fu, Bin Cui

强化学习（Reinforcement Learining）已经成为训练高性能大语言模型的重要范式。与传统监督学习不同，LLM RL 通常需要模型针对每个 prompt 生成多条响应轨迹，并基于奖励信号更新策略。整个流程主要包括 rollout 与 training 两个阶段，其中 rollout 阶段需要大量采样，往往成为训练系统的主要瓶颈。

LLM rollout 的响应长度呈现严重长尾分布：少数 prompt 会生成远长于平均水平的轨迹，使得同步 RL 系统必须等待最慢样本完成，导致批次阻塞和 GPU 利用率下降。现有方法多采用 prompt 级别的长尾调度，例如将未完成的长轨迹延后处理或截断续跑。这类方法能够缓解等待问题，但本质上仍是在“绕开长尾”，并未改变模型自身的生成长度分布。

长尾不仅存在于不同 prompt 之间，也广泛存在于同一个 prompt 的多条 rollout 内部，即 intra-prompt long-tail distribution。这说明长尾并不只是 prompt 难度差异带来的现象，而是模型生成分布本身的固有问题。更重要的是，其中相当一部分长轨迹并不带来更好的奖励或推理质量，而是冗长、低效甚至错误的生成过程。

图1：大语言模型强化学习Rollout阶段轨迹长度的长尾分布

因此，如何从源头上减少无效长尾，同时保留必要的深度推理轨迹，成为高效 LLM RL 训练亟需解决的问题。

图2：DARTS 方法示意图：针对不同 prompt 的自适应采样策略与冗余分配

我们提出 DARTS，一个面向 LLM 强化学习训练的高效 rollout 框架。其核心思想是 主动分布塑形（active distribution shaping）：不再仅仅等待或调度长尾轨迹，而是通过分布感知的采样与资源分配机制，将模型 rollout 分布逐步塑造成更加简洁、稳定且高效的形式。DARTS 主要包含三个关键组件。

1. 分布感知轨迹采样

DARTS 首先为每个 prompt 构造更大的 intra-prompt 候选轨迹池，即通过冗余 rollout 生成多于训练所需数量的响应。随后，DARTS 采用 dual-end length sampling：一方面选择最短的一部分轨迹，以鼓励模型学习简洁、直接的解答；另一方面保留少量最长但有效的轨迹，以避免过度压缩推理过程，保留必要的深度思考能力。

这种设计对应两类典型情形：对于“冗长且无效”的长尾，DARTS 会增强其训练影响，引导模型更简洁；对于“复杂且必要”的长推理，DARTS 仍会保留有效长轨迹，使模型不会因为追求短输出而损害能力。

2. 自适应冗余分配

并不是所有 prompt 都需要同等规模的冗余采样。对于响应长度分布稳定、模型较有把握的问题，过多采样会浪费计算；而对于高方差、长尾严重的问题，更大的候选空间有助于更有效地塑形分布。

因此，DARTS 提出 variance-based adaptive redundancy allocation，利用历史响应长度方差来衡量 prompt 的长尾程度与模型不确定性，并将更多采样预算分配给高方差 prompt。这样，系统能够在有限计算预算下，把资源投入到最需要分布塑形的样本上，从而同时提升算法效果与系统效率。

图3：轨迹响应长度方差越高，长尾特征越明显

3. 系统级优化

在系统实现上，DARTS 进一步引入了两类优化。第一是 variance-guided tail pruning：当某些 prompt 被识别为极端长尾时，系统切换为 shortest-only sampling，并结合主动 early stopping，在收集到足够有效轨迹后提前终止剩余超长生成，避免极端尾部拖慢整体训练。

第二是 token-level streaming：不同于传统 sample-level streaming 必须等待整条轨迹生成完毕后再送入训练，DARTS 将生成结果按 token chunk 细粒度流式传输到训练端，使 rollout 与训练能够更充分重叠，进一步减少长尾场景下的等待时间。

#3 实 验

我们在64张 NVIDIA H20 96GB集群上进行实验。我们使用Qwen2.5-3B、Qwen2.5-Math-7B、Qwen2.5-14B、Qwen2.5-32B 以及 Qwen3-30B-A3B，并在 DAPO-MATH 与 MATH-lighteval 数据集上进行了广泛的实验。结果表明，DARTS 在不同模型规模下均稳定优于现有系统。相较于 VeRL，DARTS 实现了 1.29×–1.77× 的吞吐量提升；相较于代表性 prompt-level tail scheduling 方法 Tail Batching，最高也可达到 1.43× 的加速。实验说明，DARTS 不仅能缓解系统等待，更能通过主动塑形 rollout 分布，从源头上降低长尾开销。

图4：不同模型规模下的端到端吞吐量加速对比

Case study表明，DARTS 能显著压缩rollout 长度分布，使其集中到更短、更紧凑的区域；对于冗长无效型 prompt，DARTS 可以有效消除 verbose tails；对于需要复杂推理的 prompt，DARTS 仍能保留正确长轨迹所需的推理深度。

图5：DARTS通过对rollout轨迹响应长度的主动塑造实现强化学习的有效加速

#4 总 结

本文提出了 DARTS，一个通过 Distribution-Aware Active Rollout Trajectory Shaping 加速大语言模型强化学习训练的新框架。不同于已有方法主要从 prompt 级别调度长尾，DARTS 从 rollout 分布本身出发，识别并优化 intra-prompt 长尾问题，通过分布感知轨迹采样、自适应冗余分配以及系统级流式优化，引导模型生成更加简洁、稳定且高效的响应轨迹。

实验结果表明，DARTS 在多种模型规模和数据集上均取得显著加速，最高实现 1.77× 的端到端吞吐量提升，同时保持模型收敛与下游任务表现不下降。DARTS 为大语言模型强化学习训练提供了一种新的系统优化思路：不仅要调度长尾，更要主动塑造产生长尾的分布本身。

EchoAttention: Exploiting Token-Pair Redundancy and Frame-Block Similarity for Efficient Long Video Generation

作者：Yifei Xia, Fangcheng Fu, Hao Yuan, Suhan Ling, Xupeng Miao, Huixia Li, Yuxi Ren, Xin Xia, Xuefeng Xiao, Bin Cui

近年来，基于扩散 Transformer（DiT）的视频生成模型（如 Wan2.1、CogVideoX 等）发展迅猛，但长视频推理的计算瓶颈始终制约其实用化：3D 全注意力（Full Attention）的二次方计算开销 O(n²ℓ²d) 可占据整体推理算力的 60%～82%。现有稀疏注意力方法通过剪枝 Query-Key 对来缓解这一瓶颈，但其效果受制于大量"非稀疏 Head"——这些 Head 的注意力质量无法被有效稀疏化，导致纯稀疏方法存在难以逾越的速度—质量天花板。

图1：Token 对冗余性（稀疏注意力）与帧块相似性（Echo 算子）的对比示意

为突破上述瓶颈，本文发现了视频 DiT 中一种此前被忽视的规律——帧块相似性（Frame-Block Similarity）：注意力权重矩阵中同一块行（或块列）内的各帧块分布高度相似，且差异可由轻量对角线性标定精确刻画（图2）。如图1右侧所示，只需为每行计算一个原型块，其余块通过对角矩阵标定恢复，即可将主导计算复杂度从 O(n²ℓ²) 降至 O(nℓ²)，为非稀疏 Head 的加速提供了全新手段。

图2：Wan2.1-1.3B 中帧块相似性的可视化。同一块行/列内各帧块分布高度一致，行方向 Pearson 相关系数接近 1、rMSE 普遍低于 0.1，说明轻量线性标定即可精确重建

图3：EchoAttention 整体架构与三阶段蒸馏流程。（a）双算子架构及 Echo-Row 三步算法；（b）软混合 → 直通估计 → 硬路由的三阶段训练策略

基于上述观察，本文提出 EchoAttention（图3），一个双算子注意力加速框架。它包含 Sparse 算子（处理高稀疏 Head）与 Echo 算子（Echo-Row / Echo-Col，处理帧块相似 Head），并通过两级可学习门控路由器，在去噪时间步—层—Head 三个维度上自动选择最优算子。为使路由稳定收敛，本文采用"软混合→直通估计→硬路由"三阶段蒸馏策略，配合自定义 Triton 算子实现真正的端到端推理加速。

本文在 Wan2.1-1.3B 和 CogVideoX1.5-5B 两个代表性视频 DiT 上与 VSA、SLA、SpargeAttn 等最优基线进行全面对比，EchoAttention 在 Wan2.1-1.3B 上实现 1.97× 端到端加速，在 CogVideoX1.5-5B 上实现 2.42× 加速，均为所有方法中最高，且 VBench 分数与全注意力持平、PSNR/SSIM 指标显著优于所有稀疏基线，验证了其在突破速度—质量天花板方面的有效性。

SALE:Low-bit Estimation for Efficient Sparse Attention in Long-context LLM Prefilling

作者：Xiaodong Ji，Hailin Zhang，Fangcheng Fu，Bin CUI

论文链接：https://arxiv.org/abs/2505.24179

#1 背景与挑战

随着大语言模型（LLMs）在长文档问答、长篇摘要和仓库级代码补全等复杂任务中的广泛应用，推理时上下文窗口正在被不断拉长。然而，LLM 的自注意力模块计算复杂度与序列长度呈二次关系，在预填充阶段（prefilling）成为了长上下文推理的主要时延瓶颈。

已有研究发现，LLM 的注意力图（Attention map）天然稀疏，只有少量区域对输出产生重要贡献。基于这一观察，学术界提出了一系列稀疏注意力（Sparse attention）方法，通过构造稀疏掩码跳过不重要区域的计算来加速prefilling。然而，如图1中(a) (b)两个实例所示，现有动态稀疏注意力方法在"精度—效率权衡"上仍不理想：它们对注意力图的检查方式要么过于粗糙，要么不够全面，导致无法准确识别重要的注意力区域。

图1：不同稀疏注意力方法对注意力图的检查方式对比。

本质上，稀疏注意力方法在注意力图的"检查精度"与"额外算力开销"之间存在固有的trade-off：越细粒度、越全覆盖的检查越能精准识别出真正重要的区域，但需要付出更高的额外算力；而为了压低额外开销，现有方法不得不退化为采样或粗粒度检查这类近似手段，代价是漏检或误判重要元素。如何在不显著增加开销的前提下实现细粒度、全覆盖的注意力图检查，是提升稀疏注意力精度—效率权衡的关键。

#2 方法

为解决这一问题，我们提出了 SALE（Sparse Attention via Low-bit Estimation），一种基于低比特估计的块稀疏注意力方法，能够在几乎不损失模型精度的前提下，显著加速 LLM 长上下文预填充。SALE的核心设计是：用极低比特量化的 query-key 乘积快速近似整张注意力图，再基于这一近似图做细粒度的重要性筛选。图2展示了SALE的整体工作流。

图2：SALE 的整体工作流

在 SALE 的注意力图检查阶段，我们将 Q、K 量化至 4-bit，并在此基础上计算注意力图上每个位置的近似注意力权重。该步骤利用了现代 GPU 上低比特 Tensor Core 指令的高吞吐特性，最高能达到FP16 Tensor Core指令吞吐量的8倍。

得到注意力权重之后，SALE 提出使用 Relative Attention Score（相对注意力分数）来评估每个 query-key 对的重要性。该指标源自于被大量研究证实的现象（如图3所示）：LLM每行注意力权重在 sink（序列开头）和 local（序列末尾）区域普遍偏高。基于此，SALE以当前位置权重相对于 sink-local 区域权重的比值作为重要性判据。相较于直接使用原始注意力分数进行筛选，该指标无需将整张注意力图存储到GPU显存，计算/访存开销都更小。筛选结果以块稀疏掩码的形式传递给随后的Attention计算阶段。

图3：注意力权重分布示意图

考虑到不同注意力头之间的稀疏率存在显著差异，SALE 针对关键超参数进行了逐头离线校准，使每个头都能匹配到合适的稀疏度，所得校准结果可在推理阶段直接复用；此外，SALE 为注意力图检查阶段开发了高性能的 GPU kernel，将该阶段的开销压缩至全量注意力计算的11%左右。

#3 实验

图4：SALE 与现有稀疏注意力方法的精度—加速比权衡对比

图5：SALE 与现有稀疏注意力方法的精度—加速比权衡对比。横轴为延迟，纵轴为精度得分

如图4所示，SALE 在处理 64K 以上长序列时相对全量注意力可获得至少 3.36× 的加速，同时精度损失可忽略不计；并且如图5所示，SALE在"精度—效率"权衡维度上全面优于现有方法（例如MInference，FlexPrefill和SpargeAttn），证明了低比特估计 + 细粒度检查这一设计范式的有效性。

#4 总结

本文提出一种新颖的动态稀疏注意力方法SALE，用于解决长上下文大语言模型预填充阶段稀疏注意力的精度-效率权衡问题。该方法将查询（Q）和键（K）量化至4比特以快速近似注意力图，利用低比特计算的高吞吐特性降低开销；采用相对注意力分数进行细粒度重要区域筛选，结合逐头离线校准与高性能GPU内核优化。实验表明，该方法在几乎不损失模型精度的前提下，实现了长序列预填充推理的显著加速。

实验室简介

北京大学数据与智能实验室（Data And Intelligence Research Lab at Peking Univeristy，PKU-DAIR实验室）由北京大学计算机学院崔斌教授领导，长期从事数据库系统、大数据管理与分析、人工智能等领域的前沿研究，在理论和技术创新以及系统研发上取得多项成果，已在国际顶级学术会议和期刊发表学术论文200余篇，发布多个开源项目。课题组同学曾数十次获得包括CCF优博、ACM中国优博、北大优博、微软学者、苹果奖学金、谷歌奖学金等荣誉。PKU-DAIR实验室持续与工业界展开卓有成效的合作，与腾讯、阿里巴巴、苹果、微软、百度、快手、中兴通讯等多家知名企业开展项目合作和前沿探索，解决实际问题，进行科研成果的转化落地。