Long-Context Reasoning and Retrieval: Processing Extended Information

The doubt-solving bot that forgot page 1,400

An EdTech startup in Pune builds a doubt-solving assistant for JEE aspirants. The plan is simple: feed the model the entire NCERT Physics, Chemistry, and Mathematics corpus — roughly 3,000 pages, which tokenizes to about 1.5 million tokens — so it can answer any question a student throws at it, citing the exact textbook passage. The team picks a model with a genuinely large context window, 200,000 tokens, and discovers immediately that the full corpus still does not fit in one prompt. Even if it did, a second problem shows up in testing: a question whose answer sits on page 1,400 of a document that does fit gets answered worse than one whose answer sits on page 2 or page 2,998. The model is not "forgetting" in the human sense — every token is still mathematically present in its input — yet its accuracy on facts buried in the middle of a long context measurably drops. Two separate engineering problems are tangled together here, and this chapter's job is to pull them apart: first, why extending a context window is expensive and does not scale to open-ended corpora; second, why even a fact that technically fits inside the window is not read with uniform reliability. The fix for the first problem is retrieval. The fix for the second is understanding what long context actually buys you, and what it does not.

Why you cannot just make the window infinite

Recall the transformer's self-attention mechanism from your architecture study last year: for a sequence of n tokens, every token computes an attention score against every other token via the scaled dot product QKT. That score matrix has shape n × n, so the FLOPs spent computing it — and the subsequent weighted sum over values — scale as O(n²·d), where d is the model's hidden dimension. Double the context length and you roughly quadruple the compute spent on attention specifically (the feed-forward/MLP layers, by contrast, scale linearly in n, since each token is processed independently there). This is the first wall: attention compute grows quadratically with sequence length.

The second wall is memory, and it is linear rather than quadratic, which is a distinction most students get backwards — correcting exactly this is worth doing carefully later in this chapter. During autoregressive generation, a transformer caches the key and value vectors for every token it has already processed, so it never recomputes them when generating the next token — this is the KV cache. Its size is:

KV_cache_bytes = 2 × num_layers × num_kv_heads × head_dim × seq_len × bytes_per_element × batch_size

The leading 2 accounts for storing both K and V. Every other factor multiplies in linearly with sequence length: double the tokens, double the cache. That is real memory pressure, but it is a different shape of problem from the quadratic compute cost above — and production systems attack the two separately, as you will see below.

Worked example: what 128K tokens of context actually costs in GPU memory

Take a 7-billion-parameter, Llama-2-style model: 32 transformer layers, 32 attention heads, head dimension 128 (so hidden size = 32 × 128 = 4,096), running in FP16 (2 bytes per number), with ordinary multi-head attention so num_kv_heads = num_heads = 32.

Step 1 — bytes per token.

2 × 32 layers × 32 heads × 128 dims × 2 bytes = 524,288 bytes = 512 KiB per token

Step 2 — sanity check against a published figure. The same formula applied to a 13B-parameter model (40 layers, 40 heads, head_dim 128, FP16) gives 2 × 40 × 40 × 128 × 2 = 819,200 bytes ≈ 800 KiB per token — matching the per-token KV-cache figure commonly cited for 13B-class models in the systems literature (Kwon et al., "Efficient Memory Management for Large Language Model Serving with PagedAttention," SOSP 2023, the paper that introduced the vLLM serving engine). Two independent derivations landing on the same order of magnitude and the same clean structure is good evidence the formula is right.

Step 3 — scale to real context lengths, batch size 1.

8,192 tokens:   524,288 × 8,192   = 4,294,967,296 bytes = 4 GiB
131,072 tokens: 524,288 × 131,072 = 68,719,476,736 bytes = 64 GiB

Going from an 8K to a 128K context window — a 16× increase in tokens — produces a 16× increase in KV-cache memory (4 GiB → 64 GiB), because that memory cost is linear in n. But the model's FP16 weights alone already take 13.04 GiB (≈13 GiB) (7 billion parameters × 2 bytes = 14,000,000,000 bytes, ÷1,073,741,824). At 128K context, weights plus KV cache for a single sequence is ≈77.04 GiB (≈77 GiB) — nearly the entire 80 GiB HBM on one H100, before you have served a second concurrent user or allocated any memory for activations. That is why "just widen the window" is a real economic decision, not a free knob: it is a 64 GiB memory allocation per user for a 7B model at 128K tokens, on top of the model weights, on top of activation memory for the forward pass.

