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

TL;DR

Transformers can’t read very long documents because attention cost grows with the square of the length, and models trained on short text get confused when you feed them something much longer. This paper builds a model, HSA-UltraLong, around a mechanism called Hierarchical Sparse Attention (HSA) that treats the past like a database: it chops history into 64-token chunks, learns to retrieve the top-K most relevant chunks for the current token, attends to each chunk separately, then blends the results weighted by how relevant each chunk was. Because the retrieval weights sit inside the forward pass, the model learns what to retrieve end-to-end via ordinary next-token training. The headline result: trained only up to a 32K window, it still hits >90% accuracy on most needle-in-a-haystack retrieval tasks at 16M tokens of context, while matching a normal full-attention model on short tasks. The catch: it’s an architecture and a recipe (warm-up stages, NoPE, a small sliding window), not a drop-in trick — and the custom kernel only beats FlashAttention at very long sequences.

Problem & Motivation

A modern LLM’s knowledge is frozen into its weights. Want it to remember your last six months of conversations, or read a 10-million-token codebase? You can’t — the architecture physically chokes long before that. Two walls:

1. Quadratic cost. Standard (“full”) attention has every token look at every other token. Double the length, quadruple the compute and memory. At a few hundred thousand tokens this is already painful; at millions it’s hopeless.
2. Length generalization failure. Even if you could afford it, a model pre-trained on 4K- or 32K-token windows degrades fast when shown longer inputs. It has never learned position patterns that far out, so accuracy collapses.

The authors frame this as the problem of building “machines that can remember” — and argue real long-term memory needs three properties at once:

• Sparsity — you can’t attend to everything; retrieve only the relevant fragments (like human recall).
• Random-access flexibility — you must be able to reach any past token, not a lossy summary of it.
• Length generalization — retrieval skill learned on short contexts must transfer to long ones, because you can never train on infinite length.

Prior approaches each fail at least one:

• Recurrent / linear-attention / Mamba-style models squash the whole past into a fixed-size state vector. Cheap and length-generalizing, but that fixed state is an information bottleneck — random access to a specific distant token is gone.
• Sliding-window attention only sees the last N tokens. Distant context is simply invisible.
• Sparse-attention methods (NSA, MoBA) select chunks to attend to, but the selection isn’t truly learned end-to-end. The authors show NSA picks the wrong chunk often enough that it can’t even hit perfect accuracy on an in-domain retrieval task, and degrades quickly as context grows.

HSA is the line of work that nails all three — and this paper is the first to push it to real scale (8B params, 8T training tokens) and characterize how it scales.

What’s New (Core Contribution)

This is largely a scaling-up and recipe-discovery paper for an existing mechanism (HSA, from the same group’s earlier NeurIPS 2025 work), plus sharp empirical findings about what makes length generalization actually work. The genuine contributions:

1. HSA at scale, from scratch. Before: HSA was demonstrated on small models (sub-1B) with tiny sliding windows. Now: an 8B-A1B Mixture-of-Experts model and a 0.5B dense model trained on trillions of tokens, with a usable 4K sliding window, reaching 16M-token retrieval. This is the proof that the mechanism survives real scale.

2. The “three ingredients” finding for length generalization. Before: unclear what was load-bearing. Now: they show effective extrapolation requires the combination of (a) chunk-wise separate attention, (b) retrieval-score-weighted fusion, and (c) NoPE (No Positional Encoding) on the HSA path. Drop any one and it breaks.

3. The HSA/SWA “seesaw” finding. Before: people assumed a bigger sliding window is strictly better. Now: they show a too-large sliding window actively hurts long-range generalization — because if local attention already handles short-range dependencies, HSA never gets a gradient signal to learn retrieval, so it can’t extrapolate. Counterintuitive and practically important.

4. A warm-up recipe that resolves the seesaw. A short-window + global-HSA warm-up stage (with 1% synthetic needle tasks) bootstraps HSA’s retrieval skill before widening the sliding window, getting you both good short-task performance and long extrapolation.

5. Shared-KV context memory. They share one intermediate-layer KV cache across all HSA modules as a single “context memory,” cutting the cache cost of long contexts.

The math of HSA itself (the MoE-style retrieve-attend-fuse) is not new here. The evidence that it scales and the recipe to train it is.

How It Works (Technically)

