LLM Engineering

Going from a 4k to a 128k window means 32× as many tokens: the KV-cache grows 32×, and the attention scores 1,024×. Extending a 4k-trained model with RoPE interpolation scales positions by the same factor of 32. A long window is expensive, and the cost grows faster than the length.

• Fitting text in the window is not the same as using it: accuracy drops for information in the middle of long contexts (Liu et al. 2024).
• Advertised window sizes are often tested with simple “needle” retrieval; reasoning across many parts of a long document is much harder.
• Long prompts raise cost and latency on every call; retrieval of a few relevant passages is often better.

short, long = 4_096, 131_072
tokens = long / short
attention_work = tokens ** 2                   # score matrix grows with n^2
kv_memory = tokens                             # the cache grows with n
rope_scale = long / short                      # position interpolation factor

Assumes import numpy as np. Configurations named after real models use their published shapes; other numbers are made up or simulated, as the comments say.

def sliding_window_mask(seq_len, window_size):
    """Causal + sliding window: attend to last W tokens only"""
    causal = torch.triu(torch.ones(seq_len, seq_len), diagonal=1)
    window = torch.tril(torch.ones(seq_len, seq_len),
                        diagonal=-window_size)
    mask = causal + window
    return mask.masked_fill(mask >= 1, float('-inf'))

# Mistral uses window_size=4096: each token attends to
# the previous 4096 tokens only, regardless of context lengthCopy