Tutorial 09: Transformer Block and Full Model Assembly

We have built every piece: embeddings, normalization, attention, and feed-forward networks. Now we assemble them into a complete transformer block and wire the blocks together into a full language model that can predict the next token.

---

1. The Transformer Block

A single transformer block contains two sublayers: attention and feed-forward. Each sublayer is wrapped with RMSNorm and a residual connection:

input x
  │
  ├──→ x_norm = RMSNorm(x)
  ├──→ attn_out = Attention(x_norm)
  ├──→ h = x + attn_out            ← residual connection
  │
  ├──→ h_norm = RMSNorm(h)
  ├──→ ffn_out = FFN(h_norm)
  └──→ output = h + ffn_out        ← residual connection

Why Residual Connections?

Without residual connections, the gradient signal must pass through every layer during training. With 32+ layers, this causes vanishing gradients — early layers learn almost nothing.

Residual connections create "shortcuts" that let the gradient flow directly from the output to any earlier layer. They also mean each layer only needs to learn the delta (what to add to the representation), not the full transformation.

The Pre-Norm Pattern

We apply RMSNorm before the sublayer (not after). This is the "Pre-LN" configuration, which became standard after the original "Post-LN" transformer. Pre-LN trains more stably and produces better results with deep networks.

---

2. The Block Implementation

class TransformerBlock:
    def __init__(self, layer_idx, config, weights, features):
        """
        Args:
            layer_idx: This block's index (0, 1, 2, ...).
            config: Model configuration (hidden_size, num_heads, etc.).
            weights: WeightStore for loading checkpoint weights.
            features: ModelFeatures flags for kernel selection.
        """
        prefix = f"model.layers.{layer_idx}"

        # --- Normalization layers ---
        # Two RMSNorms: one before attention, one before FFN
        self.input_layernorm = RMSNorm(
            config.hidden_size,
            weights.tensor(f"{prefix}.input_layernorm.weight"),
            eps=config.rms_norm_eps,
        )
        self.post_attention_layernorm = RMSNorm(
            config.hidden_size,
            weights.tensor(f"{prefix}.post_attention_layernorm.weight"),
            eps=config.rms_norm_eps,
        )

        # --- Attention sublayer ---
        self.self_attn = Attention(prefix, config, weights, features)

        # --- FFN sublayer ---
        # In our engine, the FFN can be either a standard MLP or a Mixture
        # of Experts, depending on whether this is a MoE layer.
        if config.is_moe_layer(layer_idx):
            self.mlp = Moe(...)   # MoE layer — covered in Tutorial 11
        else:
            self.mlp = MLP(...)   # Standard gated FFN

    def __call__(self, x, offset, cache, mask=None):
        """
        Args:
            x: Input tensor. Shape: (batch, seq_len, hidden_size)
            offset: Position offset for RoPE (used by KV cache).
            cache: Per-layer KV cache.
            mask: Attention mask ("causal", array, or None).

        Returns:
            Output tensor. Shape: (batch, seq_len, hidden_size)
        """
        # Sublayer 1: Attention with residual
        r = self.self_attn(self.input_layernorm(x), offset, cache, mask)
        h = x + r

        # Sublayer 2: FFN with residual
        r = self.mlp(self.post_attention_layernorm(h))
        return h + r

Notice how clean this is. The block is just: 1. Norm → Attention → Add 2. Norm → FFN → Add

The complexity lives inside the sublayers, not in how they connect.

---

3. The Full Model

Stack num_hidden_layers blocks on top of each other, with an embedding at the start and a vocabulary projection at the end:

Token IDs
    │
    └──→ Embedding → (batch, seq_len, hidden_size)
              │
              ├──→ TransformerBlock(0)
              ├──→ TransformerBlock(1)
              ├──→ ...
              └──→ TransformerBlock(N-1)
              │
              └──→ RMSNorm → LM Head → (batch, seq_len, vocab_size)

Implementation

class ModelForCausalLM:
    def __init__(self, config, weights, features=None):
        self.config = config
        self.features = features or ModelFeatures()
        self.num_hidden_layers = config.num_hidden_layers
        self.hidden_size = config.hidden_size
        self.vocab_size = config.vocab_size

        # --- Embedding layer ---
        embed_weight = weights.linear("model.embed_tokens")
        self.embedding = Embedding(
            vocab_size=config.vocab_size,
            embedding_dim=config.hidden_size,
            weight=embed_weight,
        )

        # --- Transformer blocks ---
        # One block per layer. Each block has its own attention + FFN + norms.
        self.layers = [
            TransformerBlock(i, config, weights, self.features)
            for i in range(config.num_hidden_layers)
        ]

        # --- Final normalization ---
        self.norm = RMSNorm(
            config.hidden_size,
            weights.tensor("model.norm.weight"),
            eps=config.rms_norm_eps,
        )

        # --- Output projection (LM head) ---
        # If embeddings are tied, we reuse the embedding matrix.
        # If not, we load a separate weight matrix.
        if not config.tie_word_embeddings and weights.has("lm_head.weight"):
            self.w_lm_head = weights.linear("lm_head")
        else:
            self.w_lm_head = None

