This is my GPU course final project in MICS600J. The main content is my attempt to implement the attention mechanism efficiently.

In order to simplify the process, I replaced the original matrix dimensions [batch_size, nheads, seq_len, headdim] with [seq_len, headdim].

N: seq_len d: headdim

The attention mechanism is well known and I won’t go into details.

1. In = N * d
2. WQ WK WV = d * d
3. Q K V = N * d
4. P = Q * K^T = N * N
5. S = SoftMax(P) = N * N
6. o = S * V = N * d
7. Out = S * W = N * d

1. Use Tensor core to compute GEMM.
2. Use Asynchronous transfer to overlap computation and communication(transfer data from global memory to shared memory).
3. Bank Conflict free.

Use make command to build program.

1. Range for N,d: N~(32, 1024), d~(32, 2048). Please get more detail from script.
2. Test on NVIDIA A100 in HKUST(GZ)-HPC Server.

Fine-tuning Llama-2-7B, when using Sparse Attention Mechanism, we found that accuracy can be improved and restored with little overhead.

1. Kernel Fusion, just like Flash attention.
2. Sparse Attention Mechanism, just like DFSS, make full use of sparse tensor core.