The cleanest way to understand HSA: it’s Mixture-of-Experts, but the “experts” are chunks of your own history.

In MoE, a router scores all experts, picks the top-K, runs the token through each independently, and blends the outputs weighted by router scores. HSA does exactly this, except instead of FFN experts it has chunks of past tokens, and instead of an FFN it runs attention against each chunk.

Step 1 — Chunk the history and give each chunk a landmark

Split the sequence into fixed chunks of size S = 64 tokens (64 chosen to align with GPU hardware). Each chunk i gets:

• its own keys/values K[i], V[i] (the actual tokens, for attention later), and
• a single landmark vector K_slc[i] — a learned summary of the whole chunk’s content (produced by a small bi-directional encoder over the chunk plus a [CLS] token).

The landmark is the “address” of the chunk — like the title on a folder.

Step 2 — Retrieve: which chunks matter for this token?

For the current token t, compute a retrieval query Q_slc[t] and score it against every chunk’s landmark via a dot product:

s[t,i] = (Q_slc[t] · K_slc[i]) / sqrt(d)      for chunks i that come before t
       = -infinity                            for future chunks (causal mask)

In plain English: “how relevant is chunk i to what I’m predicting right now?” The /sqrt(d) is the standard scaled-dot-product normalizer that keeps the numbers from exploding as the vector dimension d grows. Then keep only the top-K highest-scoring chunks:

I_t = { i : rank(s[t,i]) < K }     # the K most relevant past chunks

This is the sparsity: out of millions of tokens, token t will physically attend to only K × 64 of them.

Step 3 — Attend inside each retrieved chunk separately

For each retrieved chunk i, run ordinary attention between the token’s attention-query Q_attn[t] and that chunk’s keys/values:

Ō[t,i] = Attention(Q_attn[t], K[i], V[i])
       = Softmax( norm(Q_attn[t]) · norm(K[i])ᵀ / sqrt(d_h) ) · V[i]

The norm(·) is query-key normalization — L2-normalizing the queries and keys before the dot product. The paper stresses this is critical for stability at trillion-token scale (without it, training blows up). The key design choice: each chunk is attended independently, producing one output vector Ō[t,i] per retrieved chunk. This is the “experts run independently” half of the MoE analogy.

Step 4 — Fuse: blend the chunk outputs by retrieval score

Turn the retrieval scores into softmax weights over the selected chunks, then take the weighted sum:

w[t,i] = exp(s[t,i]) / Σ_k exp(s[t,k])        # softmax over the K retrieved chunks
O_t    = Σ_i  w[t,i] · Ō[t,i]                  # the HSA output for token t

This is the crux of the whole paper. Because the retrieval score s[t,i] is multiplied into the final output, it sits on the gradient path. During backprop, if chunk i’s content helped predict the next token, the model learns to raise s[t,i] (i.e., retrieve that kind of chunk more). Retrieval is trained by the ordinary language-modeling loss — no separate retrieval supervision, no non-differentiable “pick a chunk” step. That is what NSA/MoBA lack: their selection isn’t end-to-end learnable, so they pick badly.

Contrast the philosophies in one line:

• NSA/MoBA: select chunks, then concatenate them and attend over the union. (Selection is a hard, non-differentiable gate.)
• HSA: attend to each chunk separately, then fuse by a differentiable retrieval weight. (Selection becomes a soft, learned weight.)

The positional-encoding trick: RoPE for short, NoPE for long

RoPE (rotary position encoding) bakes relative position into attention and is great in-domain, but its rotations don’t generalize to positions never seen in training — it’s a major cause of length-extrapolation failure. The fix: the sliding-window (local) attention keeps RoPE (it only ever sees nearby positions, so RoPE is fine and helpful), while the HSA (global retrieval) path uses NoPE — no positional encoding at all. Retrieval by content doesn’t need to know the absolute distance, so removing position lets it generalize to arbitrary length.

Architecture: SWA for local, HSA for global, stacked

The model has L layers split into a lower decoder (standard Transformer layers with sliding-window attention only) and an upper decoder (grouped: each group has one layer with both SWA and HSA, followed by several SWA-only layers). The intermediate layer’s hidden states produce the shared chunk summaries and KV cache that all HSA modules reuse as a single “context memory.” MoE follows Ling-2.0 / DeepSeek-V3 conventions (a dense first layer, then MoE blocks with one shared expert, auxiliary-loss-free load balancing).

