# HyScale: Scaling Long-Context Training for Learned Sparse Attention

> HyScale accelerates long-context sparse-attention training 2x by treating selection as discrete indices, enabling on-the-fly top-k and replay without recomputing scores.

- **Source:** [arXiv](https://arxiv.org/abs/2610.32057)
- **Published:** 2026-10-31
- **Permalink:** https://picx.dev/p/xkUY5g
- **Whiteboard:** https://picx.dev/p/xkUY5g/image

## Summary

## 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$):

| Policy | Extra memory | Recomputation |
|--------|-------------|---------------|
| No recomputation | all activations | 0 |
| Full recomputation | 0 | $2nLH_Id_I + nL$ |
| Keep scores | $4nL$ | $nL$ |
| Keep top-$k$ indices | $4nk$ | $2nLH_Id_I$ |
| Keep indexer inputs | $2Ld_I + 2nH_Id_I$ | $2nLH_Id_I + nL$ |
| **Index replay** | $4nk$ | $2nkH_Id_I$ |
| **Score replay** | $8nk$ | 0 |

### 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)

| Metric | 256K | 512K | 1M |
|--------|------|------|-----|
| **Step time speedup** | 1.23× | 1.44× | **2.07×** |
| **MFU (HyScale)** | 20.9% | — | 27.8% |
| **MFU (Megatron-LM)** | 17.1% | — | 13.5% |

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)

| Architecture | Speedup |
|-------------|---------|
| CSA (DeepSeek-V4-Flash) | **2.01×** |
| QSA (Qwen3.8-Flash-Next) | **1.57×** |

### 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

---

_Markdown view of https://picx.dev/p/xkUY5g, served by PicX — AI-generated visual whiteboard summaries of research papers._