HyScale: Scaling Long-Context Training for Learned Sparse Attention
AuthorsYangshen Deng, Xinjing Huang, Yulong Wang, Yunzhuo Liu, Jun Hu, Shuyong Tan, Chun-Chih Tseng, Mingji Han, Ray Ying, Key Zhang, Luo Mai
AffiliationsUniversity of Edinburgh Tencent Hunyuan
HyScale makes training long-context sparse attention faster by avoiding large selection tables and reusing selection results during recomputation.
Key results
× faster than Megatron-LM at 1M-token context on NVIDIA H20 GPUs.
Seconds after selection replay, down from 37.0 seconds in the profiled run.
GiB per GPU per layer for the conventional score table at 1M tokens.
× over Megatron-LM on 32 NVIDIA H800 GPUs at 1M tokens.
× over Megatron-LM on 32 NVIDIA H800 GPUs at 1M tokens.
Time spent on selection in the Megatron-LM GLM-5.1 training profile.
What the paper found
HyScale tackles a new bottleneck in learned sparse attention: even when attention reads only selected tokens, the indexer scores every earlier token, and training systems can spend 47% of a step on selection. Its FlashTopK operator fuses scoring with exact top-k selection without materializing the full score table; selection replay saves the chosen indices and recomputes scores only for those entries. The system also overlaps context-parallel communication with independent computation and groups adjacent queries to reuse shared keys. On NVIDIA H20 GPUs, HyScale trains GLM-5.1 with 1M-token sequences 2.07× faster than Megatron-LM; selection replay reduces recomputed selection time from 37.0 seconds to 1.0 second in the profiled run. At that context length, the conventional score table would take 32 GiB per GPU per layer. HyScale also generalizes beyond DeepSeek-style token selection: on 32 NVIDIA H800 GPUs, it accelerates the CSA layers of DeepSeek-V4-Flash by 2.01× and the QSA layers of Qwen3.8-Flash-Next by 1.57×.
Original 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.
Read the original paperMore in Attention Mechanisms
Browse all 24 papers →Latent Sparse Attention: Towards Subquadratic Ultra-Long Context Attention
Ales Kubicek, Andrea Caciolai, Liangzu Peng, Sukrit Sumant, Benjamin Hoover, Fraser Greenlee, Ruairidh M. Battleday, Radim Urban, Torsten Hoefler, James C. R. Whittington
Latent Sparse Attention organizes keys into learned clusters so models can attend over million-token contexts more efficiently.
Subquadratic Attention with a Hierarchical Tree-Structured Indexer
Leon Lufkin, Cutter Dawes, Alexander Speicher, Kamesh Krishnamurthy, Tomás Figliolia, Beren Millidge
ZHA uses a hierarchical tree to find relevant keys efficiently, aiming to make very long-context attention practical.
Can Computation from Earlier Problems Help LLMs Solve New Ones?
Jipei He, Wenhui Tan, Xiaoyi Yu, Enver Sangineto, Fiorenzo Parascandolo, Rita Cucchiara, Ruihua Song
STAIR helps language models reuse useful computation from earlier questions, improving performance on later problems with only a tiny trainable module.