Sage Attention V2

算法流程

Smoothing Q Matrix

通过改变量化颗粒度来解决异常值的问题

上文我们用 per-block smooth QQ 矩阵的方法来处理 QQ 矩阵中异常值,我们也可以使用更加细粒度的来处理异常值的问题.

然而,如果直接使用 per-token granularity 的话,其在 dequantization 阶段会导致 significant overhead.原因是在 Attention Score 这一步矩阵乘的时候,每一个 thread 需要处理 QQ 矩阵的多个 token 和 KK 矩阵的多个 token,也就是需要读取多个 token 的 quantization scale δQ,δK\delta_Q, \delta_K,这多读取的 scale 就会导致大量 overhead

所以提出 per-thread quantization 来 mitigate 这个问题.