HyScale：面向可学习稀疏注意力的长上下文训练扩展系统
HyScale: Scaling Long-Context Training for Learned Sparse Attention
Yangshen Deng, Xinjing Huang, Yulong Wang, Yunzhuo Liu, Jun Hu, Shuyong Tan, Chun-Chih Tseng, Mingji Han 等 11 位作者
AI 速读 基于摘要
HyScale 将可学习稀疏注意力中的 top-k 选择视为索引集合而非激活，通过 FlashTopK 与选择重放加速百万级 token 的长上下文训练。
cs.AI · cs.DC
查看原文摘要
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 FlashTopK and selection replay. FlashTopK 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 tackling end-to-end challenges, HyScale speeds up GLM-5.1 training on 1M-token sequences by $2.07\times$ over Megatron-LM on 512 H20 GPUs. It also speeds up the CSA layers of DeepSeek-V4-Flash and the QSA layers of Qwen3.8-Flash-Next by $2.01\times$ and $1.57\times$ at 1M tokens on 32 H800 GPUs.