Transformer Math

Deep Dive 1 — The 1M-Token Context Window: KV Cache Math and What It Actually Costs

The Gemini 1.5 paper (arXiv:2403.05530) demonstrated a 1M-token context window with >99% recall in needle-in-haystack tests across the full window. This is architecturally significant not because of the model capability — but because of what it does to the serving cost structure.

KV cache memory derivation. For a transformer with L layers, d key/value dimension per head, and H heads, each input token requires 2 × L × d × H × bytes_per_element bytes of KV cache — factor of 2 for key and value. For Gemini Pro (architecture not publicly disclosed; using a comparable 540B-scale model as (community estimate): L ≈ 118 layers, d = 128, H = 16 KV heads for grouped-query attention, float16 = 2 bytes):

Note: with full multi-head attention (H_kv = H_q = 64), the raw number quadruples to ~3.8 MB/token and ~3.8 TB per session. Grouped-query attention (GQA), described in Ainslie et al. (2023, arXiv:2305.13245), reduces KV heads by a factor of 4–8x^ⁱ, but GQA alone does not bring 1M-token serving into HBM range — the quantization + sliding-window + prefix-sharing stack does. The Gemini report confirms use of multi-query variants but does not disclose exact group size or the production compression stack (labeled as proprietary).

Three-tier KV cache architecture. At ~1 GB effective per session, an H100 with 80 GB HBM can hold at most ~50 concurrent 1M-token sessions in HBM alone — after reserving ~30 GB for model weights. That is a hard capacity ceiling. The solution is tiered KV cache: HBM holds the most-recently-accessed layers (hot tier), DRAM holds warm context (accessed in the last few turns), and NVMe SSD holds cold context (earlier in the conversation). On cache miss, the serving node fetches from DRAM (~100 ns) or SSD (~100 µs), introducing a latency penalty proportional to the number of evicted layers that must be reloaded.

Positional encoding: YaRN extends beyond training length. Standard RoPE (Rotary Position Embedding) degrades catastrophically on sequences longer than the training context length because the rotary frequencies have not been calibrated for those positions. YaRN (Yet Another RoPE extensioN, Peng et al., 2023, arXiv:2309.00071) applies a temperature-scaled interpolation that spreads the positional encoding smoothly over the extended range. Gemini 1.5's training methodology is proprietary, but the paper acknowledges that extending to 1M tokens required specialized positional encoding treatment beyond standard RoPE (labeled in the paper as a training-time adjustment without architectural specifics).

Real-world incident: the DRAM tier eviction bug. In a hypothetical (but architecturally realistic) failure: a serving node receives a 1M-token session. The hot HBM tier holds layers 0–80. Layers 81–118 are evicted to DRAM. On turn 5, the user asks a question about context from turn 1 — content that is now in layers 81+. The eviction policy incorrectly marks those layers as cold (due to a recency-tracking bug) and moves them to SSD. Fetch from SSD adds 100 µs per layer × 37 cold layers = 3.7 ms per attention step. For a 500-token output, the total decode latency penalty is 500 steps × 3.7 ms = 1.85 seconds — a p99 TTFT regression of nearly 2 seconds with no model change. Detection: monitor KV-cache tier hit rate separately by tier; an SSD-hit spike without an HBM-size change is the signal. This failure mode is unique to long-context systems and does not appear in ChatGPT-scale case studies.

Deep Dive 2 — Multi-Modal Fusion: Image Token Budget and the 258-Token Question

The Google Gemini technical report states that images are represented as a fixed token count in the context window. The Gemini API documentation specifies approximately 258 tokens per image^ⁱ for standard (non-tiled) inputs. This number is not arbitrary — it corresponds to a patch grid.

