AI Hardware & Software Stack
What is the main idea behind fast attention implementations such as FlashAttention?
- ASplitting computation into tiles processed in on-chip memory and never writing the large n × n intermediate table to HBM, reducing memory traffic
- BSkipping attention computation to give approximate answers
- CComputing attention on the CPU
Show answer and explanation
Answer: A. Splitting computation into tiles processed in on-chip memory and never writing the large n × n intermediate table to HBM, reducing memory traffic
A naive implementation writes the n × n table to HBM, reads it back for the softmax and again for the product with the values, so memory traffic dominates. FlashAttention splits queries, keys and values into blocks loaded into shared memory and carries the computation through in one pass, updating the softmax incrementally, never writing the table to HBM. The computation is the same, but memory traffic falls sharply and it runs faster.
- A. Correct. The result matches standard attention.
- B. It is exact, not an approximation.
- C. It exploits the GPU's on-chip memory.