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× 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× and 1.57× at 1M tokens on 32 H800 GPUs.
Figures & tables
Figure 1: One training step of GLM-5.1 in Megatron-LM on 512 H20 GPUs (PP8 × CP32 × DP2 at 256K, PP8 × CP64 at 512K, PP4 × CP128 at 1M) with 8K tokens per GPU and full activation recomputation of every layer ( Korthikanti et al., 2023 ) , first pipeline stage. Step times exclude the pipeline wait, with the attention layer in color and the rest of the step in gray.
Figure 2: Execution of one learned sparse attention layer on one GPU in Megatron-LM under CP and full activation recomputation (not to scale), showing the materialized n×L score table ( \scriptsize1⃝ ) and the repeated selection in recomputation ( \scriptsize2⃝ ). Each row is one pass, with CP communication drawn below the computation.
Figure 3: A learned sparse attention layer in HyScale, drawn in the same lanes as Figure 2 (not to scale). Green numbers mark where each technique removes the inefficiency that Figure 2 marks in red. The forward pass selects while scoring with HyTopK, without writing the score table ( \scriptsize1⃝ ), the recomputation rescores only the saved keys ( \scriptsize2⃝ ), and, as a system optimization (§ 3.4 ), each CP collective runs behind computation that does not consume it.
Policy
Extra memory
Recomputation
No recomputation
all activations
0
Full recomputation
0
2nLHIdI+nL
Keep scores
4nL
nL
Keep top- k indices
4nk
2nLHIdI
Keep indexer inputs
2LdI+2nHIdI
2nLHIdI+nL
Index replay
4nk
2nkHIdI
Table 1: Costs of selection recomputation policies for one layer on one GPU, in bytes beyond the layer input and in operations rerun. The upper policies are operator-level, and the lower two are selection replay.
Figure 4: Costs of the policies in Table 1 for one GLM-5.1 layer on one GPU ( L = 1M, n = 8K, k = 2048, HI = 32, dI = 128), on log axes. Operator-level policies lie on the top or right edge, and both variants of selection replay lie in the lower-left corner.
Figure 5: Locality-aware sparse kernels. Existing kernels (a) assign each thread block a single query; HyScale (b) processes two adjacent queries in one thread block over the union of their selections. Each row in a block marks the keys that one query uses.
Figure 6: Full GLM-5.1 on 512 H20 GPUs from its released weights, with 8K tokens per GPU and 4M tokens per step; each sample is one sequence. The runs use PP8 × CP32 × DP2 × EP64 at 256K, PP8 × CP64 × EP64 at 512K, and PP4 × CP128 × EP128 at 1M. Both systems force balanced MoE routing and use momentum-free SGD with Megatron-LM’s distributed optimizer. Step times average steps 4 to 10.
Figure 7: Breakdown of one training step of the full GLM-5.1 at 1M, first pipeline stage.
Figure 8: CP scaling of one GLM-5.1 attention layer on H20. Panels (a) and (b) show strong scaling at 256K and 512K context, and (c) and (d) weak scaling at 8K and 16K queries per GPU. Layer time includes the forward pass, the recomputation, and the backward pass on the slowest rank.
Figure 9: Cumulative ablation of one GLM-5.1 layer on H20 at 8K queries per GPU.
Figure 10: Cumulative ablation of single CSA and QSA layers on 32 H800 GPUs.
Figure 11: End-to-end efficiency of HyTopK’s top- k selection for GLM-5.1 on one H20 GPU. Labels in (a) give HyTopK’s speedup over Megatron-LM with its default 1 GiB memory limit on the score table, and (b) is relative to HyTopK’s base selection time.
Figure 12: Sampling efficiency: share of failed queries at each buffer size. Labels count failures at HyTopK’s default 16K-slot buffer (red line).