60 lines
2.2 KiB
Markdown
60 lines
2.2 KiB
Markdown
# Task 05: Flash Attention Forward
|
|
|
|
## 1. Problem Statement
|
|
|
|
Implement a learning-oriented forward-only FlashAttention-style kernel in both Triton and CUDA.
|
|
|
|
## 2. Expected Input/Output Shapes
|
|
|
|
- `Q`: `[batch, heads, seq_len, head_dim]`
|
|
- `K`: `[batch, heads, seq_len, head_dim]`
|
|
- `V`: `[batch, heads, seq_len, head_dim]`
|
|
- `Output`: `[batch, heads, seq_len, head_dim]`
|
|
|
|
## 3. Performance Intuition
|
|
|
|
The goal is to reduce memory traffic by avoiding full materialization of the score matrix. Correctness comes first. Performance work only matters after the blockwise algorithm is correct.
|
|
|
|
## 4. Memory Access Discussion
|
|
|
|
This task is about staged movement:
|
|
|
|
- load a `Q` block
|
|
- iterate over `K` and `V` blocks
|
|
- compute score blocks
|
|
- update online normalization
|
|
- accumulate the output block
|
|
|
|
Track where each quantity lives: global memory, registers, or shared memory.
|
|
|
|
## 5. What Triton Is Abstracting
|
|
|
|
Triton makes block pointers, program IDs, and masked block operations compact. Those abstractions still correspond to explicit memory ownership decisions.
|
|
|
|
## 6. What CUDA Makes Explicit
|
|
|
|
CUDA exposes thread-block mapping, shared-memory staging, synchronization, and reduction details directly. This is where the same algorithm becomes visibly lower level.
|
|
|
|
## 7. Reflection Questions
|
|
|
|
- How does online softmax avoid writing out the full score matrix?
|
|
- Which loop corresponds to iterating over key/value blocks?
|
|
- Where do causal masking and normalization interact?
|
|
- How does a Triton block pointer map to a CUDA shared-memory load phase?
|
|
|
|
## 8. Implementation Checklist
|
|
|
|
- Confirm the PyTorch reference on tiny shapes
|
|
- Trace the online softmax state update
|
|
- Implement one Triton blockwise forward path
|
|
- Implement one CUDA blockwise forward path
|
|
- Test non-causal first, then causal
|
|
- Benchmark only after small-shape correctness passes
|
|
|
|
## Explicit Triton To CUDA Mapping
|
|
|
|
- Triton `program_id(axis=0)` for query tiles maps to CUDA query-tile block ownership
|
|
- Triton `program_id(axis=1)` for batch/head maps to a flattened batch-head block index
|
|
- Triton block pointer math maps to shared-memory staging and pointer arithmetic
|
|
- Triton masked edge handling maps to explicit tail checks and mask branches
|