CUDA-MODE课程笔记 第12课，Flash Attention

正文

请到「今天看啥」查看全文

，其中 是通过 Query 和 Key 的乘积得到的中间结果。

假设我们需要计算输出矩阵的一个元素 ，其中 是序列维度的token， 是隐藏层维度。

计算这个元素时，需要：

中间结果的计算

• 为了计算Softmax的权重 ，需要通过矩阵乘法 ，即Query 和 Key 的点积。

忽略 Softmax 的情况

• 暂时忽略 Softmax 的计算，讨论如何在计算中使用 分块策略，即在时间步 和序列 上进行分块计算，并在t-tile上进行循环。

分块计算的伪代码

• 将 Key 矩阵的 s-tile 和 Value 矩阵的 s-tile 加载到共享内存中
• 计算 ，得到 Softmax 的输入项 。
• 使用 Softmax 权重 更新输出

请到「今天看啥」查看全文