Principles

Tokenization Basics BPE Algorithm GPT Tokenizers BPE Training Engineering

Model Architecture

Attention Mechanisms

Position Encoding

Position Encoding Basics RoPE Math Derivation RoPE Implementation Length Extrapolation

GPU Programming Basics

GPU Architecture Basics Tensor Layout Triton Basics: Vector Add

Flash Attention Principles From Naive to Auto-Tuning Block Pointers and Multi-Dim Support Causal Masking Optimization Grouped Query Attention Backward Pass

Hands-on Training

SystemsFlashAttention

Grouped Query Attention

Premium

Add GQA/MQA support so multiple query heads share KV, reducing KV cache memory.

Log in to continue reading

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

Causal Masking Optimization

Implement causal attention for autoregressive models and skip upper-triangular compute for ~2x speedup.

Backward Pass

Implement Flash Attention gradients with recomputation for memory-efficient training.

Table of Contents

Quick GQA Recap

Problem with Standard PyTorch GQA

Efficient Flash Attention Implementation

Core Idea: Pointer Indexing, Not Data Copy

Concrete Example

Unified Support: MHA/GQA/MQA

Full Implementation

Performance Validation

Autotune Best Config

Tradeoffs and Recommendations

Quality vs Memory

Key Implementation Summary