some questions about SageAttention

[hermosayhl]

作者，你们好！读了你们的论文后，我有一些问题想请教一下：

1. "However, per-channel quantization cannot be applied for K, because quantization can only be performed at the outer axis (token dim) of the Matmul QK⊤"，请问为什么只能在 token 维度量化呢，而不能在 headdim 维度量化？
2. 如果 QK, PV 矩阵乘都使用 fp16 作为输入类型和累加器类型，SageAttention 还有多少优势呢？
3. 如果模型中没有 RoPE，没法把 mean(K) 的操作融合到 RoPE 中，SageAttention 会有多少性能损失呢？mean(K) 需要在非连续维度上求均值，相对而言开销还是比较大的，如果叠加到 SageAttention 的计时中；
4. "Notably, K exhibits distinct channel-wised outlier"，请问这里的 "outlier" 是啥意思呀？最大值吗？

谢谢！

[jason-huang03]

1. 你好，目前量化的粒度是 per block/warp，对于每个 batch 和 head 是独立的。每个量化组的大小是 (block_size, head_dim)。 常见的大概是 (128/64/32, 128/64) 这种配置。这种配置在细粒度量化之下还能保证硬件友好。
2. 目前 sageattention1 是使用 QK int8 矩阵乘 int32 累加器，PV 使用 fp16 矩阵乘 fp32 累加器。我们正是利用了硬件特性 (int8 矩阵乘更快) 达到的加速。如果 QK 也是用 fp16 矩阵乘和fp16累加器，性能必然下降，但是可能会使得结果更加准确。
3. mean(K) 的 overhead 约等于一次 global memory 的读写，时长和 .contiguous() 或者 .concat 这类的操作差不多，相比于整个模型或者 attention kernel 的运行时间开销是很小的。具体数据的话，在 cogvideoX-2b 中，一个 dit block 时长大约 50-60 ms，attention 时长大约 20-30 ms，而 mean(K) 时长大约 0.18-0.25 ms，可见性能损失是比较小的。
4. outlier 你可以理解为离群值，大概意思是一个量化组中特别大的值，具体可以参考 smoothquant 这篇文章。

[hermosayhl]

好的，谢谢，难得见这么快尽心回复的作者了。

1. 第一个问题，其实我问的是“为什么 headdim 维度不能做量化”？原因是这样硬件更友好？我以为是在 token 维度上量化能更大程度地保留信息
2. 个人建议作者加一组这种全 fp16 的实验，这样能体现出量化的效果；
3. 没问题。

[jt-zhang]

@hermosayhl 您好，回复您的问题如下：

1. 因为针对 headdim 量化会导致没办法反量化。具体来说，假设 q_int8, q_scale 和 k_int8, k_scale 是针对 headdim 量化，那 $QK^\top$ 就没办法通过“先执行 INT8 的矩阵乘法之后，再执行缩放”来反量化，也就没办法利用INT8的Tensor Core。
2. 谢谢您的建议。

[hermosayhl]

奥，是我的大脑没转过弯来，QK^t 的结果已经没有 headdim 维度了，所以没法反量化，谢谢！