Architecture & data flow

flowchart TB
  subgraph Input
    T[Current token x_t]
    H[Past history]
  end
  H --> CH[Split into 64-token chunks]
  CH --> ENC[Bi-directional encoder<br/>per chunk + CLS]
  ENC --> LM[Landmark vectors K_slc<br/>one address per chunk]
  ENC --> KV[Per-chunk KV cache<br/>shared across HSA layers]
  T --> QS[Retrieval query Q_slc]
  QS --> SCORE[Score vs every landmark<br/>s = Q_slc · K_slc / sqrt d]
  LM --> SCORE
  SCORE --> TOPK[Top-K chunks = sparsity<br/>causal mask on future]
  TOPK --> ATT[Attend to each chunk SEPARATELY<br/>NoPE, query-key norm]
  KV --> ATT
  ATT --> FUSE[Fuse: weighted sum<br/>w = softmax of retrieval scores]
  SCORE --> FUSE
  FUSE --> O[HSA output O_t]
  T --> SWA[Sliding-window attention<br/>local, keeps RoPE]
  SWA --> COMB[Combine local + global]
  O --> COMB
  COMB --> NEXT[Next-token prediction]

Schematic of HSA's retrieve-attend-fuse loop. The current token scores all past chunks by landmark relevance (top row), the top-K light up (sparsity), each is attended separately, then outputs are blended by the softmax of their retrieval scores. Click to step a token forward. Illustrative, not the paper's actual weights.

The algorithm, simplified

# HSA for one query token. Stubs: encode_chunk() -> (landmark, K, V); attention() is standard SDPA.
# Names match the paper: s = retrieval scores, I_t = selected chunks, w = fusion weights.

def hsa_step(q_slc, q_attn, chunks, top_k=64, d=128, d_h=128):
    # chunks: list of past 64-token chunks, each pre-encoded into (landmark, K, V)
    # 1) RETRIEVE: score this token's retrieval-query against every chunk's landmark
    scores = []
    for i, (landmark, K, V) in enumerate(chunks):          # causal: only past chunks present
        s_i = (q_slc @ landmark) / (d ** 0.5)              # "how relevant is chunk i to me?"
        scores.append(s_i)

    # 2) SPARSITY: keep only the top-K most relevant chunks
    selected = sorted(range(len(scores)), key=lambda i: scores[i], reverse=True)[:top_k]

    # 3) ATTEND each selected chunk SEPARATELY (NoPE here; query-key norm for stability)
    chunk_out, sel_scores = [], []
    for i in selected:
        _, K, V = chunks[i]
        o_i = attention(l2norm(q_attn), l2norm(K), V, scale=1/(d_h**0.5))  # one vector per chunk
        chunk_out.append(o_i)
        sel_scores.append(scores[i])

    # 4) FUSE: softmax the retrieval scores -> weights, blend the chunk outputs.
    #    Because w multiplies the output, retrieval is LEARNED by the LM loss via backprop.
    w = softmax(sel_scores)                                # differentiable selection weight
    return sum(w_i * o_i for w_i, o_i in zip(w, chunk_out))

The whole paper lives in that last comment: making selection a weight instead of a gate is what makes retrieval learnable, and learnable retrieval is what generalizes to 16M tokens.

Built on Prior Work

Results & Evidence

The headline (Figure 1): pre-trained with an 8K window, mid-trained to 32K, the 8B MoE hits near-perfect single-needle accuracy at 16M tokens — ~500x beyond its training length. On most in-context retrieval tasks (NIAH variants) it holds >90% accuracy out to 16M.

Short-context parity (Tables 3–4): on standard benchmarks (MMLU, GSM8K, HumanEval, etc.) within the training window, HSA-UltraLong-MoE matches an equivalent full-attention MoE (TRM-MoE) on average, and the 0.5B dense model is only ~3.3 points behind Qwen2.5-0.5B despite ~4.5x less training data. After SFT, the MoE even edges out Qwen3-1.7B on average. So sparse retrieval attention doesn’t cost you short-task quality — the usual fear with sparse attention.

The three load-bearing findings (Figure 4):

