Transformer Math

Deep dive A — Multi-Query Attention (MQA)

Why KV cache is the bottleneck

In a standard transformer serving system, the KV cache is the dominant memory consumer per active request, not the model weights. For a 12B-parameter model with 32 attention heads, 4,096 hidden dim, 128 head dim, and a 4K context window in fp16:

MQA mechanism (Shazeer, 2019, arXiv:1911.02150)

Standard multi-head attention (MHA) projects the input into H separate Q, K, V matrices — one per head. MQA keeps H separate Q projections but collapses K and V to a single shared projection (one head worth of K, one head worth of V). All H query heads attend against the same K and V. The architectural change is minimal — one line of difference in the attention code:

import torch
import torch.nn.functional as F

def mha_attention(x, Wq, Wk, Wv, num_heads, head_dim):
    """Standard multi-head attention — separate K, V per head."""
    B, T, D = x.shape
    # Project: each of num_heads heads gets its own K and V
    Q = (x @ Wq).view(B, T, num_heads, head_dim).transpose(1, 2)  # (B, H, T, d)
    K = (x @ Wk).view(B, T, num_heads, head_dim).transpose(1, 2)  # (B, H, T, d)
    V = (x @ Wv).view(B, T, num_heads, head_dim).transpose(1, 2)  # (B, H, T, d)
    # KV cache memory: B * num_heads * T * head_dim * 2 bytes * 2 (K+V)
    scale = head_dim ** -0.5
    attn = F.softmax(Q @ K.transpose(-2, -1) * scale, dim=-1)
    return (attn @ V).transpose(1, 2).reshape(B, T, -1)


def mqa_attention(x, Wq, Wk_shared, Wv_shared, num_heads, head_dim):
    """Multi-query attention — single shared K, V for all query heads."""
    B, T, D = x.shape
    Q = (x @ Wq).view(B, T, num_heads, head_dim).transpose(1, 2)  # (B, H, T, d)
    # K and V are shared: only 1 head's worth of K and V stored
    K = (x @ Wk_shared).view(B, T, 1, head_dim).transpose(1, 2)   # (B, 1, T, d)
    V = (x @ Wv_shared).view(B, T, 1, head_dim).transpose(1, 2)   # (B, 1, T, d)
    # KV cache memory: B * 1 * T * head_dim * 2 bytes * 2 (K+V) — 32x smaller!
    K = K.expand(-1, num_heads, -1, -1)   # broadcast to all query heads at attention time
    V = V.expand(-1, num_heads, -1, -1)
    scale = head_dim ** -0.5
    attn = F.softmax(Q @ K.transpose(-2, -1) * scale, dim=-1)
    return (attn @ V).transpose(1, 2).reshape(B, T, -1)

The KV cache savings are proportional to the number of attention heads (H=32 in our 12B example). In the toy single-layer example, the 64 MiB KV cache reduces to 64 / 32 = 2 MiB for that layer. For a full transformer stack, multiply by the number of cached layers before doing fleet math; the exact memory footprint depends on layer count and concurrent active sequences. This is the primary mechanism behind Character.ai's public claim that MQA, hybrid attention, and cross-layer KV sharing together cut KV-cache size by more than 20x at production scale. consumer-viable cost.

The quality trade-off: the Shazeer (2019) paper reports that MQA models trained from scratch with the same parameter budget achieve roughly 1–3% lower perplexity improvement on next-token prediction benchmarks compared to MHA models. The GQA paper (Ainslie et al., 2023, arXiv:2305.13245) showed that grouped-query attention (G=2–8 KV groups) recovers most of the quality gap while still delivering substantial KV cache savings — Character.ai chose to train with MQA (G=1) per their blog, indicating the full memory saving was worth the quality trade-off for the roleplay use case.

Interview signal: the non-obvious answer here is that MQA must be decided at model training time, not post-hoc. You cannot convert an MHA checkpoint to MQA without retraining. This is why Character.ai trained from scratch rather than fine-tuning an existing model — the architecture decision was dictated by the cost constraint, not vice versa. A candidate who confuses MQA (training-time architectural choice) with int8 quantization (post-training optimization) has misunderstood which levers are available.

Deep dive B — Int8 quantization end-to-end

What Character.ai quantizes

Most production LLM deployments use int8 for weight-only quantization (W8A16: int8 weights, fp16 activations). This saves weight memory but does not reduce KV cache memory because KV states are activations, stored in fp16. Character.ai goes further: the cost-engineering blog explicitly cites int8 KV cache storage in addition to int8 attention matmul (W8A8). This means KV cache memory halves again versus W8A16: from 2 MB/user under MQA+fp16 to 1 MB/user under MQA+int8 KV cache.

