LESSON 4.6 · Code · 180 min

FlashAttention and Triton fundamentals

FlashAttention and Triton fundamentals. Learn IO-aware tiling, online softmax, and kernel fusion and complete: Implement a simplified fused softmax. Part of the “Turn Training into

Learning objectives

  1. Explain what problem “FlashAttention and Triton fundamentals” solves without hiding behind terminology.
  2. Trace the variables and causal links across IO-aware tiling, online softmax, and kernel fusion.
  3. Complete “Implement a simplified fused softmax” and judge the result with evidence rather than intuition.

Core concepts

IO-aware tiling

IO-aware tiling is part of the lesson’s causal model. State its inputs, outputs, invariants, and failure mode; then verify it with a hand-check or a minimal experiment before moving to an optimized implementation.

online softmax

online softmax is part of the lesson’s causal model. State its inputs, outputs, invariants, and failure mode; then verify it with a hand-check or a minimal experiment.

and kernel fusion

and kernel fusion is part of the lesson’s causal model. State its inputs, outputs, invariants, and failure mode; then verify it with a hand-check or a minimal experiment.

Build and verify

Implement a simplified fused softmax

  • Predict: write the expected output, trend, or failure before running code.
  • Build: implement only the minimum components needed to answer the question.
  • Verify: compare with a baseline or trusted implementation; save seeds, parameters, and raw outputs.
  • Transfer: change one shape, dataset, scale, or workload condition and explain whether the conclusion still holds.

Open the complete interactive lesson