1. Effective context length of the training data matters more than the window size. Models trained on data whose real long-range dependencies are short don’t extrapolate, even with a 16K window; switching to genuinely long-dependency data (>32K) fixes it.
2. The HSA/SWA seesaw is real: a 512-token SWA window extrapolates better than 4K, and training from scratch with a 4K window fails to develop extrapolative HSA at all.
3. Capability scales with size on reasoning-retrieval: on pure retrieval (MQ-NIAH) the 0.5B and 8B are comparable, but on variable-tracking (retrieval + reasoning) the 8B clearly wins.

Be honest about the limits of this evidence:

• The long-context evidence is almost entirely synthetic retrieval probes (NIAH, variable tracking, RULER-style). These test “can you find the needle,” not “can you reason over a genuinely 16M-token document.” Real long-document QA / summarization at these lengths isn’t shown. The 16M number is impressive but narrow.
• No public efficiency win at normal lengths. Their custom HSA kernel (TileLang) only beats FlashAttention-3 at long sequences; at 4–16K, FA3 wins on both training and inference. So today this isn’t a free speedup — it’s a capability unlock that costs you at short lengths.
• Comparisons are somewhat apples-to-oranges: baselines were trained on far more data, and the architecture differs in expert config, so “parity” claims carry asterisks.
• Single paper, single group, single release — no external replication yet.

How You’d Use It

For an AI services company, the relevant question isn’t “should I train an 8B model” (you won’t) — it’s “what capability does this unlock, and how do I exploit it before it’s commoditized?”

• Memory backbone for agents. The paper’s own framing is “machines that remember.” A model that can keep a user’s entire interaction history in context — not a lossy vector-DB summary, but the actual tokens with random access — changes how you build personal/enterprise agents. Today you bolt on RAG to fake memory; HSA-style models make long memory native. If/when these ship as APIs (Ant is open-sourcing the code), the agent layer you build on top gets simpler: less retrieval plumbing, fewer “lost in the middle” failures.
• It reframes RAG, not replaces it (yet). RAG retrieves documents with a separate embedding model and a vector DB. HSA retrieves chunks of context with a learned, in-model retriever optimized for next-token prediction. The mental model to sell clients: “your retriever and your generator stop being two systems that disagree.” Near-term, you’ll still pair HSA-style long context with external RAG for corpora bigger than context; long-term it compresses the stack.
• Whole-codebase / whole-contract analysis. A client offering: ingest an entire repo or a full deal data-room (millions of tokens) and answer questions with true random access instead of chunked RAG that misses cross-references. The variable-tracking result (retrieval + reasoning scales with size) is the encouraging signal here.
• Where it slots in your MAS. In a multi-agent setup, a single long-context “librarian/memory” agent holding the shared history with HSA-style attention can serve other agents — replacing a brittle shared vector store. The retrieval being learned and differentiable means it degrades more gracefully than cosine-similarity lookup.

Realistic posture today: this is a research artifact you track and prototype against, not production infrastructure you resell this quarter. The moat is being early to design agent architectures that assume native long memory.

Build Your Own (Minimal Recipe)

You will not reproduce 8B/8T. But you can build a toy HSA layer to understand and demo the mechanism — and that understanding is the sellable asset.

Smallest version that captures ~80% of the value:

1. Start from a small open model (e.g., a 100–500M Transformer, or fine-tune the authors’ released code at github.com/ant-research/long-context-modeling). Reuse their kernels — don’t write CUDA.
2. Add one HSA layer alongside existing sliding-window attention:
   • Chunk size 64, top-K small (start 16–64).
   • A tiny bi-directional encoder per chunk producing a [CLS]-based landmark.
   • Implement the four steps from the pseudocode above: score landmarks → top-K → per-chunk attention → softmax-weighted fuse.
3. Get the three ingredients right or it won’t extrapolate: chunk-wise separate attention, retrieval-score-weighted fusion, and NoPE on the HSA path (keep RoPE only on the local SWA path).
4. Add query-key normalization (L2-norm Q and K before the dot product). Skipping this is the most likely cause of training instability.
5. Train with the warm-up first: small SWA window (512) + global HSA (large top-K covering the whole sequence) + ~1% synthetic needle-in-haystack samples. Only after retrieval accuracy is high do you widen SWA and shrink top-K. This sequencing is the trick that beats the seesaw.
6. Use data with genuinely long dependencies for the long-context phase — concatenated short docs won’t teach retrieval.