Int8 quantization: calibration and the outlier problem

The core challenge in int8 quantization (Dettmers et al., 2022, arXiv:2208.07339) is that large transformer models develop systematic outlier features — a small fraction of activation dimensions with values 100x larger than the median. Naive int8 quantization clips these outliers, causing severe quality degradation. LLM.int8() solves this with a mixed-precision decomposition^ⁱ: identify the outlier dimensions (typically 0.1% of hidden dims), compute their contributions in fp16, compute the remaining 99.9% in int8, and add. This preserves quality while achieving near-full int8 throughput.

Applying the same principle to KV cache storage: the K and V tensors also develop outlier dimensions. Character.ai's blog describes a per-channel int8 calibration applied to the KV cache — each channel maintains a running scale factor, and the stored int8 value is dequantized with that factor at attention time. The attention matmul itself runs in int8 (leveraging NVIDIA tensor core int8 paths, which run at 2x the throughput of fp16 on A100).

The calibration drift failure mode

If the calibration dataset used to compute int8 scale factors does not represent the token distribution of production traffic, the scale factors are miscalibrated. The failure manifests gradually: as the model encounters out-of-distribution tokens (a viral foreign-language character, a newly trending topic), the scale factors for affected attention layers are wrong, the quantized values saturate, and output quality degrades — not catastrophically on any single response, but systematically across a class of inputs. This is the third runbook entry below: an online recalibration that encounters distribution shift mid-shift and progressively degrades a serving shard before the monitor fires. The mitigation is a KL-divergence guard on the online calibration buffer — if the incoming token distribution diverges from the calibration distribution by more than 3-sigma, freeze the scale factors and alert.

Combined savings

All GPU counts are for KV cache memory only; decode compute adds additional GPU requirements. Numbers derived from first principles using disclosed model architecture (community estimate for exact parameter count).

Deep dive C — Character prefix caching at consumer scale

The structural opportunity

Character.ai's traffic pattern is qualitatively different from ChatGPT's. In ChatGPT, the system prompt is the same for all users of the same subscription tier — but the user conversation history is unique to each session. In Character.ai, millions of users might be chatting with the same popular character (e.g., a top-100 anime character might have 10M+ chats from hundreds of thousands of concurrent users at peak). The personality prefix for that character — 4,000 to 16,000 tokens of backstory, speech patterns, and relationship context — is identical for every one of those users. The KV states for that prefix can be computed once and reused for every single conversation, not once per user session.

This is a fundamentally different caching structure from per-session prefix caching. In ChatGPT, the prefix cache saves the system prompt (~500–1,500 tokens) per session. In Character.ai, the prefix cache saves the character personality (~4,000–16,000 tokens) per unique request — and the cache entry is shared across an arbitrary number of users simultaneously.

The hit rate ceiling and the routing dependency

The cache hit rate for a given request depends on the ratio of shared prefix tokens to total context tokens. If a character has a 6K-token personality prompt and an average user conversation history of 4K tokens (at turn 5–8 of a session), the shared fraction is 6/(6+4) = 60%^ⁱ — exactly matching the cited figure from the Character.ai blog. For longer sessions (turn 15+, ~10K history), the hit rate drops to 6/16 = 37.5%. For first-turn requests, it approaches 100% (no user history yet). The fleet-level 60% average is plausible given the distribution of session lengths.

The routing dependency is the load-bearing constraint: for the KV states to be reusable, the request must land on the GPU shard that holds those states in HBM. This requires character-to-shard affinity routing — the cache router must maintain a mapping of “character ID → primary serving shard” and route all requests for that character to the same shard. When a shard is added or removed (scaling event), the affinity mapping must be updated carefully: migrating a popular character's cache to a new shard is a warm-up event that temporarily drops the hit rate for that character to zero and causes a brief latency spike. The mitigation is a shadow-warm protocol: start serving the new shard at 0% of traffic, warm the personality prefix KV states, then ramp traffic gradually — identical to a cache warm-up before a Redis migration.

Cross-link: PagedAttention and KV cache management

The prefix cache described here sits on top of a vLLM-style PagedAttention engine^ⁱ (see the ChatGPT case study for the PagedAttention deep dive). PagedAttention manages KV memory in non-contiguous pages, enabling the serving engine to store the shared personality prefix in a single set of immutable pages that are reference-counted across all concurrent users of that character. When a user's session context extends the shared prefix (by appending their conversation history), the engine allocates new pages for the user-specific suffix without copying the shared prefix pages — the same copy-on-write semantics used in OS virtual memory management. This means the character prefix is stored exactly once per GPU shard regardless of how many concurrent users are accessing it, making the memory cost per additional user approaching zero for the shared portion.