Patch grid math.A ViT-style encoder divides an image into non-overlapping patches. For a 224×224 input image with 14×14 patches, the grid is 16×16 = 256 patches. Adding 2 special tokens (image start/end or CLS + SEP) gives 258. Larger images (e.g., 1024×1024) can be tiled — Google's documentation mentions variable token counts for high-resolution inputs. At the maximum, a high-resolution image can consume ~1,290 tokens (5× tiling). For a 10-image product catalog at standard resolution: 10 × 258 = 2,580 image tokens. That is 2.6K of a 32K context window — manageable. But at high resolution or with video (24 fps × N seconds × 258 tokens/frame), the context window fills rapidly.

Video token budget example.A 10-second video at 1 frame per second = 10 frames × 258 tokens = 2,580 tokens. At 24 fps: 240 frames × 258 tokens = 61,920 tokens — nearly 62K tokens for 10 seconds of video. This is a quadratic attention cost and a linear KV-cache cost. For a 1-minute video at 24 fps: 86,400 tokens — consuming 8.6% of the 1M-token window. The system must decide whether to downsample frame rate (lossy) or charge the full token cost (expensive). Google's approach is adaptive frame sampling — the encoder selects the most semantically distinct frames rather than sampling uniformly (described qualitatively in the Gemini 1.5 paper; exact algorithm is proprietary, labeled as (inferred from paper)).

Cross-modal attention vs late fusion. Gemini's architecture uses a natively multimodal approach — image tokens are interleaved with text tokens in the same context window and attend to each other via standard self-attention from the earliest layers (per the Gemini report). This is architecturally different from late-fusion models (e.g., LLaVA-style) that encode images separately and inject the embedding at a single layer. Native multimodal attention enables richer cross-modal grounding but means the full attention matrix must cover both modalities, making the cost proportional to (text_tokens + image_tokens)^2 in the naive case — mitigated by flash attention and GQA.

Real-world incident: image token count inflation. Google's Vision API had a documented incident in 2023 where a change to the image resizing pipeline caused images to be upscaled before tokenization, inflating token counts by 4–5x for common web images. In a Gemini-equivalent system, this would: (1) silently exceed context limits for users sending multiple images; (2) inflate billing by 4–5x per image without user notification; (3) increase GPU memory usage 4–5x per image, reducing maximum concurrent sessions. The unit-test mitigation is an exact-token-count assertion on a canonical test image — catching this class of regression before production.

Deep Dive 3 — TPU v5p vs H100: Different Performance Cliffs and What They Mean for System Design

Google's primary inference hardware is TPU v5p (and the newer v6e “Trillium”). OpenAI and most external inference providers use NVIDIA H100. These are not interchangeable — they have different performance cliffs that change system design decisions.

TPU v5p vs H100 comparison (community estimates from public specs):

The practical cliff: TPU decode throughput degrades at small batch sizes. TPU's matrix unit is optimized for large, regular tensor shapes. For autoregressive decode with batch size 1 (common in low-latency consumer scenarios), the TPU's systolic array is severely underutilized — each decode step processes a 1×d_model matrix multiplication instead of a B×d_model one. H100 handles small batches more efficiently due to its more flexible execution model. This means TPU inference is more cost-efficient at high QPS (large batches) and less efficient at low-latency single-request scenarios. Google's response: continuous batching that aggregates individual decode steps across multiple concurrent requests into a single large-batch decode step — the same technique as vLLM's continuous batching, but implemented in XLA.

Real-world incident: TPU pod configuration change causing decode regression. Google published a paper (Jouppi et al., 2023, “TPU v4: An Optically Reconfigurable Supercomputer for Machine Learning with Hardware Support for Embeddings”) describing how inter-chip link (ICI) topology changes between TPU generations can require recompilation of XLA compute graphs. In a serving system, a fleet upgrade from v4 to v5p requires re-profiling the optimal tensor parallel degree and batch size — a configuration that was optimal on v4 (e.g., 4-way tensor parallel, batch 64) may be suboptimal on v5p (optimal: 8-way TP, batch 128). A fleet upgrade without re-profiling can deliver 30–40% lower throughput than the hardware upgrade suggests, temporarily degrading SLOs during the transition window.