---
title: 'HSA-UltraLong: Ultra-Long Context LLM'
url: https://www.emergentmind.com/topics/hsa-ultralong-model
type: topic
---

# HSA-UltraLong: Ultra-Long Context LLM

HSA-UltraLong is an 8B-parameter Mixture-of-Experts (MoE) decoder architecture for Large Language Models (LLMs), integrating the Hierarchical Sparse Attention (HSA) mechanism to yield efficient ultra-long context modeling up to 16 million tokens. Developed as a foundation for “Machines that Can Remember,” HSA-UltraLong achieves near-linear compute scaling with respect to context length while preserving strong in-domain performance and high-accuracy retrieval from ultra-long contexts [2511.23319].

## 1. Hierarchical Sparse Attention (HSA) Mechanism

The HSA mechanism supersedes conventional quadratic attention by enforcing a two-level sparsity pattern, addressing the core requirements of ultra-long context modeling: sparsity, random-access flexibility, and length generalization. Contexts of length $L$ are partitioned into $M = L/S$ non-overlapping fixed-size chunks of $S$ tokens. Each chunk $H_{[i]}$ is associated with:

- A learned landmark key vector $K^{\mathrm{slc}_i} \in \mathbb{R}^d$,
- Per-chunk KV cache $(K_{[i]}, V_{[i]})$ of shape $S \times h \times d_h$ (with $h$ heads and $d_h = d/h$).

For each decoding position $t$, a retrieval query $Q^{\mathrm{slc}_t} \in \mathbb{R}^d$ computes scores $s_{t,i}$ by dot product with the landmark keys of all preceding chunks, normalized by $\sqrt{d}$. The top-$k$ scoring chunk indices $\mathcal{I}_t$ are selected via $\mathrm{TopK}$, providing random-access selection.

Intra-chunk multi-head attention is performed between $Q_t^{\mathrm{attn}}$ and the KV cache of each selected chunk. The outputs $\bar{O}_{t,i}$ are fused by retrieval score-based weights $w_{t,i} = \exp(s_{t,i}) / \sum_{j \in \mathcal{I}_t} \exp(s_{t,j})$, yielding the final token output as $O_t = \sum_{i \in \mathcal{I}_t} w_{t,i} \bar{O}_{t,i}$.

This construction ensures that each token attends only to $k \ll M$ chunks (sparsity), can randomly access arbitrary historical information (random-access flexibility), and, with NoPE positional encoding, generalizes retrieval behavior across context lengths during inference (length generalization).

Computational complexity is reduced from $O(L^2 d)$ (traditional attention) to
$O(L d + k S d)$ per token, which results in near-linear scaling for $L \gg kS$.

## 2. Model Architecture and Integration

HSA-UltraLong is built as an $L$-layer decoder divided into two sequential halves:

1. **Lower decoder ($L/2$ layers):** Standard Transformer blocks with Sliding-Window Attention (SWA), using window size up to 4K tokens and Rotary Positional Embedding (RoPE). MLPs are dense.
2. **Upper decoder ($L/2$ layers):** Organized into $G$ groups, each consisting of:
   - One hybrid layer with both SWA and an HSA module (HSA operates with NoPE),
   - Followed by additional SWA-only layers.

At the midpoint ($l = L/2$), chunk hidden states $H^l_{[i]}$ are bi-encoded (with prepended [CLS]) to produce embeddings $E_{[i]}$ and landmark vectors $L_i$, reused as chunk representations in all HSA layers.

MLP blocks in layers $2$ and higher are replaced by MoE blocks following the Ling-2.0 / DeepSeek-V3 protocol, with 64 experts per layer, 4 experts activated per token (top-4 gating), halved expert dimension to maintain overall parameter budget ($\approx$8B, with only 1B active per token). A shared “global expert” operates across all tokens to stabilize training. Token-expert allocation is balanced by an auxiliary, training-free loss function.

## 3. Training Regimen and Data Pipeline

Training proceeds in multi-phase curriculum:

- **Phase 1: Pretraining on 10T tokens (multi-domain, deduplicated)**
  - Web 50%, Code 14.4%, Math 12%, Code-NLP 5.6%, Reasoning 5%, Multilingual 4%, Books 2%, Wikipedia 1.5%, Others 5.5%
  - Dense variant: 4T tokens, MoE variant: 8T tokens
- **Phase 2: Long-text mid-training**
  - 175B tokens (sequence length 32K)
- **Phase 3: Reasoning-focused training**
  - 400B tokens

Supervised fine-tuning (SFT) uses 8K context data per Wu et al. (2025).

