HyScale speeds up training for long-context sparse attention AI models
HyScale: Scaling Long-Context Training for Learned Sparse Attention
Artificial IntelligenceDistributed, Parallel, and Cluster Computing
Summary
Training AI models to understand very long texts is slow because they have to pick important parts from a huge amount of previous information. The authors introduce HyScale, a system that makes this selection faster by not storing unnecessary full score tables and by reusing previous computations. This leads to over twice the speed for training certain large AI models on sequences with up to one million tokens. Their method significantly reduces the training bottleneck for models using learned sparse attention.
What this means in practice
- •For ai model engineers: Train large-scale long-context AI models faster by optimizing sparse attention selection steps in transformer architectures.
- •For high performance computing teams: Deploy more efficient GPU training pipelines for billion-token sequence models, reducing training time and resource use.
Authors
Yangsheng Deng, Xinjing Huang, Yulong Wang, Yunzhuo Liu, Jun Hu, Shuyong Tang, Junchi Zeng, Mingji Han, Ray Ying, Key Zhang, Luo Mai
Abstract
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 tackling end-to-end training 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.