Summary (Overview)

• HyScale is a long-context training system for learned sparse attention, addressing the training bottleneck of the selection step (scoring all previous tokens per query and selecting top-$k$).
• Core insight: The output of selection is a discrete set of indices, not a continuous activation. This enables two key techniques: HyTopK (on-the-fly top-$k$ selection without materializing a full score table) and selection replay (reusing forward-pass selection during recomputation instead of repeating it).
• Results: HyScale accelerates end-to-end GLM-5.1 training on 1M-token sequences by 2.07× over Megatron-LM on 512 H20 GPUs, with speedups growing with context length.
• Generality: HyScale generalizes to block-level sparse attention architectures, speeding up CSA layers of DeepSeek-V4-Flash by 2.01× and QSA layers of Qwen3.8-Flash-Next by 1.57× at 1M tokens on 32 H800 GPUs.
• HyScale is deployed in production, powering large-scale LLM training at million-token contexts.

Introduction and Theoretical Foundation

Background

Long context has become a core capability of large language models, driving training sequence lengths from 128K to 1M tokens. At these lengths, full attention dominates training cost due to quadratic computation growth.

Learned Sparse Attention

Learned sparse attention reduces this cost by splitting attention into two steps:

1. Selection: A lightweight indexer scores every previous token for each query and keeps the top-$k$.
2. Sparse attention: Each query attends only to the selected tokens.

The selection process is formalized as:

$$I_t = \text{Score}(q^I_t, k^I_u) \big|_{u \in U_t}, \quad T_t = \text{Top-}k(I_t)$$

where $q^I_t$ is the indexer's query, $k^I_u$ is the indexer's key of entry $u$, $I_t \in \mathbb{R}^{|U_t|}$ is the score vector, and $T_t \subseteq U_t$ is the index set of size $k$.

Sparse attention then attends only to selected entries:

$$o_t = \text{Attn}\left(q_t, \{c_u\}_{u \in T_t}\right)$$

The Selection Bottleneck

The paper identifies two key inefficiencies in existing systems:

1. Materialized scores: Current systems write an $n \times L$ score table to GPU memory before selecting, even though top-$k$ keeps only $k$ entries per query. At 1M tokens ($n=8$K, $k=2048$), the table takes 32 GiB per GPU, whereas selected indices take only 64 MiB.

2. Repeated selection: Under full activation recomputation, selection runs twice per training step—once in forward, once in recomputation.

Key Insight

The output of selection is not a continuous activation, but rather a discrete set of indices.

Two properties support this view:

• The score table serves only to find these indices; nothing reads the scores of non-selected keys.
• Top-$k$ is non-differentiable, so no gradient flows through selection; the backward pass needs only the selected indices and their scores.

Methodology

HyTopK: On-the-Fly Top-$k$ Selection

HyTopK uses optimistic thresholding—the threshold need not be a guaranteed lower bound, as long as failures are rare and cheap to recover from.

Three key components:

1. Optimistic strided sampling (C1): Samples $M = \min(\max(4096, L_t/128), 8192)$ evenly spaced keys from each query's own context, setting the threshold at the $\lceil 3kM/L_t \rceil$-th best sample.

2. Lightweight host-free recovery (C2): A check kernel marks failed queries (candidates overflowed or fewer than $k$), then recovery kernels rescore up to 64 failed queries over their whole context.

3. Deterministic selection (C3): Uses radix select for tied keys at the $k$-th position and radix sort to order selected entries by score, breaking ties by position.

Selection Replay

Selection replay is a fine-grained recomputation policy that records forward-pass selection results and replays only the $k$ selected entries during recomputation.

Two variants:

• Index replay: Keeps only $k$ indices per query, rescores only those entries (default).
• Score replay: Also keeps the $k$ scores, rescores nothing.

Cost comparison (per layer, one GPU at $L=1$M, $n=8$K, $k=2048$):

Holistic System Optimizations

1. Decoupled CP Overlap: Transfers indexer keys and KV states separately, overlapping each with computation that doesn't consume it.

2. Locality-aware sparse kernels: Process two adjacent queries in one thread block over the union of their selections (adjacent queries share ~71% of top-$k$ keys), reducing L2 traffic and atomic operations.

3. Implementation details: CP layout cache, sequence boundaries cache, kernel fusion, and permutation-free KV gather.

Empirical Validation / Results

End-to-End Training (GLM-5.1 on 512 H20 GPUs)

At 1M context, HyTopK cuts forward selection from 35.9s to 20.8s, and selection replay cuts recomputed selection from 37.0s to 1.0s.

Ablation Study (1M context, 128 GPUs)

• HyScale saves 66% of Megatron-LM's layer time
• HyTopK saves 37% (largest single contribution)
• Selection replay saves another 18%
• Implementation optimizations save the remaining 11%

Generality (Block-level architectures, 1M tokens, 32 H800 GPUs)

HyTopK Efficiency

• Peak memory: 1.5–1.9 GiB per layer (vs. 3.1–48 GiB for Megatron-LM)
• Deterministic sort adds only 0.8% overhead
• Host-free recovery: 1.8% overhead for 1 failed query, 4.3% for 64
• Optimistic strided sampling fails on 5,000× fewer queries than LiteTopK/DoSA at 16K buffer slots

Theoretical and Practical Implications

Theoretical Implications

• Rethinking activation semantics: The paper challenges the assumption that all intermediate computation outputs must be treated as continuous activations. Discrete index sets require fundamentally different treatment in training systems.
• Recomputation policies: Operator-level recomputation granularity is insufficient for non-differentiable operations; entry-level policies can achieve orders-of-magnitude improvements.
• Threshold-based selection: Guaranteed lower bounds are unnecessary; optimistic thresholds with cheap recovery are more efficient in practice.

Practical Implications

• Production deployment: HyScale is deployed in production, demonstrating practical viability at scale.
• Architecture generality: Works across token-level (DSA) and block-level (CSA, QSA) sparse attention designs.
• Memory efficiency: Eliminates the memory–time trade-off of score table materialization, enabling longer contexts on existing hardware.
• Scaling behavior: HyScale's advantage grows with context length, making it increasingly valuable as models move toward million-token contexts.

Conclusion

HyScale addresses the training bottleneck of learned sparse attention through the key insight that selection outputs indices, not activations. Its two core techniques—HyTopK (on-the-fly top-$k$ selection via optimistic thresholding) and selection replay (fine-grained recomputation policy)—eliminate the score table materialization and repeated selection costs, respectively.

The system achieves 2.07× end-to-end speedup for GLM-5.1 training at 1M tokens on 512 H20 GPUs, with speedups growing as context lengths increase. HyScale generalizes across diverse learned sparse attention architectures and is deployed in production.

Future directions implied by this work include:

• Extending these techniques to other non-differentiable operations in training
• Further optimization of sparse kernels for even longer contexts
• Adaptation to emerging sparse attention architectures as they evolve