Flash Attention in CUDA
Follow the data, decisions, feedback, and validation boundaries before writing the full system.
How to read this diagram
Read left to right for the forward path: raw information becomes a representation, passes through the project’s main computational ideas, and produces an output that can be measured. Then follow the feedback path back toward the trainable or decision-making components.
Tiling
Tiling moves a reusable working set from slow global memory into shared memory and registers. Correct tile boundaries come first; performance then depends on coalescing, bank conflicts, occupancy, synchronization, and arithmetic intensity.
Boundary check: document its accepted input, output shape, mutable state, failure modes, and the metric that proves this stage is correct before connecting it downstream.
Online softmax
Online softmax maintains a running row maximum and rescaled exponential sum. When a new tile has a larger maximum, earlier partial sums are corrected, allowing exact stable softmax without storing the complete score matrix.
Boundary check: document its accepted input, output shape, mutable state, failure modes, and the metric that proves this stage is correct before connecting it downstream.
Memory hierarchy
Tiling moves a reusable working set from slow global memory into shared memory and registers. Correct tile boundaries come first; performance then depends on coalescing, bank conflicts, occupancy, synchronization, and arithmetic intensity.
Boundary check: document its accepted input, output shape, mutable state, failure modes, and the metric that proves this stage is correct before connecting it downstream.
Architecture review checklist
- Every arrow has a documented shape, dtype, unit, or schema.
- Training and evaluation paths cannot leak information into each other.
- Randomness is seeded and captured in experiment metadata.
- Expensive stages expose timing, memory, throughput, and error metrics.
- Each feedback loop has a stop condition and a rollback strategy.
- Small reference implementations exist for numerical comparisons.