Repository navigation
DevLog
Total Attention 是对标准 FlashAttention 的一个功能扩展。它返回 attention weight 沿着 q 的序列轴求和的结果。亦即 k 对应的每一个 token 在所有的 q 中得到的 attention 的总和。
给定输入 q(B, H, M, D), k(B, H, N, D), v(B, H, N, D), FlashAttention 的计算过程中会产出 log_normalizer 和 output.
- log_normalizer(B, H, M) 是 attention 中的 softmax 运算的分母，形状为
- output(B, H, M, D) 是 v 以 softmax 输出的概率分布 p 为权重加权求和的结果。
它们都是沿着 k 的序列轴规约的结果。因此都必须迭代完一行才能计算出来，但在计算过程中并不会算出 softmax 的输出 p。
在这个扩展中，需要额外计算的是 total_attention(B, H, N). 这是 softmax 的输出 p 沿着 q 的序列轴求和的结果，亦即沿着列求和。而计算 p 又需要先完成一行的迭代。而 FlashAttention 的 kernel 的 grid 划分方案是，一个 CTA 计算 a row of tiles。如果是在这个 kernel 的基础上用简单的方式实现，要对 p 沿着列求和，那就要求在 CTA 之间通信，这么做并不方便。
因此，可以利用 FlashAttention 计算过程中输出的 log_normalizer. 额外设计一个 kernel，使用不同的 grid 划分方案，其中一个 CTA 计算 a column of tiles. 因为有了 log_normalizer(L), softmax 运算就可以变成一个 pointwise 运算。只需要重新计算 Attention score S=dot(Q,K), 利用 log_normalizer, 就可以得到 Attention weight P=exp(S-L). 然后再对 P 沿着 q 的序列轴求和即可。
计算示意图如下：