Verified with code:

def kv_cache_bytes(num_layers, num_kv_heads, head_dim, seq_len,
                    bytes_per_element=2, batch_size=1):
    return 2 * num_layers * num_kv_heads * head_dim * seq_len * bytes_per_element * batch_size

llama7b_8k   = kv_cache_bytes(32, 32, 128, 8192)
llama7b_128k = kv_cache_bytes(32, 32, 128, 131072)
print(llama7b_8k / (1024**3), "GiB")     # 4.0 GiB
print(llama7b_128k / (1024**3), "GiB")   # 64.0 GiB

This is exactly why production systems rarely ship plain multi-head attention at long context. Two levers bring the number down without shrinking the window. Grouped Query Attention (GQA) shares each key/value head across several query heads — Llama-3 and Mistral-class models use 8 KV heads instead of 32 query-matched ones. Re-running the formula with num_kv_heads = 8 at 128K tokens gives 2 × 32 × 8 × 128 × 131,072 × 2 = 17,179,869,184 bytes = 16 GiB — a 4× cut, exactly matching the 4× cut in KV heads (32/8), because the formula is linear in that term. Lower-precision caching (FP8 instead of FP16) is the other lever: halving bytes-per-element on the original 32-head model gives 34,359,738,368 bytes = 32 GiB at 128K — exactly half, again because bytes-per-element is a linear multiplier. Neither trick touches the quadratic attention-compute cost from the previous section; they only shrink the linear memory term. That is the production reality: long-context serving is an engineering budget spent across several independent linear and quadratic terms, not a single "context length" dial.

Lost in the middle: the recall problem that window size alone does not fix