Curricular schedule:

1. **Warm-up:** 16K windows, SWA=512, large HSA top-$k$; synthetic RULER retrieval tasks introduced until reliable in-window retrieval.
2. **LM pretraining:** 16K window, SWA=4K, reduced HSA top-$k$.
3. **Long-context adaptation:** 32K window, increased HSA top-$k$ for improved long-range retrieval adaptation.
4. **Annealing:** on high-quality data, 32K window.
5. **SFT:** 8K window.

Optimizers employ AdamW ($\beta_1=0.9$, $\beta_2=0.95$, weight decay = 0.01, $\|\nabla\|_\infty = 1.0$), with learning rates and batch sizes set separately for MoE and dense pretraining. LR follows a schedule: linear warm-up, constant, cosine decay (SFT).

## 4. Empirical Evaluation and Key Results

Performance is validated on an array of in-domain and ultra-long context benchmarks:

- **In-domain ($\leq$32K):** MMLU, CMMLU, C-Eval, ARC-C, AGIEval, HellaSwag, PIQA, BBH, GSM8K, MATH, CMATH, MATH-500, OlympiadBench, HumanEval, HumanEval+, MBPP, MBPP+, CRUX-O, IFEval
- **Ultra-long (up to 16M):** Needle-in-a-Haystack variants, Variable Tracking in RULER

### Selected Quantitative Results

| Model Variant         | In-domain AVG (%) | Ultra-long MQ-NIAH @16M (%) | RULER Variable Tracking @16M (%) |
|----------------------|------------------|-----------------------------|----------------------------------|
| HSA-UL-MoE 8B/1B     | 63.1             | >95                         | ~85                              |
| TRM-MoE baseline     | 56.6             | N/A                         | N/A                              |
| Qwen2.5/0.5B         | 41.1             | N/A                         | N/A                              |
| Qwen3/0.6B           | 48.4             | N/A                         | N/A                              |
| Dense-0.5B           | N/A              | 0 (@1M MQ-NIAH)             | ~50                              |

In small-scale experiments (0.5B PG19), SWA+HSA with self-copy warm-up closes the gap in last-4K perplexity (16.5 vs 16.8 full-attention baseline) and yields MQ-NIAH accuracy of 93% at both 64K and 1M, compared to 5%/0% for standard BaseLM.

Ultra-long retrieval performance exceeds 90% on Single-NIAH and >95% on Multi-Query NIAH at 16M tokens, maintained at scale after mid-training and annealing. Variable Tracking shows HSA-UL-MoE sustaining ~85% at 16M tokens, while dense models drop to ~50%.

On computational efficiency, HSA kernels surpass FlashAttention-3 in speed for contexts $\gg$32K, though lag at contexts <8K.

## 5. Empirical Insights and Open Challenges

Empirical findings indicate that initial warm-up is critical to teaching HSA short-range retrieval prior to training SWA on large windows. A notable “seesaw” interaction arises: excessively wide SWA windows attenuate gradient signal for HSA retrieval, impeding extrapolation to longer contexts. Effective training demands early inclusion of long contiguous passages (>32K) in data.

Model scaling improves the reasoning+retrieval composite, with MoE-8B outperforming smaller dense models.

### Open Problems

- Supervised fine-tuning can degrade extrapolation if SWA windows are overly large; this seesaw effect highlights a delicate balance in architectural scheduling.
- HSA’s current implementation demands a 16:1 query/KV head ratio, imposing bottlenecks; kernel-level optimization is required.
- At short contexts (<8K), custom HSA kernels are slower than FlashAttention-3, necessitating further engineering.
- Promising directions include adaptive chunk sizing, dynamic $k$ selection, and hybrid routing (combining learned and data-driven mechanisms) to advance scalable infinite-length memory.

## 6. Significance and Context

HSA-UltraLong represents a synthesis of end-to-end sparse retrieval, hierarchical attention structure, and large-scale MoE, addressing the computational and functional challenges of ultra-long context LLMs. The model’s approach—hierarchical chunk retrieval followed by intra-chunk attention and output fusion—effectively replaces quadratic attention overhead with a mechanism characterized by sparsity, differentiable random-access, and capacity for length generalization.

This architecture provides a practical and principled solution to memory- and retrieval-centric AI, enabling LLMs to scale memory footprint and context reasoning to unprecedented lengths without sacrificing in-domain performance or computational tractability. Its open challenges suggest further avenues for scalable memory in deep learning, including optimization at both kernel and curriculum levels [2511.23319].

Source: https://www.emergentmind.com/topics/hsa-ultralong-model