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)
SystemsFlashAttention

Causal Masking Optimization

Premium

Implement causal attention for autoregressive models, achieving ~2x speedup by skipping the upper-triangular computation.

Get code access

In the previous three chapters, we implemented a complete Flash Attention kernel that supports arbitrary sequence lengths and multi-dimensional parallelism. But we have not yet optimized for one very important application scenario: autoregressive generation.

This is precisely the core computation pattern of decoder-only models such as GPT and LLaMA. By introducing causal masking, we can gain a further speedup of about 2x.

Quick Review of Causal Attention

Causal attention is the core mechanism of autoregressive language models: when predicting the iii-th token, the model can only see information from positions 0 to i-1, and cannot see future tokens.

In the attention matrix, this is implemented by setting the upper-triangular part to −∞-\infty−∞:

         K₀  K₁  K₂  K₃
    Q₀ [ ·  -∞  -∞  -∞ ]  ← Q₀ can only see K₀
    Q₁ [ ·   ·  -∞  -∞ ]  ← Q₁ can only see K₀,K₁
    Q₂ [ ·   ·   ·  -∞ ]  ← Q₂ can only see K₀,K₁,K₂
    Q₃ [ ·   ·   ·   · ]  ← Q₃ can see everything

If you are not yet familiar with the concept and principles of causal attention, we recommend first reading the detailed introduction in Attention Mechanism Explained.

This chapter focuses on how to efficiently implement causal masking in Flash Attention to gain a performance improvement.

Log in to continue reading

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

Block Pointers and Multi-Dim Support

Scale from single sequence to Batch/Head parallelism and simplify pointer math with block pointers.

Grouped Query Attention

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

Table of Contents

Quick Review of Causal Attention
The Performance Opportunity of Causal Masking
The Mathematical Principle of Halving the Compute
Visualization: The Skipped Computation Region
Code Implementation Walkthrough
Change 1: Coarse-Grained Skip — Loop Bound Optimization
Change 2: Fine-Grained Mask — Boundary Handling Inside a Block
The Necessity of Two Layers of Masking
Performance Comparison and Validation
Numerical Correctness Validation
Performance Comparison
Speedup Across Different Sequence Lengths
Summary of Implementation Techniques
The Clever Use of Triton Compile-Time Constants
Summary