Suppose you accept the GPU bill and fit the whole relevant passage inside the window. You are still not done. Liu, Lin, Hewitt, Paranjape, Bevilacqua, Petroni, and Liang ("Lost in the Middle: How Language Models Use Long Contexts," arXiv 2023, published in Transactions of the Association for Computational Linguistics, 2024) ran a controlled experiment: place a single fact needed to answer a question at different positions within an otherwise-identical long context, and measure how often the model retrieves it correctly. Across the models they tested, accuracy is highest when the fact sits near the very start or very end of the context, and dips when the fact sits in the middle — a U-shaped curve, not a flat one. The exact accuracy numbers differ by model, task, and context length and are not reproduced here as data (a specific model's specific score on a specific benchmark run is not something to restate from memory as a precise figure), but the qualitative shape — strong primacy and recency, weak middle — is the reproducible finding, and it recurs across many later long-context evaluations, including "needle-in-a-haystack" style tests run informally across model releases since.

Why does this happen at all, given that attention can in principle look at any token? Two compounding reasons. First, softmax attention is a competition for a fixed budget of weight that sums to 1 across all n tokens — as n grows, the average weight available per token shrinks, and the model's learned attention patterns (shaped by what training data usually rewarded) tend to allocate disproportionate weight to positions near the boundaries of a document, because instruction-style and retrieval training data disproportionately places the answer at the start of a passage (a topic sentence) or the end (a conclusion). Second, positional encodings themselves were often trained on shorter sequences than they are deployed on; a purely learned or naively extrapolated position signal degrades for positions far outside its training distribution, which further blurs the model's sense of "where" a middle-of-document token sits relative to the query.

The other lever: retrieval-augmented generation

If the corpus is larger than any affordable window (the NCERT case: 1.5M tokens against a 200K window) or if you simply do not want to pay quadratic attention cost and linear KV-cache cost for tokens that are irrelevant to a given question, the answer is not a bigger window — it is retrieval. Retrieval-Augmented Generation (Lewis, Perez, Piktus, Petroni, Karpukhin, Goyal, Küttler, Lewis, Yih, Rocktäschel, Riedel, and Kiela, "Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks," NeurIPS 2020, Facebook AI Research) keeps the context window small and instead searches an external index for the handful of passages actually relevant to the current question, injecting only those into the prompt.

The pipeline has an offline (ingestion) half and an online (query-time) half. Offline: split the corpus into overlapping chunks (a few hundred tokens each, with modest overlap so a fact split across a chunk boundary is not lost entirely), pass every chunk through an embedding model to get a dense vector, and store all vectors in an index built for fast approximate nearest-neighbour search (ANN structures such as HNSW, used inside libraries like FAISS, avoid comparing the query against every chunk one by one). Online: embed the incoming question with the same embedding model, search the index for the k chunks whose vectors are closest to the question's vector, and splice those chunks into a short prompt alongside the question before calling the LLM. "Closest" is almost always measured with cosine similarity — the cosine of the angle between two vectors, which ignores their magnitude and captures only how aligned in direction they are:

cosine_similarity(a, b) = (a · b) / (‖a‖ × ‖b‖)

Worked example: retrieving the right chunk by hand

Use toy 4-dimensional embeddings (real embedding models use hundreds or thousands of dimensions; four is enough to trace the arithmetic by hand). Say the query is "co-lending risk-sharing clause" and three chunks have already been embedded:

query = [0.8, 0.1, 0.5, 0.2]
chunk A (on-topic passage)     = [0.75, 0.15, 0.45, 0.25]
chunk B (unrelated passage)    = [0.10, 0.90, 0.05, 0.30]
chunk C (loosely related)      = [0.50, 0.30, 0.60, 0.10]

For chunk A: dot product = 0.8×0.75 + 0.1×0.15 + 0.5×0.45 + 0.2×0.25 = 0.6 + 0.015 + 0.225 + 0.05 = 0.89. Norms: ‖query‖ = √(0.64+0.01+0.25+0.04) = √0.94 ≈ 0.9695; ‖A‖ = √(0.5625+0.0225+0.2025+0.0625) = √0.85 ≈ 0.9220. Cosine similarity = 0.89 / (0.9695 × 0.9220) ≈ 0.9957. Running the same three steps for chunks B and C gives 0.2753 and 0.9181 respectively. Ranked: A (0.9957) > C (0.9181) > B (0.2753) — chunk A is retrieved first, chunk C second if k = 2, and the clearly off-topic chunk B is correctly excluded. Verified in code:

import math

def cosine_similarity(a, b):
    dot = sum(x * y for x, y in zip(a, b))
    norm_a = math.sqrt(sum(x * x for x in a))
    norm_b = math.sqrt(sum(x * x for x in b))
    return dot / (norm_a * norm_b)

query = [0.8, 0.1, 0.5, 0.2]
chunks = {
    "A": [0.75, 0.15, 0.45, 0.25],
    "B": [0.10, 0.90, 0.05, 0.30],
    "C": [0.50, 0.30, 0.60, 0.10],
}

scores = {name: cosine_similarity(query, vec) for name, vec in chunks.items()}
for name, score in sorted(scores.items(), key=lambda item: item[1], reverse=True):
    print(f"{name}: {score:.4f}")
# A: 0.9957
# C: 0.9181
# B: 0.2753

Every number printed above was independently computed and matches the hand derivation to four decimal places — this is the trace, not just an assertion of output.

Both mechanisms in one picture

Making long context itself cheaper: the engineering responses

RAG and window-extension are not competitors so much as answers to different questions ("which few passages matter?" versus "how do I make attending to many tokens affordable and reliable?"), and production systems typically use both together. On the window-extension side, four techniques matter:

FlashAttention (Dao, Fu, Ermon, Rudra, and Ré, "FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness," NeurIPS 2022) does not reduce the O(n²) FLOP count of attention at all — it computes the exact same mathematical result. What it changes is memory traffic: naive attention materializes the full n × n score matrix in GPU high-bandwidth memory (HBM), which is slow to read and write; FlashAttention tiles the computation into blocks that fit in the much faster on-chip SRAM and uses an online (running) softmax so the full n × n matrix is never written out. The FLOPs are unchanged; the wall-clock time drops sharply because HBM bandwidth, not compute, was the bottleneck. This is a clean example of a fact students often merge incorrectly: reducing memory traffic and reducing FLOP count are different optimizations, and FlashAttention is squarely the former.

Grouped Query Attention, used in the KV-cache worked example above, cuts the linear memory term by sharing KV heads across query heads — an architecture choice made at training time, not something applied afterward.

Sparse and local attention (e.g. sliding-window attention, where each token only attends to a fixed nearby window plus a few global tokens, as in Longformer-style designs) trades some full-context expressivity for attention cost that grows linearly rather than quadratically in n.

Positional-encoding extrapolation addresses the "trained short, deployed long" mismatch directly. Rotary Position Embedding (Su, Lu, Pan, Wen, and Liu, "RoFormer: Enhanced Transformer with Rotary Position Embedding," 2021) encodes position as a rotation applied to query/key vectors; Press, Smith, and Lewis's ALiBi ("Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation," ICLR 2022) instead adds a fixed distance-based penalty directly to attention scores so a model trained at, say, 4K tokens degrades gracefully at 16K; and Chen, Wong, Chen, and Tian's Position Interpolation (Meta AI, 2023) rescales RoPE's rotation frequencies so a model trained at one context length can be fine-tuned cheaply to a much longer one without the positional signal falling outside its trained distribution.

The misconception to retire

The specific error to correct: "Modern models advertise 128K–1M token context windows, so the right move is always to paste the entire document or corpus in, rather than bothering with retrieval — more context is strictly better." Three things are wrong with this. First, cost: as derived above, doubling context roughly quadruples attention FLOPs and doubles KV-cache memory, so "just include everything" is frequently the most expensive way to answer a question that a handful of relevant chunks would answer just as well. Second, recall is not uniform across the window — the lost-in-the-middle effect means a fact that technically fits can still be retrieved less reliably than the same fact would be if it were the only thing in a short, targeted prompt. Third, and most decisively for the opening example: even a 1M-token window does not fit a 1.5M-token corpus, and corpora keep growing while any fixed window stays fixed — retrieval is what lets the system scale past the window's ceiling at all, not merely what makes it cheaper below that ceiling. The corrected mental model: context-window size determines how much material a single call can hold at once; retrieval quality determines whether the right material was selected to be there. A frontier context window with poor retrieval, or a small window with excellent retrieval, both underperform a well-tuned combination of a moderate window and precise retrieval.

Active recall

Attempt each question before reading its answer.

1. Using kv_cache_bytes from the worked example, what is the KV-cache size for the Llama-2-7B configuration (32 layers, 32 heads, head_dim 128) at a 128K context, if the cache is stored in FP8 (1 byte per element) instead of FP16, batch size 1?
2. Same 128K context, but the model instead uses Grouped Query Attention with 8 KV heads (still FP16, still 32 layers, head_dim 128). What is the new per-token cache size and the new total at 128K tokens? By what factor did it shrink relative to the original 32-KV-head, FP16 case, and why does that factor match a ratio you can read directly off the two configurations?
3. Take the cosine-similarity worked example. If the query embedding changes to q' = [0.2, 0.8, 0.1, 0.3] (chunks A, B, C unchanged), which chunk is now retrieved as top-1? Show the recomputed dot products and norms for at least the winning chunk.
4. A classmate argues: "Lost-in-the-middle can't be a real architectural limit, because attention lets every token look at every other token directly — there's no distance penalty built into the math like there is in an RNN." Is the classmate's premise about the raw math correct, and if so, where does the effect actually come from?
5. Explain, in one or two sentences, why FlashAttention speeds up attention computation without changing the number of floating-point operations performed.
6. The JEE doubt-bot from the opening (3,000-page / ~1.5M-token corpus, 200K-token model window) needs to answer a question whose source fact is on page 1,400. Design the retrieval approach in three concrete steps, including one design choice (chunk size, overlap, or k) you would tune and why.

Answers.

1. Halving bytes-per-element halves the memory linearly: the FP16 total was 68,719,476,736 bytes (64 GiB), so FP8 gives 34,359,738,368 bytes = 32 GiB. Formula check: 2 × 32 × 32 × 128 × 131,072 × 1 = 34,359,738,368. Confirmed by direct computation, matching the FP16 value divided by exactly 2.

2. Per-token: 2 × 32 × 8 × 128 × 2 = 131,072 bytes = 128 KiB (versus 512 KiB for the 32-head case). At 128K tokens: 131,072 × 131,072 = 17,179,869,184 bytes = 16 GiB (versus 64 GiB). The shrink factor is exactly 4×, which is precisely num_kv_heads going from 32 to 8 (32/8 = 4) — because num_kv_heads is a linear multiplier in the formula, its ratio directly gives the memory ratio, with every other factor (layers, head_dim, seq_len, bytes/element) held fixed.

3. Dot(q', A) = 0.2×0.75 + 0.8×0.15 + 0.1×0.45 + 0.3×0.25 = 0.15+0.12+0.045+0.075 = 0.39; ‖q'‖ = √(0.04+0.64+0.01+0.09) = √0.78 ≈ 0.8832; cos(A) = 0.39/(0.8832×0.9220) ≈ 0.479. Dot(q', B) = 0.2×0.10+0.8×0.90+0.1×0.05+0.3×0.30 = 0.02+0.72+0.005+0.09 = 0.835; cos(B) = 0.835/(0.8832×0.9552) ≈ 0.990. Dot(q', C) = 0.2×0.50+0.8×0.30+0.1×0.60+0.3×0.10 = 0.1+0.24+0.06+0.03 = 0.43; cos(C) = 0.43/(0.8832×0.8426) ≈ 0.578. New ranking: B (0.990) > C (0.578) > A (0.479) — the top-1 result flips entirely from A to B, because q' now points much closer to B's direction. The lesson generalizes: retrieval quality is only as good as the embedding — a small shift in the query vector's direction can completely reorder which chunks are considered "relevant," which is why embedding-model choice and query formulation matter as much as the retrieval algorithm itself.

4. The classmate's premise about the raw math is correct — there is no hard-coded distance penalty in vanilla scaled dot-product attention; any token can in principle attend fully to any other token regardless of distance. The effect is not a hard architectural ceiling but a learned and positional-encoding artifact: models are trained on data (and often fine-tuned on instruction data) where the answer-bearing sentence is disproportionately near the start of a passage or its conclusion, so attention patterns implicitly optimize for those positions; separately, positional encodings can degrade for relative distances that were rare or absent during training, further weakening the signal for middle-of-document tokens. So it is real, reproducible, and measurable — just not "built into the math" the way an RNN's vanishing gradient is.

5. FlashAttention restructures the same O(n²) computation into tiles that fit in fast on-chip SRAM and uses an online softmax so the full n × n score matrix is never written to slow HBM; since attention on modern GPUs is usually memory-bandwidth-bound rather than compute-bound, cutting HBM reads/writes speeds up wall-clock time even though the FLOP count is identical to naive attention.

6. (a) Offline: chunk the 3,000-page corpus into passages of a few hundred tokens each with modest overlap, and embed every chunk into a vector index. (b) At query time: embed the student's question with the same embedding model and retrieve the top-k chunks by cosine similarity — since the corpus is textbook material with long explanations, a slightly larger k (say 6–10) and generous overlap (e.g. 15–20% of chunk length) are worth tuning, because a single physics derivation can span a chunk boundary and losing its second half would silently break the answer. (c) Assemble a short augmented prompt from the question plus the retrieved chunks — well under the 200K window even with generous k — and send that to the LLM instead of the 1.5M-token corpus; the fact on page 1,400 never needs its own dedicated slot in the window, because it is only pulled in on the turn a student actually asks about it.

Think About It

Think about this: How would you explain long-context reasoning and retrieval: processing extended information to a friend who has never seen a computer? What real-world analogy would you use? Imagine you had to build a system using these concepts — what would be your first step? Try this: before moving on, write down three things you learned and one question you still have.

Practice Exercises

Now it is time to practice! Complete these challenges to solidify your understanding:

• Exercise 1: Write a short program that demonstrates the core concept from this chapter. Test it with at least 3 different inputs.
• Exercise 2: Find a real-world example where long-context reasoning and retrieval: processing extended information is used in an Indian company (like TCS, Infosys, Flipkart, or ISRO). Write a paragraph explaining the connection.
• Exercise 3: Create a mind-map connecting long-context reasoning and retrieval: processing extended information to at least 3 other topics you have studied.

Key Takeaways — Summary and Recap

Let us recap what we covered: the core ideas behind long-context reasoning and retrieval: processing extended information, how they connect to real-world applications, and why they matter for your journey in computer science. Remember these key points as you move forward. For competitive exam preparation (CBSE, JEE, BITSAT), focus on understanding the WHY behind each concept, not just the WHAT.