The Forward Pass

    def __call__(self, inputs, offset=0, cache=None, logits_to_keep=None):
        """
        Full forward pass: tokens in, logits out.

        Args:
            inputs: Token IDs. Shape: (batch, seq_len)
            offset: Starting position for RoPE.
            cache: List of per-layer KV caches (created if None).
            logits_to_keep: If set, only compute logits for the last N positions.
                           Saves memory when we only need the next-token prediction.

        Returns:
            Logits. Shape: (batch, seq_len, vocab_size)
        """
        # Create cache if not provided
        if cache is None:
            cache = [TinyKvFullCache() for _ in range(self.num_hidden_layers)]

        # Step 1: Embed tokens into dense vectors
        h = self.embedding(inputs)

        # Step 2: Determine if we need a causal mask
        # Mask is only needed during prefill (seq_len > 1).
        # During decode (seq_len == 1), the single token attends to everything.
        mask = None if inputs.shape[1] == 1 else "causal"

        # Step 3: Process through all transformer layers
        for layer_idx in range(self.num_hidden_layers):
            h = self.layers[layer_idx](h, offset, cache[layer_idx], mask=mask)

        # Step 4: Optionally trim to only the last positions
        if logits_to_keep is not None:
            h = h[:, -logits_to_keep:, :]

        # Step 5: Final normalization
        h = self.norm(h)

        # Step 6: Project to vocabulary logits
        if self.w_lm_head is not None:
            return linear(h, self.w_lm_head)
        else:
            return self.embedding.as_linear(h)  # tied embeddings

---

4. Understanding the Output

The output is a tensor of shape (batch, seq_len, vocab_size). Each value is a logit — an unnormalized score for that token being next in the sequence.

# Example: batch=1, seq_len=3, vocab_size=5
# Logits for the 3 input tokens
logits = [
    [0.1, 2.3, 0.5, 0.0, 1.2],   # predictions after seeing token 1
    [1.0, 0.3, 0.1, 4.0, 0.2],   # predictions after seeing tokens 1,2
    [0.0, 0.8, 1.5, 0.3, 3.1],   # predictions after seeing tokens 1,2,3
]

# To get the next token, we apply softmax to the last row:
probs = softmax(logits[0, -1, :])  # probabilities for position 3
next_token = argmax(probs)          # token with highest probability

During generation, we typically only need the last position's logits (the prediction for the next token). The logits_to_keep parameter handles this optimization — we skip computing logits for earlier positions to save memory.

---

5. Shape Flow Through the Full Model

Input token IDs:      (batch, seq_len)            = (1, 10)
                          │
After embedding:       (batch, seq_len, hidden)    = (1, 10, 4096)
                          │
After each block:      (batch, seq_len, hidden)    = (1, 10, 4096)  ← same shape
                          │
After final norm:      (batch, seq_len, hidden)    = (1, 10, 4096)
                          │
After LM head:         (batch, seq_len, vocab)     = (1, 10, 151936)

The shape only changes at the boundaries (embedding and LM head). Every transformer block preserves the shape. This is why transformers are so modular — you can add, remove, or rearrange blocks without changing the interface.

---

6. KV Cache Creation

The model creates caches for autoregressive generation:

def create_kv_cache(self, mode="dense", page_size=128):
    if mode == "dense":
        # Simple cache: each layer gets a concat-based cache
        return [TinyKvFullCache() for _ in range(self.num_hidden_layers)]
    if mode == "paged":
        # Paged cache: shared page pool per layer for memory efficiency
        if self._page_pools is None:
            self._page_pools = [
                TinyKvPagedPool(page_size=page_size)
                for _ in range(self.num_hidden_layers)
            ]
        return [TinyKvPagedCache(pool=pool) for pool in self._page_pools]

Each layer gets its own cache because each layer has its own attention with different K/V weights. During autoregressive decode, the cache stores the K/V vectors from all previous tokens so we do not need to recompute them.

---

7. Model Loading

Loading a model from a checkpoint involves: 1. Reading config.json for model dimensions 2. Loading safetensors files for weights 3. Constructing the model with those weights

def load_model(model_path, features=None):
    from pathlib import Path

    # Step 1: Parse the model configuration
    model_path = Path(model_path)
    config_file = model_path / "config.json"
    config = Config.from_json_file(config_file)

    # Step 2: Load weight tensors from safetensors
    store = WeightStore.from_path(model_path)

    # Step 3: Construct the model
    return ModelForCausalLM(config, store, features)

---

8. Summary

Component	Input Shape	Output Shape	Purpose
Embedding	(B, L) token IDs	(B, L, D)	Token → vector
TransformerBlock	(B, L, D)	(B, L, D)	Process tokens (N times)
Final RMSNorm	(B, L, D)	(B, L, D)	Stabilize before output
LM Head	(B, L, D)	(B, L, V)	Vector → vocab scores

The full model is a pipeline: embed → process → normalize → project. Every intermediate tensor has the same shape, which makes the model easy to reason about and easy to extend.

Next: Tutorial 10 — Quantization: Fitting Large Models in Memory