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).
Long contexts have become standard in pretrained LLMs, yet they remain expensive to run: prefill compute grows quadratically with sequence length, and every decode step re-reads a key-value cache that grows linearly with it. Sparse attention cuts these costs by attending only to a relevant subset of past tokens, but selecting that subset is itself expensive. We present SpotAttention, a lightweight selector that attaches to a frozen pretrained transformer and learns by KL distillation to estimate its attention distribution. The selector picks the top-K keys each query attends to, and because its estimate is a calibrated distribution, a dual top-p rule reads the per-query, per-layer budget directly from it. Across Qwen3 (dense, 4B-32B) and Qwen3.5 (hybrid linear/full attention, 4B-9B), SpotAttention matches dense accuracy at contexts up to 128K tokens, eight times the training length. Decode at L=128K runs 3.9x faster than FlashAttention and 1.8x faster than Twilight, the strongest training-free baseline. Quantizing the selector's K-cache to INT4 or FP4 microscale shrinks it 3.5x at no accuracy cost.
Scaling modern large language models (LLMs) to long contexts is limited by the quadratic computation cost, and poor length extrapolation of dense attention. Chunk-wise sparse attention offers a promising alternative, but all existing methods fall short of full attention because of their inaccurate chunk selection. We propose Hierarchical Landmark Sparse (HiLS) Attention, a chunk-wise sparse attention mechanism that learns chunk selection end-to-end under the language-modeling (LM) loss. HiLS factorizes attention hierarchically: each query performs attention independently with each retrieved chunk to extract chunk-specific information, and the resulting outputs are fused according to chunk retrieval scores. By incorporating retrieval scores into the forward attention computation, HiLS optimizes them directly with the LM loss, enabling end-to-end retrieval learning and native sparse training. Experimental results show that HiLS-Attention achieves performance comparable to, and in some cases better than, full attention at in-domain context lengths. Meanwhile, HiLS-Attention extrapolates more than 64× the training context length with 90% retrieval accuracy, far beyond full attention. Moreover, existing full-attention models can be converted to HiLS-Attention with lightweight continued pretraining, preserving in-domain performance while acquiring ultra-long-context extrapolation. Together with its sparse KV access and computation, HiLS-Attention breaks the usual efficiency-performance trade-off, enabling long-context LLMs that are both more efficient and more effective on general long-context tasks than their full-attention counterparts.
Xiang Hu, Xinyu Wei, Hao Gu +10
Tencent HY Team · ShanghaiTech University · The Hong Kong University of Science and Technology +1
Token-level sparse attention mechanisms, exemplified by DeepSeek Sparse Attention (DSA), achieve fine-grained key selection by scoring every historical key for each query through a lightweight indexer, then computing attention only on the selected subset. While the downstream sparse attention itself scales favorably, the indexer must still scan the entire prefix for every query, introducing an per-layer bottleneck that grows prohibitively with context length. We propose HISA (Hierarchical Indexed Sparse Attention), a plug-and-play replacement for the indexer that rewrites the search path from a flat token scan into a two-stage hierarchical procedure: (1) a block-level coarse filtering stage that scores pooled block representations to discard irrelevant regions, followed by (2) a token-level refinement stage that applies the original indexer exclusively within the retained candidate blocks. HISA preserves the identical token-level top-sparse pattern consumed by the downstream Sparse MLA operator and requires no additional training. On kernel-level benchmarks, HISA achieves up to speedup at 64K context. On Needle-in-a-Haystack and LongBench, we directly replace the indexer in DeepSeek-V3.2 and GLM-5 with our HISA indexer, without any finetuning. HISA closely matches the original DSA in quality, while substantially outperforming block-sparse baselines.