为什么需要 ring attention
ring attention 将输入序列拆分为多个块，每个块的推理计算运行在不同的 GPU 卡上，用于单卡上的显存消耗，使得可以支持更长的序列。降低单卡显存主要体现在以下两方面：
（1）可以减少 qkv 计算的中间过程临时显存占用，避免推理过程中用尽显存。$Q*K^T$ 会产生 $seq_len^2$ 个矩阵元素，对长序列来说很占用显存，并且使用 TP 无法解决，TP 只减少做矩阵乘法时的加和元素个数，可以减少计算量，无法减少结果矩阵元素的个数。
（2）减少单卡的 KV Cache 占用，避免显存用尽，分两类场景
- 单机 KV Cache 随着 TP 扩大而减小的模型，虽然可以扩 TP 来减少单机 KV Cache 显存占用，但 TP size 受限于 HEAD size 无法无限扩大，超长序列产出的 KV Cache 仍然可能将单卡显存用尽；
- 单机 KV Cache 无法随着 TP 扩大而减小的模型（譬如 MLA 结构的模型），每个 TP rank 上都有所有 token 的完整 KV Cache，超长序列可以将显存用尽；
原理简介
大模型的一层推理有两个核心计算环节：Attention 和 FFN。先说 FFN，对于输入序列分块的 ring attention 模式来说，每个卡（或者多个卡组成的小 TP 集群，下同）都会加载完整的 FFN 权重，因此 FFN 可以在本卡上独立完成计算。再看看 attention，attention 计算中的 K、V 依赖完整的序列，部分序列在其他卡上，所以需要网络通信获得其他卡上的 K、V 数据，如果通过 all-gather 一次性获取其它卡上的 K、V 数据，单卡还是需要存储完整序列的 K、V，并且注意力计算的中间结果也还是完整序列长度平方，这就达不到降低显存的目的，因此 ring attention 采用的是“滚动 KV”的方式，这也是它叫做“ring”的原因。
其分块做注意力计算过程可以用以下为代码描述：
output = zeros_like(Q)
for q_chunk in Q.split(chunk_size):
# 对每个Q块，逐块处理KV并累加结果
attn_scores_chunk = 0
for kv_chunk in KV_cache.split(chunk_size):
# 只计算当前Q块与KV块的交互
attn_scores_chunk += matmul(q_chunk, kv_chunk_k.transpose())
# 应用Softmax并累积输出
output_chunk = matmul(softmax(attn_scores_chunk), kv_chunk_v)
output += output_chunk
上述每个 chunk 都处于不同的卡上，因此计算之前需要先“滚动KV”。
计算流程图
下面通过计算流程图描述 ring attention 的推理过程：
核心要点：
- 自回归推理模型里只有注意力计算部分需要考虑序列的其他 token，需要通过 “kv流转”实现。
- 投影、FFN、残差连接都只和参与计算的 q 有关系，是 local 的。
- local 注意力计算，也需要考虑掩码，考虑计算均衡，拆分的时候会考虑计算的负载均衡；
分块注意力计算公式推导
分块后注意力计算的差异
$$
\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^\top}{\sqrt{d_k}}\right)V
$$
从公式可以看到，注意力计算分数有两部分用到全局信息，一个是 K，一个是 V，只有局部信息算出来的局部注意力结果和全局注意力结果是什么关系呢？肯定不是简单加和，因为本地 attention 计算过程中的 softmax 公式的分母不是全局的，局部 attention 直接加和后得到的值会虚高。准确的做法是将局部注意力拆分为分子和分母，分子和分母都做累加，下面对公式做推导。
基础公式
（1）注意力计算公式
$$
\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^\top}{\sqrt{d_k}}\right)V
$$
（2）softmax 公式
一行向量的第 i 个元素的 softmax 公式：
$$
\text{softmax}(x)_i = \frac{\exp(x_i)}{\sum_j \exp(x_j)}
$$
softmax 是对一个行向量做归一化操作，得到的结果也是一个行向量，行扩展到矩阵即是批量操作，计算 [N, D] 矩阵的 softmax 时，分母是一个 [N, 1] 列向量，分子是 [N, d]的向量。
注意力计算过程拆解
未分块时单个 token 的注意力
（1）标准计算过程：
- 第 i 个 token 的向量为行向量$Q_i$，即 [1, d] 矩阵，d 为 hidden state 维度
- $K^T$ 为 [d, N] 矩阵，N 为 序列长度
- $Q_i* K^T$ 结果为 [1, N] 矩阵，除以 K 向量的维度 $\sqrt{d_k}$ 后还是 [1, N] 矩阵
- softmax 运算，对 [1, N] 矩阵的每一个元素求自然指数，作为分子，所有元素的自然指数加和结果作为分母
- $V$ 为 [N, d] 矩阵， 乘以 V 后得到的结果是 [1, d] 矩阵
（2）为后面做序列分块做准备，计算过程调整和细化：（TODO，补充介绍怎么从原始公式转换为 Qi 公式，核心：分子部分是将矩阵QKT矩阵乘法结果的单个元素拿出来考虑，表现为做QKT 矩阵乘法时，K 按 token 拆开一个个考虑，这也就和 ring attention 的拆序列对上了。）
乘以 V 矩阵可以和 softmax 的分子合并到一起计算，同时把 K 和 V 进一步拆分为 N 个向量对应到完整序列的 N 个 token，最终公式如下：
$$
\text{Attention}(Q_i, K, V) = \frac{\sum_{j=1}^{N} \exp\left(\frac{Q_i K_j^\top}{\sqrt{d_k}}\right) \cdot V_j}{\sum_{m=1}^{N} \exp\left(\frac{Q_i K_m^\top}{\sqrt{d_k}}\right)}
$$
分子部分
- j 表示 K 和 V 的第 j 个 token 对应的向量，形状都是 [1, d]。
- exp 的参数 $\frac{Q_i K_j^\top}{\sqrt{d_k}}$ 是一个标量（[1,d]*[d,1]=1），exp 后得到的标量乘以向量 $V_j$ 得到向量，最外侧的加和是向量同位置加和
- 最后得到的结果是向量
分母部分
- m 表示 K 的第 m 个 token 对应的向量
- exp 的参数$\frac{Q_i K_m^\top}{\sqrt{d_k}}$ 是标量，exp 后得到的也是标量，最外侧的加和是标量值加和
- 最后得到的结果是标量
容易困惑的地方
转换后的公式，第一眼看起来好像不对：softmax 公式原来是只有分母加和，分子只有一个元素，转换之后怎么分子和分母都加和了？
因为 softmax 公式是在描述一个 [1, d] 向量里的单个元素是怎么算出来的；而 Attention 公式是在描述 [1, d] 的向量是怎么算出来的，包括里 d 个元素。
分子的公式要这么看：
- Qi 向量 [1, d] 和 KT 向量 [1, d] 做乘法得到一个数字再除以根号d，再做 exp 运算，也还是一个数字，再乘以 V，得到一个形状为 [1, d] 的向量，这个向量代表：在指定 Q 向量的情况下，对某一个 K token 的计算结果。然后加和指的是 N 个 [1, d] 向量加和，这里的 N 对应到 N 个 K 矩阵。最后得到的形状还是是 [1, d]。
分母的公式要这么看：
- Qi 向量 [1, d] 和 KT 向量 [1, d] 做乘法得到一个数字再除以根号d，再做 exp 运算，也还是一个数字，这个数字代表：在指定 Q 向量的情况下，对某一个 K token 的计算结果。多个 K token 的计算 结果加和得到分母。
这个拆解方式也和 CP 对应得上：每个 CP rank 只有部分 token，在做注意力计算的时候，指定某个 Q 对本 CP rank 下的 token 做计算，然后加和。得到的也就是：在指定 Q 向量的情况下，对某一小批 K token 的计算结果。 多个 CP rank 下的分子和分母分别加和，然后再相除，得到的就是不分 CP 时的结果。
分块时单个 token 的注意力计算
将上述的公式的分子和分母分别定义为 A 和 B：
- 分子：$A = \sum_{j=1}^{N} \exp\left(\frac{Q_i K_j^\top}{\sqrt{d_k}}\right) \cdot V_j$
- 分母：$B = \sum_{m=1}^{N} \exp\left(\frac{Q_i K_m^\top}{\sqrt{d_k}}\right)$ 因分子和分母都是序列长度加和，所以分块之后计算也是每个分块的分子和分母分别加和：
$$
\frac{A}{B} = \frac{\sum_{i=1}^{n} A_i}{\sum_{i=1}^{n} B_i}
$$
实际应用：多卡在分别计算出 A_i 和 B_i 之后，可能会进一步计算出 attn_i = A_i / B_i，这是为了让 A_i 向量的值更小，减少精度损失。
举例验证
场景假设
- 序列长度 seq_len = 4
- 一个 q 向量 = [1, 1]
- K 矩阵 = [[1,0], [0,1], [1,1], [0,0]]
- V 矩阵 = [[1, 1], [2, 2], [3, 3], [4, 4]]
- 序列分为两个块
注：为方便计算，假设自然指数为 2，$\sqrt{d_k}$ 为 1。
不分块注意力计算
$$
\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^\top}{\sqrt{d_k}}\right)V
$$
- $QK^T$ = [1, 1, 2, 0]
- softmax 结果：[2/9, 2/9, 4/9, 1/9]
- 整体结果：[2/9, 2/9, 4/9, 1/9] * [[1, 1], [2, 2], [3, 3], [4, 4]] = [22/9, 22/9]
分块注意力计算
$$
A = \sum_{j=1}^{N} \exp\left(\frac{Q_i K_j^\top}{\sqrt{d_k}}\right) \cdot V_j
$$
-
A1
- $QK^T$ = [1, 1] * [[1, 0], [0, 1]] = [1, 1]
- exp = [2, 2]
- 整体结果： [2, 2] * [[1, 1], [2, 2]] = [6, 6]
-
A2
- $QK^T$ = [1, 1] * [[1, 1], [0, 0]] = [2, 0]
- exp = [4, 1]
- 整体结果： [4, 1] * [[3, 3], [4, 4]] = [16, 16]
$$
B = \sum_{m=1}^{N} \exp\left(\frac{Q_i K_m^\top}{\sqrt{d_k}}\right)
$$
-
B1
- $QK^T$ = [1, 1] * [[1, 0], [0, 1]] = [1, 1]
- 整体结果：4
-
B2
- $QK^T$ = [1, 1] * [[1, 0], [1, 0]] = [2, 0]
- 整体结果：5
$$
\frac{A}{B} = \frac{\sum_{i=1}^{n} A_i}{\sum_{i=1}^{n} B_i}
$$
- A/B = ([6, 6] + [16, 16]) / (4 + 5) = [22/9, 22/9]
不分块和分块计算的结果一致。
其他理解思路
（1）从矩阵的维度变化来理解为什么多个分块计算得到的分子向量要按位置加和
分块的核心计算分为两步：
- Q K 运算：[S, d] * [d * S] = [S, S]
- *V 运算：[S, S] * [S * d] = [S * d]
可以看到最后一步得到的结果是完整的形状 [S * d]，每一个 token 都有 d 个元素。细看最后一步 *V 的计算过程，是 S 次乘法得到 S 个元素后加和，而未分块时，这里是 N 次乘法得到 N 个元素加和，分块之后这里加和少了一部分元素，因而分块分别计算的结果需要按元素加和得到完整结果。
扩展：从 ring attention 到 Context Parallelism
ring attention 是从降低计算时的显存占用出发设计的
q 乘以 K的转置会产生 N（token 个数） 的平方元素，再乘以 V 会产生 N * d_model（hidden state 维度） 个元素。拆分 sequence 之后，每一个 卡 上变成 S 平方个元素 和 S * d_model 个元素，当 N >> S 时，峰值显存的占用会大幅度降低。
注：如果这些卡不用来做 ring attention ，而是用来做 TP，因为 TP 是按 HEAD 拆分，也就是 q 乘以 k 的转置这个矩阵计算的中间维度减少了，但计算结果的 N 的平方个中间元素还是会存在，最后输入时，每一个 TP 都需要完整的 hidden state，N*d_model 个元素因为无法减少，因此 TP 无法降低计算过程的峰值显存。
ring attention 会导致 TTFT 升高
当我们有超长的输入序列时，继续按 TP 模式拆分有几方面的缺陷：
（1）【峰值显存高】q 乘以 k 转置产出 N 的平方个中间元素，显存峰值很高；
（2）【拆分方式有限】TP 对权重矩阵做按列拆分，每一个卡只加载部分权重，拆分方式只能按 HEAD 个数去拆，有最大值限制，不能无限水平扩展；
（3）【通信耗时缺陷】假设按 TP 拆分为很多个 tp rank，更多个 tp 之间的 all-reduce 可能耗时会增加，主要来自几方面：参与通信的设备变多、整体通信量变多；
如果按照 ring attention 的思路，假设在极限情况下，一个 token 由一个 GPU 卡完成，对比 TP 模式：
（1）【单卡权重变大】ring attention 模式一个卡需要加载完整的权重，tp 模式可以只加载部分权重，因此 ring attention 模式权重占用的显存会增加；
（2）【矩阵计算持平】ring attention 的一个卡需要完成一个 token 的推理过程中的所有矩阵计算，而 tp 模式需要完成 N 个 token 的部分矩阵计算，这两计算量接近；
（3）【网络通信量变少】ring attention 模式推理过程中只有注意力计算的 k、v 传递涉及多卡通信，其他的计算都是本机完成；tp 模式下，所有矩阵计算都涉及网络通信。
（4）【计算耗时增加】ring attention 模式下的多个小块的注意力是串行计算的，在 N 个 token 分散到 N 个卡上执行时，每个卡都需要串行执行 N 次局部注意力计算。另外，局部 k、v 的传递也是串行的，但这部分时间和局部注意力计算 overlap 了，可以忽略；
CP：用显存换 TTFT 下降
计算耗时增加的优化思路：这部分耗时增加是由计算串行引起的，想办法把串行改成并行：局部 k、v 的流转修改为一次性的 all-reduce，这时候所有的卡都有完整的 k、v，而 q 是局部的，注意力计算一次性完成。这就是放弃了显存峰值优化，而获得耗时下降，这个思路假设叫 CP（Context Parallelism）模式，CP 模式和纯 TP 模式对比：
（1）CP 模式相比 TP 模式，单卡权重变大；
（2）CP 模式和 TP 模式的矩阵计算持平；
（3）CP 模式相比 TP 模式网络通信大幅度变少；通信频率减少，对于每一层来说，CP 只需要一次 all-gather，而 TP 至少需要两次 all-reduce：注意力计算结果、FFN 计算结果；通信量减少，CP 模式传的量是：[Batch_Size, Sequence_Length, Hidden_Size]，TP 模式传的量是：[Batch_Size, Sequence_Length, Hidden_Size]，单次传输的计算量是一样的，但因 TP 需要传两次，所以 TP 模式传的量更大。
（4）CP 模式和 TP 模式在注意力计算上，计算耗时持平：在单卡上，一次 k、v 的 all reduce 网络通信之后，一次性完成 sequence chunk 的注意力计算；对比 TP 多卡，多个卡都需要执行局部注意力计算，每个卡也是只执行一次注意力计算。
（5）显存峰值，CP 模式和 TP 模式一样存在显存峰值过高的风险。
（6）CP 模式在 GQA 模型结构下，可以降低计算量：TP 模式，如果 4 个 Query HEAD 共用一个 KV HEAD，而每一个 TP 负责一个 Query HEAD，那么 KV HEAD 会被重复计算 4 次。而 CP 模式，假设一个卡负责一个 token，那么这个卡在处理一个 token 时，只会算一次 KV HEAD。
（7）CP 模式在 MLA 模型结构下，可以降低计算量：即便是开启了 dp_attention ，一个 dp rank 下有 多个 TP，这些 TP 也都会算一次矩阵的潜在表示，这些是重复计算的。而CP 模式，也只会算一次。
（8）如果 CP 模式下，模型既不是 GQA，也不是 MLA，那就只有降低网络通信的优势。
TP VS CP 在网络通信方面的对比
CP 是把 TP 集群划分为多个称为 CP rank 的 TP 小集群，每个 TP 小集群完成一部分 token 的推理计算。CP 像是在走向 TP 的反面：尽量用更少的卡完成计算，而不要使用全局所有卡。纯 TP 模式，每一层推理至少需要两次全集群 all-reduce（注意力和 FFN），而 CP 模式只需要一次全集群 all-reduce（KV Cache）以及 CP rank 内部的两次 all-reduce（和 TP 模式类似），当 CP rank 是单机或者是网络更近的小集群时，CP 模式会有网络通信的优势。
ring attention 和 flash attention 的关系
flash attention 负责的是单卡上的分块计算，而 ring attention 实现的是多卡间的分块计算，因此他们两是两个层次的优化。对于超长序列来说，ring attention 是必要的，否则 KV Cache 都可能把显存用尽，而对于中等长度的序列，如果 KV Cache 无压力，则用 flash attention 就足够了。
总结：两者都可以优化注意力计算中间结果的显存占用，也可以正交的使用。除了优化中间显存占用，Ring Attention 还可以突破单卡显存限制，实现多卡分布式 KV Cache 扩展，而 flash attention 还可以优化计算速度。
浙公网安备 33010602011771号