Source-linked AI summary
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, Ray Ying, Key Zhang, Luo Mai
TL;DR
At million-token contexts, learned sparse attention still faces a selection bottleneck because it scores every previous token and current systems materialize score tables and repeat selection. HyScale combines on-the-fly top-k selection, selection replay, and system optimizations to address this cost. It speeds up GLM-5.1 training by 2.07× at 1M tokens and also reports speedups for CSA and QSA layers.
Problem
Learned sparse attention still scores every previous token per query, and current systems materialize a full score table and repeat selection during recomputation.
Source: §1 · p. 1, §3.1 · p. 4Method
HyScale combines exact on-the-fly top-k selection without a score table, replay of forward-pass selections during recomputation, and system optimizations for learned sparse attention.
Source: §1 · p. 2Results
2.07× faster GLM-5.1 training on 1M-token sequences than Megatron-LM on 512 H20 GPUs; at 1M tokens on 32 H800 GPUs, CSA layers are 2.01× faster and QSA layers 1.57× faster.
Source: §1 · p. 2Takeaways & Limitations
HyScale supports token-level DSA and block-level CSA and QSA architectures for long-context training.
Source: §1 · p. 2, §6 · p. 11
Abstract
from arXiv · showhide
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.
1 Introduction
At million-token contexts, learned sparse attention reduces attention reads but leaves selection costly because it scores every previous token and existing systems materialize scores and repeat selection. HyScale addresses these bottlenecks with on-the-fly exact top-k selection, selection replay, and system optimizations, achieving substantial training speedups across token- and block-level architectures.
1 Introduction
- Learned sparse attention scores every previous token per query but reads only the selected top-k, with token-level and block-level variants used in production models.Block-level variants select contiguous token blocks.Source: §1 · p. 1
- At 1M tokens, selection consumes 47% of a Megatron-LM GLM-5.1 training step, excluding pipeline wait, and takes 6.6× as long as sparse attention.The cost grows because selection scores all previous tokens, materializes a full score table, and repeats selection during recomputation.Source: §1 · p. 1
- Selection outputs discrete indices rather than activations, so backward needs neither the intermediate score table nor a rerun of top-k search.Sparse attention uses selected indices, while indexer loss differentiates through only the selected scores.Source: §1 · p. 1
- HyTopK performs exact, deterministic top-k selection while scoring, avoiding the score table through optimistic thresholding with rare, cheap recovery and full context-parallel compatibility.Existing sampling-based top-k approaches can fail in training because sampling conflicts with context parallelism, thresholds can overflow fixed buffers, and selection can become nondeterministic.Source: §1 · pp. 1–2
- Selection replay records forward selection and replays only the k selected entries during recomputation, eliminating repeated selection with negligible memory overhead.HyScale integrates both techniques with decoupled context-parallel overlap and locality-aware sparse kernels, and supports DSA, CSA, and QSA.Source: §1 · p. 2
- 2.07× faster GLM-5.1 training on 1M-token sequences than Megatron-LM on 512 H20 GPUs; at 1M tokens on 32 H800 GPUs, CSA layers are 2.01× faster and QSA layers 1.57× faster.The CSA result is for DeepSeek-V4-Flash and the QSA result is for Qwen3.8-Flash-Next.Source: §1 · p. 2
2 Background and Motivation
Learned sparse attention selects a small set of entries for each query, but selection and its recomputation become major costs at long context lengths. HyScale’s key insight is that selection produces indices, not activations, so it need not retain or regenerate a full score table.
2.1 Learned Sparse Attention
- Selection scores every candidate entry and keeps the top k, while sparse attention reads only those selected entries.Per query, selection costs O(|U_t|), up to O(L), while sparse attention costs O(k), with k ≪ L.Source: §2.1 · p. 2
2.2 Long-Context Training
- Context parallelism splits a sequence across GPUs; activation recomputation trades compute for memory by rerunning layer operations during backward.With full recomputation, each layer keeps only its input and reruns its forward pass before backward.Source: §2.2 · p. 3
2.3 The Selection Bottleneck
- At 1M tokens, selection takes 47% of the step and 6.6× as long as sparse attention.The attention layer takes 49% of the step at 256K tokens and 61% at 1M, excluding pipeline wait.Source: §2.3 · p. 3
- At 1M tokens, materializing scores for one layer takes 32 GiB per GPU, although selected indices take only 64 MiB.With n=8K and k=2048, top-k keeps 0.2% of scores, so the full n × L table wastes memory and memory traffic.Source: §2.3 · p. 3
- Full activation recomputation repeats selection before backward, rescoring all L keys and selecting top-k again.At 1M tokens, this rerun accounts for half of the 47% of the step spent on selection.Source: §2.3 · p. 3
2.4 Key Insight
- Because the backward pass needs only selected indices and their scores, HyScale can find indices without a full score table and replay selection by rescoring only selected entries.The score table is used to find the indices; top-k is non-differentiable, so backward does not need that table or other selection intermediates.Source: §2.4 · p. 4
3 HyScale Design
HyScale combines on-the-fly top-k selection and fine-grained replay with communication and kernel optimizations for learned sparse attention. Its design targets both selection costs and the system overheads exposed when selection is accelerated.
3.4 Holistic System Implementation
- HyScale combines HyTopK and selection replay with system optimizations that overlap CP communication and improve sparse-kernel efficiency.It is built on Megatron-LM and supports DSA, CSA, and QSA.Source: §3.1 · p. 4, §3.4 · p. 6
3.2 HyTopK
- HyTopK estimates each query’s threshold from its own strided context sample, recovers failed queries on GPU, and restores deterministic exact top-k output.The recovery handles queries with too few candidates or buffer overflow without host synchronization.Source: §3.2 · p. 5
3.3 Selection Replay
- Selection replay stores each query’s forward-pass indices and rescores only their k entries, cutting recomputation from O(L) to O(k) per query.Index replay keeps indices; score replay also keeps the selected scores and therefore rescores nothing.Source: §3.3 · pp. 5–6
- At 1M tokens, index replay costs at most 1.6% of a training step and saves about 5 GiB per GPU versus score replay; operator-level policies cost at least 100× more in memory or recomputation.Both replay variants occupy the lower-left of the cost plot; index replay is the default.Source: Figure 4 · p. 6, §3.3 · p. 6
3.4.1 Decoupled CP Overlap for Selection and Attention
- Decoupled CP overlap transfers indexer keys and KV states separately, hiding each transfer behind computation that does not consume that data.The forward-pass and recomputation gathers overlap query projections and selection; backward reduce-scatters overlap other backward computation.Source: §3.4.1 · p. 7
3.4.2 Locality-Aware Sparse Kernels
- Adjacent GLM-5.1 queries share 71% of their top-2048 keys on average, so locality-aware kernels process each pair over a union and reuse shared-key loads and gradient updates.Each query masks out keys selected only by the other, preserving exact outputs and gradients.Source: §3.4.2 · p. 7, Figure 5 · p. 7
3.4.3 Implementation Details
- HyScale caches CP layouts and sequence boundaries per micro-batch, fuses elementwise operators, and adjusts KV indices instead of reordering gathered data.These techniques target host synchronizations and wasted GPU work.Source: §3.4.3 · pp. 7–8
4 Evaluation
HyScale improves long-context training performance across GLM-5.1 and block-level sparse-attention architectures, while its selection techniques account for most of the measured layer-time savings. The evaluation also examines context-parallel scaling, memory use, and selection reliability.
4.1 Setup
- The evaluation trains full GLM-5.1 from released weights on 512 H20 GPUs, with a 4M-token global batch and 8K tokens per GPU across 256K, 512K, and 1M-token sequences.The testbed includes H20 and H800 clusters.Source: §4.2 · p. 8, §4.1 · p. 8
4.2 End-to-End Training
- HyScale is up to 2.07× faster than Megatron-LM, with speedup increasing from 1.23× at 256K to 2.07× at 1M context, at about the same peak GPU memory.With the same 4M tokens per step, from 256K to 1M, HyScale’s step time grows 1.27× and MFU rises from 20.9% to 27.8%; Megatron-LM’s step time grows 2.13× and MFU falls from 17.1% to 13.5%.Source: §4.2 · p. 8
4.3 Context-Parallel Scalability
- At 512K context, HyScale’s speedup over Megatron-LM grows from 1.79× to 8.1× as GPU count increases; Megatron-LM stops improving as early as 128 GPUs and slows beyond 256 GPUs.Up to 256 GPUs, each doubling cuts HyScale layer time by 1.65–1.96×; beyond 256 GPUs, gains fall to 1.18–1.54× because each GPU holds too few queries to hide CP transfers.Source: §4.3 · p. 9
- With 16K queries per GPU, HyScale’s layer time grows 11× from 128K to 8M context versus Megatron-LM’s 190×, and HyScale is 22.3× faster at 8M.Layer time includes the forward pass, recomputation, and backward pass on the slowest rank.Source: §4.3 · p. 9, Figure 8 · p. 9
4.4 Ablation Study
- At 1M context, HyScale saves 66% of Megatron-LM’s layer time, with HyTopK and selection replay accounting for 84% of that saving.HyTopK saves 37% and selection replay a further 18%; implementation optimizations save the remaining 11%.Source: §4.4 · pp. 9–10, Figure 9 · p. 9
4.5 Generality
- For learned sparse block-level attention on 32 H800 GPUs, HyScale runs single CSA and QSA layers 2.01× and 1.57× faster, respectively, than Megatron-LM.Every ablation step speeds up both layers; HyTopK’s gain is smaller than on token-level DSA because block-level layers score 4× fewer keys per query.Source: §4.5 · p. 10, Figure 10 · p. 10
4.6 HyTopK Efficiency
- HyTopK uses 1.5–1.9 GiB per layer and is faster than Megatron-LM at every tested context, avoiding its score-table memory–time trade-off.With the default 1 GiB score-table limit, Megatron-LM is 1.22–1.71× slower; using the whole table requires 12–48 GiB and remains 3–10% slower.Source: §4.6 · p. 10, Figure 11 · p. 10
- At 16K slots, HyTopK’s optimistic threshold fails on over 5,000× fewer queries than fixed-threshold LiteTopK and DoSA, and on none with random weights.With random weights, HyTopK fits every query in 5.9K slots, versus 10.5K for LiteTopK and 12.3K for DoSA.Source: §4.6 · pp. 10–11, Figure 12 · p. 11
5 Related Work
Related work covers methods that accelerate top-k selection, approaches that reduce selection work through reuse or reduced scoring, and HyTopK’s distinctions for training.
5 Related Work
- HPC-Ops Top-K (Liu et al., 2026b), GVR (Cheng et al., 2026), DeepSelect, and Prof-K (Dziarmaga et al., 2026) accelerate top-k but still write score tables to GPU memory.LiteTopK (Yin et al., 2026) fuses scoring and selection using sampled thresholds, while DoSA (Meng & Fang, 2026) also skips some scoring through score bounds; both target inference and face challenges in training.Source: §5 · p. 11
- LiteTopK samples keys selected by nearby queries and DoSA uses a related sampling method, whereas HyTopK samples evenly from each query’s own context.The paper attributes the baselines’ higher failure rates to thresholds that must be guaranteed lower bounds and to weak top-k overlap among nearby queries at random initialization.Source: §5 · p. 11
- HyTopK targets training with per-query optimistic-threshold sampling, host-free recovery, and deterministic selection.PIVOT (Liu et al., 2026a), IndexCache (Bai et al., 2026), HISA (Xu et al., 2026), and MISA (Zhou et al., 2026) are orthogonal approaches that reuse selections or reduce scoring across queries, layers, blocks, or heads.Source: §5 · p. 11
6 Conclusion
HyScale addresses selection overhead in long-context learned sparse attention by exploiting that selection outputs indices rather than activations. It combines two selection techniques and reports faster GLM-5.1 training at 1M tokens.
6 Conclusion
- HyScale trains GLM-5.1 on 1M-token sequences 2.07× faster than Megatron-LM on 512 H20 GPUs.The system is also reported to generalize across learned sparse-attention architectures.Source: §6 · p. 11
- HyTopK selects top-k entries without writing a score table, while selection replay skips repeated selection during recomputation.Both techniques build on the observation that selection outputs indices rather than activations.Source: §6 · p. 11
- HyScale’s conclusion is that its approach applies across various learned sparse-attention architectures.Source: §6 · p. 11
A Backward Pass of Learned Sparse Attention
The appendix derives the backward pass for learned sparse attention and describes how selection, sparse-attention gradients, and context-parallel gradient aggregation use selected entries.
A Backward Pass of Learned Sparse Attention
- Selection replay recomputes only the k scores at selected entries because the backward pass needs those scores rather than scores for all L keys.The indexer’s inputs are detached, so gradients reach only its projection weights.Source: A Backward Pass of Learned Sparse… · p. 13
- Sparse-attention backward is restricted to the selected entries, and the detached target prevents the indexer loss from sending gradients into sparse attention.Source: A Backward Pass of Learned Sparse… · pp. 13–14
- Under context parallelism, GPUs all-gather keys and states, compute gradients for their local queries, then reduce-scatter partial sums to return gradients for the positions they own.Source: A Backward Pass of Learned Sparse… · p. 14