flash-attention-minimal

flash-attention-minimal是对Flash Attention算法在CUDA和PyTorch环境下的一个最小化重实现。Flash Attention是一种优化的注意力机制实现，能显著加速深度学习模型中的注意力计算过程。

功能总结

• 整个前向传播（forward pass）过程在flash.cu文件中用约100行代码完成。
• 变量命名遵循原始论文的规则。
• 它提供一个途径，在CUDA环境下，用PyTorch框架来实现快速的注意力计算。

适用情况

你可能会在以下情况下使用flash-attention-minimal:

• 当你是CUDA编程的初学者：官方的Flash Attention实现可能对于初学者来说太复杂了。这个最小化重实现试图提供一个简单而有教育意义的替代方案。
• 需要加速注意力机制计算：基于性能基准测试，flash-attention-minimal有效地减少了CPU和CUDA的总计算时间，提供了比手动实现更快的处理速度。
• 资源有限，但想尝试Flash Attention：如果没有GPU资源，项目还提供了一个在线Colab演示，让你可以尝试和体验Flash Attention的加速效果。

如何使用

1. 前提条件：安装PyTorch（支持CUDA）和Ninja（用于在C++中加载）。
2. 性能基准测试：通过执行python bench.py比较手动注意力计算和minimal flash attention的墙钟时间，以验证加速效果。

注意事项

• 目前没有实现反向传播（backward pass）。反向传播要比前向更复杂，但前向足以展示共享内存的使用来避免大量的N^2读/写操作。
• 在内循环中，每个线程被分配到输出矩阵的一行上。这与原始实现有所不同。
• 这种线程分配到行的简化操作使得矩阵乘法变得非常慢。这可能是为何在处理更长的序列和更大的块大小时，相比于手动实现，速度会变慢。
• 输入的Q,K,V是float32格式的，而原始实现使用的是float16。
• 编译时固定了块大小为32。

待完成的任务

• 添加反向传播。
• 加快矩阵乘法的计算速度。
• 动态设置块大小。

flash-attention-minimal是一个简化版的Flash Attention实现，非常适合CUDA编程的新手和想要快速了解或实验Flash Attention加速效果的研究者或开发者。