LogoCookLLM Docs
LogoCookLLM Docs
HomeCookLLM

Principles

Tokenization
Tokenization BasicsBPE AlgorithmGPT TokenizersBPE Training Engineering
Model Architecture
Transformer LM
From token ids to logitsEmbedding and LM Head
Attention Mechanisms
From Self-Attention to GQAAttention Sink
Position Encoding
Position Encoding BasicsRoPE Math DerivationRoPE ImplementationLength Extrapolation
GPU Programming Basics
GPU Architecture BasicsTensor LayoutTriton Basics: Vector Add
FlashAttention
Flash Attention PrinciplesFrom Naive Implementation to Auto-TuningBlock Pointers and Multi-Dim SupportCausal Masking OptimizationGrouped Query AttentionBackward Pass Implementation
Distributed Training
Data ParallelismZeRO OptimizerFully Sharded Data ParallelTensor ParallelismPipeline ParallelismMulti-Dimensional Hybrid Parallelism

Hands-on Training

Overview
Pretraining
Pretraining DataTokenizer TrainingModel ArchitectureData PipelineTraining LoopMonitoring and Validation
X (Twitter)
FundamentalsModel ArchitectureAttention Mechanisms

From Self-Attention to GQA

Premium

Starting from Self-Attention, unpack the design trade-offs of Multi-Head, Causal Masking, and GQA / MQA in turn

What Is the Attention Mechanism

Attention is one of the most important innovations in modern deep learning. It allows a model to dynamically focus on information at different positions when processing sequence data.

In the Transformer architecture, the core idea of attention is: let each position's token perform a weighted sum over all other tokens according to its relevance to them.

The Standard Self-Attention Computation

Given an input sequence, Self-Attention computes its output through the following steps:

  1. Linear projection: project the input XXX into Query (QQQ), Key (KKK), and Value (VVV)

Log in to continue reading

This is premium content. Please log in to access the full article.

Attention Mechanisms

From MHA / Causal / GQA to Attention Sink and Gated Attention, understand the design, flaws, and evolution of attention

Attention Sink

Why the first token absorbs most attention: the mechanism and cost of this phenomenon, and why eliminating it is deferred to Gated Attention

Table of Contents

What Is the Attention Mechanism
The Standard Self-Attention Computation
PyTorch Reference Implementation
Multi-Head Attention (MHA)
Why Multiple Heads?
The Structure of MHA
PyTorch Implementation
Advantages and Challenges of MHA
Causal Attention
What Is Causal Attention?
Mathematical Representation
Why Do We Need Causal Masking?
PyTorch Implementation
The Performance Opportunity in Causal Masking
Application Scenarios
Grouped Query Attention (GQA)
The Evolution from MHA to GQA
The Memory Problem of MHA
Multi-Query Attention (MQA)
Grouped Query Attention (GQA)
The Mathematics of GQA
PyTorch Implementation
Comparison of the Three Mechanisms
Summary
Q=XWQ,K=XWK,V=XWVQ = XW_Q, \quad K = XW_K, \quad V = XW_VQ=XWQ​,K=XWK​,V=

  • Compute attention scores: compute similarity through the dot product of Query and Key

    S=QKTDS = \frac{QK^T}{\sqrt{D}}S=D​QK

  • Apply Softmax: convert scores into a probability distribution

    A=softmax(S)A = \text{softmax}(S)A=softmax(S)

  • Weighted sum: weight the Value with the attention weights

    O=AVO = AVO=AV

    • Intuition:

    • QQQ (Query): what information do I want?
    • KKK (Key): what information do I provide?
    • VVV (Value): what information do I actually contain?
    • Attention score: measures how well "what I need" matches "what you provide"

    PyTorch Reference Implementation

    import torch
import torch.nn.functional as F

def self_attention(X, W_q, W_k, W_v):
    """
    X: (batch, seq_len, d_model)
    W_q, W_k, W_v: (d_model, d_head)
    """
    Q = X @ W_q
    X
    WV​
    T
    ​

    where DDD is the head dimension, and dividing by D\sqrt{D}D​ is for numerical stability (scaled dot-product)

    # (batch, seq_len, d_head)
    K = X @ W_k
    V = X @ W_v
    # Compute attention scores
    scores = Q @ K.transpose(-2, -1) # (batch, seq_len, seq_len)
    scores = scores / (K.shape[-1] ** 0.5) # scale
    # Softmax normalization
    attn_weights = F.softmax(scores, dim=-1)
    # Weighted sum
    output = attn_weights @ V # (batch, seq_len, d_head)
    return output, attn_weights