llm-internals-learning

What this skill does

```markdown
---
name: llm-internals-learning
description: Skill for understanding and implementing LLM internals concepts from tokenization to attention mechanisms to inference optimization
triggers:
  - "explain how attention works in transformers"
  - "implement KV cache for LLM inference"
  - "show me how tokenization works in LLMs"
  - "how does flash attention work"
  - "implement causal masking in attention"
  - "explain mixture of experts architecture"
  - "how does byte pair encoding work"
  - "implement transformer attention from scratch"
---

# LLM Internals Learning Guide

> Skill by [ara.so](https://ara.so) — Daily 2026 Skills collection.

A structured, step-by-step educational resource for learning how Large Language Models work internally — from tokenization and embeddings through attention mechanisms, transformer architecture, and inference optimization techniques like KV Cache and Flash Attention.

---

## What This Project Covers

| Topic | Concept |
|---|---|
| Tokenization | BPE, subword splitting, vocabulary |
| Attention | Q, K, V matrices, scaled dot-product |
| Masking | Causal masking, attention masks |
| Architecture | Transformer encoder/decoder, FFN, MoE |
| Optimization | KV Cache, Paged Attention, Flash Attention |
| Training | Backpropagation, gradient descent |

---

## Core Concept Implementations

### 1. Byte Pair Encoding (BPE) Tokenization

```python
from collections import defaultdict

def get_vocab(text: str) -> dict:
    """Build initial character-level vocabulary with end-of-word marker."""
    vocab = defaultdict(int)
    for word in text.strip().split():
        chars = " ".join(list(word)) + " </w>"
        vocab[chars] += 1
    return dict(vocab)

def get_pairs(vocab: dict) -> dict:
    """Count all adjacent symbol pairs in the vocabulary."""
    pairs = defaultdict(int)
    for word, freq in vocab.items():
        symbols = word.split()
        for i in range(len(symbols) - 1):
            pairs[(symbols[i], symbols[i + 1])] += freq
    return dict(pairs)

def merge_vocab(pair: tuple, vocab: dict) -> dict:
    """Merge the most frequent pair throughout the vocabulary."""
    new_vocab = {}
    bigram = " ".join(pair)
    replacement = "".join(pair)
    for word in vocab:
        new_word = word.replace(bigram, replacement)
        new_vocab[new_word] = vocab[word]
    return new_vocab

def run_bpe(text: str, num_merges: int = 10) -> list:
    """Run BPE algorithm for num_merges iterations."""
    vocab = get_vocab(text)
    merges = []

    for _ in range(num_merges):
        pairs = get_pairs(vocab)
        if not pairs:
            break
        best_pair = max(pairs, key=pairs.get)
        vocab = merge_vocab(best_pair, vocab)
        merges.append(best_pair)
        print(f"Merged: {best_pair} -> {''.join(best_pair)}")

    return merges

# Example usage
text = "low lower lowest newer newest"
merges = run_bpe(text, num_merges=8)
```

---

### 2. Scaled Dot-Product Attention (Q, K, V)

```python
import numpy as np

def scaled_dot_product_attention(Q: np.ndarray, K: np.ndarray, V: np.ndarray,
                                  mask: np.ndarray = None) -> tuple:
    """
    Compute scaled dot-product attention.

    Args:
        Q: Query matrix  (seq_len, d_k)
        K: Key matrix    (seq_len, d_k)
        V: Value matrix  (seq_len, d_v)
        mask: Optional causal mask (seq_len, seq_len)

    Returns:
        output: Attention-weighted values
        weights: Attention weight matrix
    """
    d_k = Q.shape[-1]

    # Step 1: Compute raw attention scores
    scores = Q @ K.T  # (seq_len, seq_len)

    # Step 2: Scale by sqrt(d_k) to stabilize gradients
    scores = scores / np.sqrt(d_k)

    # Step 3: Apply causal mask (if provided)
    if mask is not None:
        scores = np.where(mask == 0, -1e9, scores)

    # Step 4: Softmax to get attention weights
    # Subtract max for numerical stability
    scores_shifted = scores - scores.max(axis=-1, keepdims=True)
    exp_scores = np.exp(scores_shifted)
    weights = exp_scores / exp_scores.sum(axis=-1, keepdims=True)

    # Step 5: Weighted sum of values
    output = weights @ V  # (seq_len, d_v)

    return output, weights


