Deep Dive: LLM Concepts

The Flow

1. Token Embedding: Convert token IDs to vectors
2. Position Encoding: Add position information (RoPE applied inside attention)
3. N Decoder Layers, each containing:
   • RMSNorm → Self-Attention → Residual Add
   • RMSNorm → MLP (SwiGLU) → Residual Add
4. Final RMSNorm
5. LM Head: Project back to vocabulary (often weight-tied with embedding)

The Flow with RoPE

Traditional (GPT-2 style):
  token_embed + position_embed → layers

With RoPE:
  token_embed → layers (RoPE applied inside attention)

Token: "The"  "cat"  "sat"  "on"   "mat"
Pos:     0      1      2      3      4

Embedding lookup (NO position added):
  [vec0] [vec1] [vec2] [vec3] [vec4]
           |
      Decoder Layer
           |
      Q, K, V projection
           |
    +---------------------------+
    |  RoPE: rotate Q and K     |
    |  Q_0 rotated by 0*theta   |
    |  Q_1 rotated by 1*theta   |
    |  Q_2 rotated by 2*theta   |
    |  (same angles for K)      |
    +---------------------------+
           |
      Attention: Q @ K.T (now position-aware!)
    

RoPE Math

For each pair of dimensions (2i, 2i+1) at position m:

Forward:
    theta_i = rope_theta ^ (-2i / head_dim)
    angle = m * theta_i

    x'[2i]   = x[2i] * cos(angle) - x[2i+1] * sin(angle)
    x'[2i+1] = x[2i] * sin(angle) + x[2i+1] * cos(angle)

Backward:
    Simply rotate by negative angle (transpose of rotation matrix)
    

Our implementation: rope_kernels.c - precomputes cos/sin cache, applies in-place to Q and K

Attention Math

1. Project input to Q, K, V:
   Q = input @ W_q    # What am I looking for?
   K = input @ W_k    # What do I contain?
   V = input @ W_v    # What do I offer?

2. Compute attention scores:
   scores = Q @ K.T / sqrt(d_k)    # How relevant is each position?

3. Apply causal mask (decoder only):
   scores[i][j] = -inf if j > i    # Can't look at future tokens

4. Softmax to get weights:
   weights = softmax(scores)       # Normalize to probabilities

5. Weighted sum of values:
   output = weights @ V            # Aggregate information
    

Online Softmax Trick

The key insight: you can compute softmax incrementally without seeing all values first.

Traditional: softmax(x) = exp(x) / sum(exp(x))  # Need all x first

Online (Flash):
    For each new block of scores:
        1. Update running max: m_new = max(m_old, max(block))
        2. Rescale previous sum: sum *= exp(m_old - m_new)
        3. Add new block contribution: sum += sum(exp(block - m_new))
        4. Update running output with correction factor
    

Key Insight: Flash Attention reduces memory traffic from O(N²) to O(N), making it valuable for any hardware where memory bandwidth is a bottleneck. The algorithmic improvement (online softmax + tiling) applies to both CPU and GPU.

Why GQA?

GQA is the sweet spot: 4x smaller KV cache with minimal quality loss.

Why Remove Mean Centering?

Research found that the re-centering (subtracting mean) in LayerNorm isn't necessary for good performance. RMSNorm keeps just the scaling, which is what matters for gradient flow.

Used in: Llama, Mistral, Qwen, Gemma (most modern LLMs)

Why Gating Works

The gate learns to selectively activate different features:

gate = activation(x @ W_gate)  # Swish (SwiGLU) or GELU (GeGLU)
value = x @ W_up            # The actual content

output = gate * value       # Gate controls information flow
    

This gives the network more expressive power: it can learn to completely shut off certain dimensions for certain inputs.

What Gets Shared

Input embedding:  token_id → vector    E[vocab, hidden]
Output projection: vector → logits     W[vocab, hidden]

With weight tying: W = E  (same matrix!)
    

Savings: For vocab=128K, hidden=4096: saves 2 GB of parameters