Implement FlashAttention-2 in Triton in PyTorch
Expertgpu-systemsauto-graded
Write the tiled, fused attention kernel that computes exact attention in O(n) memory using online softmax in Triton.
Solve it
Check your answer
The grader verifies properties of your implementation, so a correct solution written differently from ours still passes.
pip install torchleet
from torchleet import check
check("flash-attention-triton", flash_attention_pytorch)