[论文] HyScale: Scaling Long-Context Training for Learned Sparse Attention

论文概要

研究领域: ML 作者: Yangsheng Deng, Xinjing Huang, Yulong Wang, Yunzhuo Liu, Jun Hu, Shuyong Tang, Junchi Zeng, Mingji Han, Ray Ying, Key Zhang, Luo Mai 发布时间: 2026-10-27 arXiv: 2610.32057

中文摘要

我们提出HyScale，一个用于学习稀疏注意力（learned sparse attention）的长上下文训练系统。学习稀疏注意力是百万token上下文的一种有前景的架构，自DeepSeek-V3.2以来已广泛用于生产。然而，选择操作逐渐成为训练瓶颈。对于每个查询，它需要对每个先前的token进行评分，而稀疏注意力仅读取k个；更糟糕的是，当前系统物化完整的评分表并在重计算中重复选择。

我们的关键洞察是：选择的输出不是激活，而是一组索引。基于此，HyScale引入了HyTopK和选择回放（selection replay）。HyTopK在没有评分表的情况下即时找到每个查询的精确top-k，而选择回放则在重计算中复用前向传播的选择。结合解决端到端训练挑战的整体系统优化，HyScale在512个H20 GPU上将GLM-5.1在100万token序列上的训练加速了2.07倍（相比Megatron-LM）。它还在32个H800 GPU上将DeepSeek-V4-Flash的CSA层和Qwen3.8-Flash-Next的QSA层在100万token处分别加速了2.01倍和1.57倍。

原文摘要

We present HyScale, a long-context training system for learned sparse attention. Learned sparse attention is a promising architecture for million-token contexts, widely adopted in production since DeepSeek-V3.2. Selection, however, grows into the training bottleneck. For each query, it scores every previous token whereas sparse attention reads only k; worse, current systems materialize a full score table and repeat selection during recomputation. Our key insight is that the output of selection is not an activation, but rather a set of indices. Building on it, HyScale introduces HyTopK and selection replay. HyTopK finds each query's exact top-k on the fly without a score table, while selection replay reuses forward-pass selection in recomputation. Together with holistic system optimizations...

#论文 #arXiv #ML #小凯