Context Window and Memory Mechanics in LLMs

Every time you interact with an LLM, you are constrained by its Context Window: the maximum number of tokens the model can read, process, and retain in memory during a single inference call.

For data engineers designing pipelines, the context window governs everything: whether an entire SQL table schema fits into prompt memory, how much history an agent can maintain, and how quickly API costs explode.

1. The Quadratic Scaling Problem ($O(N^2)$)

In standard self-attention, every token compares itself to every other token. If an input prompt contains $N$ tokens:

$$\text{Attention Matrix Size} = N \times N$$

flowchart LR
    subgraph AttentionGrowth["Memory & Computation vs Sequence Length"]
        N1000["1,000 Tokens<br/>1,000,000 operations"]
        N4000["4,000 Tokens<br/>16,000,000 operations (16x)"]
        N32000["32,000 Tokens<br/>1,024,000,000 operations (1024x!)"]

        N1000 -->|4x Length| N4000 -->|8x Length| N32000
    end

    style N1000 fill:#10b981,color:#fff
    style N4000 fill:#f59e0b,color:#fff
    style N32000 fill:#ef4444,color:#fff

When context length doubles, the memory required to store intermediate attention activations quadruples ($N^2$). This was why early LLMs (GPT-3) were limited to just 2,048 tokens.

How Modern LLMs Scaled to 1M+ Tokens:

1. FlashAttention (Dao et al.): Uses GPU SRAM tiling and kernel fusion to compute exact attention without materializing the massive $N \times N$ matrix in high-bandwidth memory (HBM).
2. RoPE (Rotary Position Embeddings): Encodes position information via rotation in complex vector space, allowing seamless extrapolation to longer contexts via YaRN or frequency scaling.

2. The KV Cache (Key-Value Cache)

During inference, LLMs generate text autoregressively, one token at a time.

If the model outputs token 150, does it recompute the Query, Key, and Value vectors for all preceding 149 tokens from scratch? No. Doing so would be catastrophic ($O(N^3)$ total generation compute).

Instead, the model caches the Key and Value tensors for all previous tokens in GPU memory. This is called the KV Cache:

sequenceDiagram
    participant User as Input Prompt
    participant GPU as GPU Compute Core
    participant KVCache as GPU VRAM (KV Cache)

    User->>GPU: Prompt: 'SELECT name FROM customers' (4 tokens)
    GPU->>KVCache: Compute & store K, V for tokens [0, 1, 2, 3]
    GPU->>User: Generates next token: 'WHERE'

    Note over GPU,KVCache: Step 2: Generating subsequent token
    GPU->>GPU: Compute Q, K, V ONLY for 'WHERE'
    GPU->>KVCache: Read past K, V vectors (fast lookup)
    GPU->>KVCache: Append new K, V for 'WHERE'
    GPU->>User: Generates next token: 'active'

The VRAM Cost of the KV Cache:

$$\text{KV Cache Memory (Bytes)} = 2 \times \text{Layers} \times \text{Heads} \times \text{HeadDim} \times \text{SequenceLength} \times \text{BatchSize} \times \text{BytesPerElem}$$

For a 70B model serving a 32,000 token request, the KV cache alone can consume over 20 GB of VRAM for a single user!

3. The "Lost in the Middle" Phenomenon

Just because a model supports a 128k or 1M token context window does not mean it pays equal attention to all tokens.

Research (Liu et al., Stanford) revealed that LLMs suffer from a structural retrieval bias:

flowchart TD
    subgraph ContextDistribution["Attention Retrieval Accuracy Across Context"]
        Start["Beginning of Prompt (System Prompt)<br/>⭐⭐⭐⭐⭐ High Accuracy (Primacy Effect)"]
        Middle["Middle of Prompt (Tokens 20,000 to 80,000)<br/>⚠️ Degradation & Hallucinations ('Lost in the Middle')"]
        End["End of Prompt (Latest Query / Context)<br/>⭐⭐⭐⭐⭐ High Accuracy (Recency Effect)"]

        Start --> Middle --> End
    end

    style Start fill:#10b981,color:#fff
    style Middle fill:#ef4444,color:#fff
    style End fill:#10b981,color:#fff

4. Massive Context vs. Retrieval-Augmented Generation (RAG)

With Gemini offering 2M token context windows, data engineers often ask: "Is RAG obsolete? Can I just dump my entire corporate database schema and documentation into the prompt?"

5. Hands-on: Context Truncation and Buffer Management

When streaming logs or historical database transactions into an LLM, implement deterministic windowing to protect against context overflow:

import tiktoken

class RollingContextManager:
    """
    Manages prompt memory to guarantee that input payloads 
    never breach the model's context ceiling.
    """
    def __init__(self, model_name: str = "gpt-4o-mini", max_allowed_tokens: int = 4096):
        self.tokenizer = tiktoken.encoding_for_model(model_name)
        self.max_allowed_tokens = max_allowed_tokens
        self.system_prompt = "You are a specialized SQL optimization assistant."
        self.system_tokens = len(self.tokenizer.encode(self.system_prompt))
        self.history = []

    def count_tokens(self, text: str) -> int:
        return len(self.tokenizer.encode(text))

    def add_message(self, role: str, content: str):
        tokens = self.count_tokens(content)
        self.history.append({"role": role, "content": content, "tokens": tokens})
        self._prune_history()

    def _prune_history(self):
        """Prune oldest user/assistant turns if total exceeds threshold (FIFO)."""
        total_tokens = self.system_tokens + sum(m["tokens"] for m in self.history)

        while total_tokens > self.max_allowed_tokens and self.history:
            removed = self.history.pop(0)
            total_tokens -= removed["tokens"]
            print(f"[Prune] Removed oldest {removed['role']} turn ({removed['tokens']} tokens) to free memory.")

    def get_payload(self):
        return [
            {"role": "system", "content": self.system_prompt},
            *[{"role": m["role"], "content": m["content"]} for m in self.history]
        ]

if __name__ == "__main__":
    manager = RollingContextManager(max_allowed_tokens=80)

    print("Initial System Tokens:", manager.system_tokens)
    manager.add_message("user", "Explain how partition pruning works in Apache Spark SQL.")
    manager.add_message("assistant", "Partition pruning skips reading directories on disk that do not match the WHERE clause.")
    manager.add_message("user", "Now explain dynamic partition pruning introduced in Spark 3.0.")

    print("\nFinal Context Messages retained:")
    for msg in manager.get_payload():
        print(f"[{msg['role'].upper()}]: {msg['content'][:60]}...")

Key Takeaways

• Standard self-attention scales quadratically ($O(N^2)$) in memory, making long context computationally expensive.
• The KV Cache accelerates autoregressive generation by caching Key and Value states in GPU VRAM, but consumes gigabytes of memory under high concurrency.
• The "Lost in the Middle" effect degrades reasoning quality when critical details are buried in lengthy prompts.
• RAG remains essential: While 1M+ token windows are incredible for deep document analysis, RAG is 100x cheaper and 10x faster for production search and data engineering workflows.