# Example: 3-token sequence, d_k=4, d_v=4
np.random.seed(42)
seq_len, d_k, d_v = 3, 4, 4

# Learnable projection matrices
W_Q = np.random.randn(d_k, d_k)
W_K = np.random.randn(d_k, d_k)
W_V = np.random.randn(d_k, d_v)

# Token embeddings
X = np.random.randn(seq_len, d_k)

Q = X @ W_Q
K = X @ W_K
V = X @ W_V

output, weights = scaled_dot_product_attention(Q, K, V)
print("Attention weights:\n", np.round(weights, 4))
print("Output shape:", output.shape)
```

---

### 3. Causal Masking

```python
import numpy as np

def create_causal_mask(seq_len: int) -> np.ndarray:
    """
    Create a causal (autoregressive) mask.
    Position [i,j] = 1 if token i can attend to token j (j <= i).
    Position [i,j] = 0 if token i cannot attend to token j (j > i).
    """
    mask = np.tril(np.ones((seq_len, seq_len)))
    return mask

def masked_attention(Q, K, V, seq_len):
    mask = create_causal_mask(seq_len)
    output, weights = scaled_dot_product_attention(Q, K, V, mask=mask)
    return output, weights

# Visualize the causal mask for a 5-token sequence
seq_len = 5
mask = create_causal_mask(seq_len)
print("Causal Mask (1=can attend, 0=blocked):")
print(mask.astype(int))

# Token 0 can only see itself
# Token 4 can see all previous tokens
```

---

### 4. Multi-Head Attention

```python
import numpy as np

class MultiHeadAttention:
    def __init__(self, d_model: int, num_heads: int):
        assert d_model % num_heads == 0, "d_model must be divisible by num_heads"
        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_heads

        # Initialize projection weights
        np.random.seed(0)
        self.W_Q = np.random.randn(d_model, d_model) * 0.01
        self.W_K = np.random.randn(d_model, d_model) * 0.01
        self.W_V = np.random.randn(d_model, d_model) * 0.01
        self.W_O = np.random.randn(d_model, d_model) * 0.01

    def split_heads(self, x: np.ndarray) -> np.ndarray:
        """Split last dimension into (num_heads, d_k)."""
        seq_len = x.shape[0]
        x = x.reshape(seq_len, self.num_heads, self.d_k)
        return x.transpose(1, 0, 2)  # (num_heads, seq_len, d_k)

    def forward(self, X: np.ndarray, mask: np.ndarray = None) -> np.ndarray:
        seq_len = X.shape[0]

        # Project to Q, K, V
        Q = X @ self.W_Q  # (seq_len, d_model)
        K = X @ self.W_K
        V = X @ self.W_V

        # Split into heads
        Q = self.split_heads(Q)  # (num_heads, seq_len, d_k)
        K = self.split_heads(K)
        V = self.split_heads(V)

        # Attention per head
        head_outputs = []
        for h in range(self.num_heads):
            out, _ = scaled_dot_product_attention(Q[h], K[h], V[h], mask)
            head_outputs.append(out)

        # Concatenate heads
        concat = np.concatenate(head_outputs, axis=-1)  # (seq_len, d_model)

        # Final projection
        output = concat @ self.W_O
        return output


# Usage
d_model, num_heads, seq_len = 512, 8, 10
mha = MultiHeadAttention(d_model=d_model, num_heads=num_heads)
X = np.random.randn(seq_len, d_model)
mask = create_causal_mask(seq_len)
output = mha.forward(X, mask=mask)
print("Multi-head attention output shape:", output.shape)
```

---

### 5. KV Cache Implementation

```python
import numpy as np
from typing import Optional

class KVCache:
    """
    Key-Value cache for autoregressive LLM inference.
    Avoids recomputing K and V for previously seen tokens.
    """
    def __init__(self, max_seq_len: int, d_k: int, num_heads: int):
        self.max_seq_len = max_seq_len
        self.d_k = d_k
        self.num_heads = num_heads
        self.current_len = 0

        # Pre-allocate cache buffers
        self.k_cache = np.zeros((num_heads, max_seq_len, d_k))
        self.v_cache = np.zeros((num_head