在Pytorch的2.2版本更新文档中，官方重点强调了通过实现FlashAtteneion-v2实现了对scaled_dot_product_attention约2X左右的加速。
今天抽空亲自试了下，看看加速效果是否如官方所说。测试前需要将Pytorch的版本更新到2.2及以上，下面是测试代码，一个是原始手写的Self-Attention的实现，一个是使用Pytorch官方的scaled_dot_product_attention接口：
import time
import torch
import torch.nn.functional as F
def main():
repeat = 100
device = torch.device("cuda:0")
dtype = torch.float16
query = torch.rand(32, 8, 128, 64, dtype=dtype, device=device)
key = torch.rand(32, 8, 128, 64, dtype=dtype, device=device)
value = torch.rand(32, 8, 128, 64, dtype=dtype, device=device)
scale_factor = 0.125
ori_time_list = []
for _ in range(repeat):
torch.cuda.synchronize(device=device)
time_start = time.perf_counter()
# 原始Self-Attention实现
res = torch.softmax(query @ key.transpose(-2, -1) * scale_factor, dim=-1) @ value
torch.cuda.synchronize(device=device)
time_end = time.perf_counter()
ori_time_list.append(time_end - time_start)
fa_time_list = []
for _ in range(repeat):
torch.cuda.synchronize(device=device)
time_start = time.perf_counter()
with torch.backends.cuda.sdp_kernel(enable_math=False):
# 使用Pytorch官方提供的FA实现
res_fa = F.scaled_dot_product_attention(query, key, value, scale=scale_factor)
torch.cuda.synchronize(device=device)
time_end = time.perf_counter()
fa_time_list.append(time_end - time_start)
diff = (res - res_fa).abs().max()
ratio = [ori_time_list[i] / fa_time_list[i] for i in range(repeat)]
avg_ratio = sum(ratio[1:]) / len(ratio[1:])
print(f"max diff: {diff}")
print(f"avg speed up ratio: {avg_ratio}")
if __name__ == '__main__':
main()
执行以上代码，终端输出如下：
max diff: 0.00048828125
avg speed up ratio: 2.2846881043417118
这里使用的设备是RTX4070，跑了很多次发现确实加速2X左右，看来以后训练或者推理时可以考虑直接使用官方的scaled_dot_product_attention接口了。但是这里也发现了两个问题，一个是原始手写的Self-Attention的计算结果和直接调用scaled_dot_product_attention接口得到的结果差异有点大（注意，这里计算的Tensor都是FP16精度的），如果我切换到FP32精度差异会再小两个数量级。第二个问题是如果使用FP32的话实测没有明显加速，这个就很奇怪了，官方文档里并没有说专门针对FP16精度优化的（后面找了个A100的GPU试了下，发现FP32也是有加速的）。
