Every Token Counts: Generalizing 16M Ultra-Long Context in Large Language Models

• HSA retrieves top-k past chunks using landmark representations, attends to each chunk separately, then fuses outputs using retrieval scores.This chunk-wise design integrates retrieval scores into the forward pass so they receive gradient updates during backpropagation.Source: §1 · p. 2
• HSA-UltraLong combines sliding-window attention with HSA; the study trains 0.5B dense and 8B-A1B MoE models on 8 trillion tokens.Training is followed by long-context extension and annealing; a 4K-token sliding window is used to preserve in-domain capability.Source: §1 · p. 2
• Effective length generalization requires chunk-wise attention, retrieval score–based fusion, and NoPE; all three are essential.Source: §1 · p. 3
• An overly large sliding window can weaken HSA’s learning of short-range dependencies and degrade long-range generalization.HSA generalizes by learning retrieval over short contexts and transferring that ability to longer contexts.Source: §1 · p. 3
• The effective context length in the training corpus strongly influences length extrapolation.Source: §1 · p. 3
• 16M-token extrapolation succeeds from a 32K pre-training window on trillion-token-scale data.The study analyzes in-domain and extrapolation performance across model size, attention-window size, and data scaling.Source: §1 · p. 3

2.1 Limitations of Chunk Selection in NSA

• NSA’s chunk selection does not always retrieve the most relevant chunk and misses perfect in-domain Multi-Query NIAH accuracy on RULER.Table 1 reports an NSA ablation at a 4K in-domain length.Source: §2.1 · p. 3, Table 1 · p. 3
• NSA’s in-context retrieval degrades rapidly as context grows, while NoPE supports extrapolation better than RoPE.Source: §2.1 · p. 3

2.2 Attention with Chunk Retrieval

• HSA jointly learns chunk selection and attention end-to-end, attending to each past chunk separately and fusing attention results via retrieval scores.Source: §2.2 · p. 4
• HSA splits sequences into fixed 64-token chunks with landmark representations, then selects the K most relevant chunks for each token.Source: §2.2 · p. 4
• Query-Key Normalization is important for HSA stability in practical trillion-token-scale training.Source: §2.2 · p. 4

3 Methodology

• SWA handles local retrieval and HSA handles global retrieval, with intermediate-layer KV cache shared across HSA modules as context memory.Source: §3 · p. 4

3.1 Model Architecture

• The model stacks a lower SWA decoder beneath grouped upper layers that begin with SWA+HSA and continue with SWA-only layers.Chunk summaries and KV cache are shared across HSA modules; the MoE design uses a dense first-layer MLP and one shared expert per MoE block.Source: §3.1 · pp. 4–5, Figure 3 · p. 5

3.2 Training

• Training begins with 512-token SWA warm-up and global HSA, then proceeds through 4K-SWA pretraining, 32K mid-training and annealing, and 8K supervised fine-tuning.The authors hypothesize that training from scratch with 4K SWA prevents meaningful short-context learning for HSA retrieval.Source: §3.2 · p. 5

4 Experiments

• The evaluation spans general, math, coding, and alignment tasks; MoE and dense models processed 8T and 4T tokens, followed by 175B long-text and 400B reasoning-data phases.The general-pretraining dataset totaled 10T tokens; the report also describes model-training hyperparameters and evaluation benchmarks.Source: §4 · pp. 5–6

4.2 Scaling and Evaluation of Pretrained and Fine-tuned Models

• At standard contexts, HSA-UltraLong-MoE matched TRM-MoE in average scores and, after SFT, averaged 1.3 points above Qwen3-1.7B (Non-thinking); the models extended effective context to 16M tokens.The dense variant was 3.3 points below Qwen 2.5-0.5B before SFT and approximately 4 points below Qwen3-0.6B after SFT, despite using a dataset one-ninth the size.Source: Table 3 · p. 7, Table 4 · p. 8, §4.2 · p. 8

4.1 Small-scale Preliminary Experiments

• Self-copy warm-up yielded the best length extrapolation, albeit with some negative impact on in-domain performance; full HSA + short SWA warm-up maintained in-domain performance while delivering reasonable length extrapolation.Initial direct training with such configuration had limited length extrapolation capabilities.Source: §4.1 · pp. 6–7

4.3 Long-context Evaluation

• Retrieval extrapolation improved with training data above 32K effective context; a 512-token SWA supported better HSA extrapolation than 4K, while the 8B MoE excelled over the 0.5B model on variable tracking, not pure retrieval.The comparison covers Single-NIAH, Multi-Query NIAH, and variable tracking across Figure 4 panels.Source: §4.3 · p. 9, Figure 4 · p. 9

4.4 Training/Inference Efficiency

• FlashAttention-3 led at shorter sequence lengths, while HSA gained an efficiency advantage only with longer contexts.The authors attribute HSA's short-sequence disadvantage to more memory accesses and FlashAttention-3's CUDA implementation better leveraging Hopper features.Source: §4.4 · pp. 9–10, Figure 5 · p. 10

5 Conclusion

• The remaining challenges are the HSA/SWA seesaw after short SFT, a 16:1 query-to-key–value head ratio that creates an information bottleneck, and no clear short-sequence efficiency advantage over FlashAttention-3.The authors point to kernel-level optimization as needed for the head-ratio constraint and short-sequence efficiency.Source: §5 · p. 10
• HSA attends chunk by chunk and fuses results via retrieval scores rather than selecting and concatenating chunks before attention.The conclusion describes this as the core insight behind the approach.Source: §5 · p. 10
• The authors frame their experiments as a step toward handling infinite-long context and advancing machine long-term memory research.The report also describes its experimental insights and open problems as a foundation for future ultra-long-context research.Source: §5 · p. 10