The two genuinely hard parts:

• The kernel. Naive per-chunk attention in PyTorch is fine for a toy but unusably slow at scale; the real win needs a fused sparse kernel (they used TileLang). This is the deep-systems effort.
• The training schedule. Warm-up → pre-train → long mid-train → anneal → SFT, each flipping SWA size and top-K. Getting the transitions right is finicky and is most of the paper’s real know-how.

Libraries/models to reach for: their open-source repo first; otherwise PyTorch + FlashAttention as the SWA backbone, a small encoder for landmarks, FSDP2 for any multi-GPU run.

How to Improve It

The limitations are the roadmap — each is a concrete, testable bet:

1. Kill the short-sequence penalty. HSA loses to FlashAttention-3 below ~16K. A hybrid that routes by length — plain attention for short prompts, HSA only past a threshold — would remove the only reason not to use it. Testable: measure latency crossover and gate on it.
2. Break the 16:1 head-ratio bottleneck. The paper admits HSA currently needs 16 query heads per KV head, a “severe information bottleneck.” Kernel-level work to support richer KV (e.g., grouped-query variants tuned for HSA) is an open, measurable target.
3. Fix the SFT-induced extrapolation decay. They note short SFT data degrades long-range ability (the seesaw biting again). Mixing long-context / synthetic-needle samples into SFT, or a length-curriculum during SFT, is an obvious experiment — measure NIAH@16M before vs. after.
4. Test real long reasoning, not just retrieval. The big evidence gap. Build/evaluate on genuine multi-hop QA, code-edit-impact, or contract-cross-reference tasks at 1M+ tokens. If HSA holds there, the commercial story is real; if it doesn’t, the 16M number is mostly a benchmark trophy.
5. Learn the chunk size / top-K instead of fixing them. Chunk size 64 and a fixed top-K are hardware-driven, not learned. A model that adapts top-K per token (attend to more chunks when uncertain) could trade compute for accuracy dynamically — a natural extension of the MoE analogy (variable expert count).
6. Combine with external RAG cleanly. Use HSA’s learned in-context retriever as a reranker over candidates fetched by a cheap external retriever, getting beyond-context-length scale with learned relevance. Differentiable rerank is an attractive, ownable feature.

Glossary

• HSA (Hierarchical Sparse Attention) — attention that chunks the past, retrieves the top-K relevant chunks per token, attends to each separately, and fuses by retrieval score.
• Sliding-window attention (SWA) — attention restricted to the last N tokens; cheap, captures local context only.
• Sparsity (in attention) — attending to a small selected subset of past tokens instead of all of them, to escape quadratic cost.
• Random access — the ability to reach any individual past token, as opposed to a lossy fixed-size summary.
• Length generalization / extrapolation — performing well on context lengths far longer than those seen in training.
• Landmark — a single learned vector summarizing a chunk’s content; the “address” used for retrieval.
• Top-K — keep only the K highest-scoring items (here, chunks). The sparsity knob.
• MoE (Mixture-of-Experts) — a layer where a router selects top-K expert sub-networks per token and blends their outputs; HSA is structurally analogous with chunks as experts.
• A1B / 8B-A1B — an 8B-parameter MoE with ~1B parameters activated per token (the rest are dormant experts).
• RoPE (Rotary Position Embedding) — encodes relative position by rotating query/key vectors; strong in-domain, weak at unseen lengths.
• NoPE (No Positional Encoding) — using no explicit position signal; lets content-based retrieval generalize to any length.
• Query-Key Normalization — L2-normalizing Q and K before their dot product; stabilizes large-scale training.
• KV cache — stored keys/values of past tokens so attention doesn’t recompute them; its size growth is the long-context memory bottleneck.
• NIAH (Needle-in-a-Haystack) — a probe task: hide a fact in a long context and test whether the model can retrieve it.
• RULER / BabiLong — long-context benchmark suites combining retrieval, tracking, and reasoning probes.
• FlashAttention-3 — a highly optimized exact-attention GPU kernel; the efficiency baseline here.
• TileLang — a tiled GPU programming framework used to implement the custom HSA kernel.
• FSDP2 — Fully Sharded Data Parallel (v2), PyTorch’s memory-sharding strategy for training large models across GPUs.
• Annealing (training) — a final phase on high-quality data with a decaying or low learning rate to polish the model.