矩阵运算
随着模型参数的增长，推理 GPU 资源需求从 1/N 卡、单卡、多卡、多节点、再到超节点，规模不断增加，之前在训练里面常用的 DP/TP/SP(CP)/EP/PP/Zero 等并行方法也在推理中逐步使用起来。
推理的并行还有着自己的一些特点，比如 PD 分离场景下 P 和 D 的部署方式存在差异、PP 场景下 Attention 与 FFN 数量可以不等。
在展开推理并行优化前，对一些基础知识进行了解是必要的。
1. 矩阵运算
并行底层的运算涉及二维矩阵拆分运算、多维矩阵切分与合并运算。
在神经网络中线性运算基本是矩阵乘法：C = A x B。矩阵运算的输入、输出可以是多维矩阵（>2 维），但维度之间有两个关键约束：
- 若 A 尺寸为 [M, N]，B 的第一维度必须为 N。B 为 [N, K] 时，计算得到的 C 尺寸为 [M, K]。
- 多维矩阵的乘法都会变为二维矩阵的乘法，倒数第 1 和第 2 维必须满足第一条的 MNK 关系。
- 多维乘法认为是批量的二维矩阵乘法。
对第二点有个补充：多维矩阵乘除了最后两维，其它维度必须相等（或者进行广播），举个例子：A size: [10, 3, 2]、B size: [2, 4]，计算 A x B
- B 广播变为 [10, 2, 4]，
- 对应 Ai [i, :, :] 与 Bi [i, :, :] 进行相乘，
- 得到的 C 维度为 [10, 3, 4]。
知道尺寸的约束关系后，接下来分析矩阵的分块运算。
先看二维矩阵的运算切分场景。下面的举例中都进行的矩阵二分，矩阵的多次切分可以看成二分的扩展。
1.1. 场景一：B 矩阵列切 (column split)
B 矩阵二分（列切）后，子矩阵 B1、B2 的尺寸为 [N, K/2]，切分计算不改变局部元素值，但需要进行一次列维度的拼接才能获得完整矩阵 C。
特点：计算量减半，多了一次 allgather 集群通信；显存空间变化不一定，若是非原地操作，则变化为 (-K/2 * N + M * K/2) * element_memory。
应用：在并行运算中，通常用两个 GPU 来完成这一个拆分运算，比如当 A 矩阵为输入激活值，B 矩阵为权重，最后结果通过 allgather 进行拼接。
矩阵切分运算可用 numpy 做个辅助验证，这里提供一个示例：
import numpy as np
# 1. 定义整数输入矩阵 (M, N) 和 (N, K)
M, N, K = 3, 4, 6
A = np.random.randint(0, 10, size=(M, N)) # 随机整数矩阵 [0, 10)
B = np.random.randint(0, 10, size=(N, K))
print("A:\n", A, "\nshape:", A.shape)
print("\nB:\n", B, "\nshape:", B.shape)
# 2. 对 B 按列切分（均分）
num_splits = 3 # 切分块数
B_splits = np.split(B, num_splits, axis=1) # 沿列切分
print("\nB 分块结果:")
for i, B_i in enumerate(B_splits):
print(f"B_{i}:\n", B_i, "\nshape:", B_i.shape)
# 3. 模拟并行计算：每个进程计算 A @ B_i
local_results = [A @ B_i for B_i in B_splits]
print("\n局部乘积结果:")
for i, C_i in enumerate(local_results):
print(f"C_{i} (A @ B_{i}):\n", C_i, "\nshape:", C_i.shape)
# 4. 模拟 allgather：拼接所有局部结果
C_final = np.concatenate(local_results, axis=1)
print("\n合并后的 C_final:\n", C_final, "\nshape:", C_final.shape)
# 5. 验证结果与直接乘法的等价性
C_ground_truth = A @ B
print("\n标准乘法结果 (A @ B):\n", C_ground_truth)
print("\n验证一致性:", np.array_equal(C_final, C_ground_truth))
1.2. 场景二：A 矩阵行切 (row split)
A 矩阵二分（行切）后，子矩阵 A1、A2 的尺寸为 [M/2, N]，切分不改变局部元素值，但需要进行一次行维度的拼接才能获得完整矩阵 C。
特点：计算量减半，多了一次集群通信（allgather）；若是非原地操作，显存变化：(K-N) * M/2 * element_memory。
应用：这类切分比较常见，如 A 矩阵是激活值，B 矩阵为权重，如果后面的还是线性运算或者元素行运算，allgather 步骤可以放到最后且只需进行一次。
如：A[M/2, N] x B x E x F … -> allgather 与原计算得到结果一致。这就是常见的激活值进行序列切分的场景：[sequence/sp, hidden_dim]。
Note：这里有个值得注意的是，仅切分一个矩阵的方式，适合任何线性运算，但不一定能够带来性能/存储空间的收益。
1.3. 场景三：A 矩阵列切 + B 矩阵行切
A、B 矩阵都进行二分后，输出矩阵 C 的大小不变，所有 C 矩阵进行元素求和，才能让结果与之前计算方式相等。
特点：计算量减半，多了一次集群通信（allreduce）；显存变化：输入、权重减半，输出值取决于是否为原地操作。
这种切分方式的使用场景并不多，更多的是下面这种方式：
1.4. 场景四：B 矩阵列切 + C 矩阵行切
在 A x B x C 矩阵乘法场景，B 矩阵列切，C 矩阵行切。
特点：在场景三基础上，中间值的存储大小减半。
在模型并行运算的过程中，大部分都是这些应用场景的组合与变形，需要根据模型层的特点设计合适的矩阵切分策略。
2. 线性层的并行计算
线性层的并行运算，包括层内切分、层间切分以及参数冗余消除，并行方式包括了数据并行（DP）、序列并行 (SP)、张量并行（TP）、层并行（PP）、参数冗余消除（Zero）。
2.1. 层内并行
可根据输入激活值的切分维度对应不同并行策略，一般切 batch 为 DP、切序列为 SP、切隐藏层尺寸为 TP。
关键特点：
- DP 可以看成批量矩阵乘法，不会改变二维矩阵的乘法形态，所以对应的权重不需要变化。
- SP 切分不会影响最后一维的计算，如 LayerNorm、softmax。
- TP 切分激活值后，权重必须做对应的切分。
DP 与 SP/TP 可以解耦独立使用，SP 与 TP 一般可以配合使用（在当前推理应用中一般设计成：DP*TP=world_size，SP=DP）。
又因为：
- allreduce 是可以被拆分为 ReduceScatter 和 AllGather 两个步骤
- SP 不影响元素操作的运算，如 LayerNorm，softmax
所以在上面矩阵运算场景四的基础上，对于最后的集群通信操作有如下改进，将 allreduce 运算进行拆分，变为两个步骤。并且将 allgather 操作移动到 softmax 计算后：
这样能够降低元素运算相关计算的数据大小、降低过程中激活值的大小（示例中 GPU 上的 LayerNorm 输入减半），即能节省显存、也能提升速度。
在 megatron 论文里面介绍了这种场景的应用，MLP 层中用 SP-TP-SP 的组合方式：
2.2. 层间并行
模型的各个层分散放在不同 GPU 上面，激活值按照流水的方式在层间传递，这种方式称为流水线并行 PP(Pipeline Parallelism)，具体可参看 megatron2 论文。
在推理中，由于特殊计算场景，比如 FFN 为稀疏的 MoE 结构时，往往 attention 与 FFN 对算力需求不同，层间并行可以变为异构模式，如 AFD，attention 层与 MoE 分离部署到不同设备上：
2.3. 冗余参数消除
冗余参数消除的方式，如 DeepSpeed 的 Zero，将参数分散存储，当参数需要使用时通过集群通信将参数召回。
根据 Zero 的思想，这里介绍两种消除某一层权重冗余存储的方式。
一般而言大模型由多个 blocks 构成，blocks 中有重复的层结构。当有多个设备时，可以把这些 blocks 的权重均分到各个设备中。如下图所示，以 blocks 中线性层 B 为例，每个 GPU 存储模型的两层权重，当 GPU 需要使用到某一层权重时，先对该层的权重进行广播。
由于 GPU 支持多流，所以可以做计算、广播通信的掩盖，比如计算到 x 层时，提前把 x+1 层的权重先广播出去（权重预取）。
这种方式适合于模型层权重偏小的场景，模型层权重比较大时，或者说更常使用的是权重切片（shard）的场景。
即将每一层的权重进行 n 等分，每个 GPU 设备上面存一份，当需计算时将其 allgather 回来。
3. Attention 层的并行计算
Attention 层的并行计算与 FFN 或者说模型中其它线性层的计算的主要不同之处：
- 在计算 attention 之前，会多一个 head 的维度，这个 head 维度是在线性映射（投影）运算中产生的。[bs, seq_len, heads * head_dim] -> [bs, heads, seq_len, head_dim]
- 序列（seq_len 维度）并行需要修正运算。
在并行维度上面，将 Head 维度的切分也算到 TP 里面，而序列维度并行上称之为 CP(context parallel) 或 SP。
Attention 的层间并行、冗余参数消除方式与线性层的方式一致，层内并行的主要差异是 TP 和 SP，所以重点关注到这两类上面。
3.1. 序列并行
在经过线性映射计算后，模型输入变为多头的 Q、K、V 值，其尺寸表达为：[bs, heads, seq_len, head_dim]，序列并行是切分 seq_len 维度。
有个前提问题：QKV 的序列并行跟线性层的差异在哪，是否能直接拆分运算后合并？
关于 Scaled Dot-Product Attention，我们分三个子问题来分析：
- 如果只切 Q 的序列，KV 保持不变，结果拼接后是否相等？
- 如果只切 K 的序列，QV 保持不变，结果拼接后是否相等？
- 如果只切 V 的序列，QK 保持不变，结果拼接后是否相等？
先看问题 1，Q 切分后的尺寸为 [bs, heads, seq_len/SP, head_dim]，按照 attention 计算：
- step1：求解 score，Q x K 相当于左矩阵行切（场景二），score 尺寸：[bs, heads, seq_len/SP, seq_len]
- step2：softmax 求解的是最后一个维度，计算元素值相同，得到 attention_weights，
- step3：attention_weights 与 V 进行矩阵乘，还是左矩阵行切运算，元素值相同，计算得到 O 的分块结果
- step4：将计算的 O 进行 allgather，结果相等。
所以：单独切 Q 的序列值后进行结果拼接，与原计算相等。我们可以通过如下例子来验证：
import numpy as np
def softmax(x):
e_x = np.exp(x - np.max(x, axis=-1, keepdims=True))
return e_x / np.sum(e_x, axis=-1, keepdims=True)
def attention(Q, K, V):
scores = np.matmul(Q, K.transpose(0,1,3,2)) / np.sqrt(K.shape[-1])
weights = softmax(scores)
return np.matmul(weights, V)
# 设置随机种子（确保可复现）
np.random.seed(42)
# 定义形状：batch_size=1, num_heads=1, seq_len=4, hidden_dim=3
bs, heads, seq_len, hidden_dim = 1, 1, 4, 3
# 随机生成 Q, K, V
Q = np.random.randn(bs, heads, seq_len, hidden_dim)
K = np.random.randn(bs, heads, seq_len, hidden_dim)
V = np.random.randn(bs, heads, seq_len, hidden_dim)
# 全局注意力计算
global_out = attention(Q, K, V)
print("全局注意力输出:\n", global_out)
# 将 Q 按 seq_len 切分为两部分
Q1 = Q[:,:,:2,:] # 前两行
Q2 = Q[:,:,2:,:] # 后两行
# 分片计算注意力（K/V 保持完整）
local_out1 = attention(Q1, K, V)
local_out2 = attention(Q2, K, V)
# 合并分片结果（模拟 all_gather）
local_out = np.concatenate([local_out1, local_out2], axis=2)
print("分片注意力输出:\n", local_out)
# 验证一致性
print("结果是否一致:", np.allclose(global_out, local_out, atol=1e-6))
问题 2：单独切 K 的序列。
第一步，QK 乘法得到的 score 尺寸为：
[bs, heads, seq_len, seq_len/SP]
进一步计算 softmax，由于最后一个维度的数据只有之前的一半长度，而 softmax 的计算跟整个序列相关，直接拼接会导致结果不相等。
所以，单独切 K 序列后拼接，结果不等。
问题 3：单独切 V 的序列。
softmax 之前的计算保持不变，得到 attention_weights 尺寸完整，计算 attention_weights x V，因为 V 矩阵被行切，所以 attention_weights 需要列切，根据前面的场景三，最后 allreduce 能够获得完整结果。
综上看到 attention 序列切分的关键：如果需要进行 Q、K、V 在序列维度的同时切分，要解决 softmax 的分块计算。
当前比较常用的解决方式是：QKV 按照相同方式切分，然后对计算结果进行修正，这个方法在 Blockwise Multi-Head Attention、ring attention、tree attention、merge attention 中都有应用，softmax 修正：
关键点：
- 每个切分 Qi 要与所有的 Ki、Vi 进行一次计算，得到 Oi；
- 最后的分块 Oi 结果需要进行修正；
有了 softmax 的修正方式，结合推理的计算场景实现序列并行：
prefill：训练 forward/推理的 prefill 阶段，Q 的序列长度与 KV 保持一致，开启 SP 后 GPU 之间需要交换 KV 值与 Q 进行运算（ring-attention 方式），计算过程的 CP/SP 并行可以参看：[并行训练]Context Parallelism 的原理与代码浅析。至于是用 KV 传递还是 Q 传递，可以参看这篇分析。
decode：与 prefill 阶段不同的是，在 decode 阶段 Q 的长度为 1，若对应的 KV 值分散在各个 GPU 设备中，可以将 Q 复制 N 份与切片的 KV 值进行 attention 计算后，最后将结果 gather 到一个设备上再计算修正结果。
这个方式意味着 GPU0 需要完成最大值、修正运算、求和运算，GPU0 可能成为瓶颈。
一种 Tree attention 算法对过程做了改进，提升了通信效率。就是把 gather 换成多次 allreduce，这种方式在跨节点场景中优势明显。
prefix cache：在 prefill 阶段，还有个特殊场景 prefix cache，若开启了 prefix cache 功能，表示有一段 KV cache 可以复用。
把序列 seq_len 拆分成两部分讨论，prefix_cache_seq_len、tail_seq_len。Q/K/V 的激活值尺寸：
[bs, heads, prefix_cache_seq_len + tail_seq_len, head_dim]
其中 KV 值的 prefix_cache_seq_len 已被缓存，prefix_cache_seq_len KV 的计算量更少。Q 序列只需计算 tail_seq_len 的结果，加上前面的分析（Q 序列单独切分，不会改变元素计算结果）可知，计算分为两类场景：
- Q tail_seq_len 与 prefix_cache_seq_len KV 计算 attention；
- Q tail_seq_len 与 tail_seq_len KV 计算 attention；
在有 prefix cache 的序列切分时要考虑的：
- 问题 1：两类计算如何分配，保证负载均衡？
- 问题 2：已有的 prefix cache 如何从存储位置传递到计算位置（prefix cache 可能不在 GPU 上）。
这里举例一种简单的实现方式，将 prefix_cache_seq_len 和 tail_seq_len 进行均分，每个 GPU 拿取一份，同时进行 Q 分片的轮转，如下图所示：
Q-pass ring attention 优势：当切分比例保持不变且 KV 值每次计算后都在本地显存中，则无须进行 KV 分片的搬运（解决了问题 2）。
当 prefix cache 在其它存储位置时比如内存中，依然要考虑重新切分和 cache 的传输问题。
该方式在计算的时候可能需要进一步拆分，因为按照 SP 切分后的 KV 序列与 Q 序列不一定相等：prefix_cache_seq_len/SP != tail_seq_len/SP。若相等，理想情况下一次迭代计算中只需要两次 attention 运算。
3.2. 张量并行
Attention 的张量并行分为两个部分：QKV 运算、线性投影运算。
QKV 运算的 TP 切分一般针对 heads 维度，相当于矩阵的并行的批量运算。attention 计算切 heads，每个 rank(GPU) 拿到不同的 head，最后进行 heads 还原（allgather）；attention 前后的线性运算按照 hidden_dim/heads 方式切分。
Attention 的张量并行与其它层的并行可结合使用降低通信/计算量，这里给出 megatron 中 TP-SP 结合的例子：
目前的 DeepSeek 中用的 attention 模块 MLA 与 MHA 的计算在线性层的处理中有些不同，导致其线性层的 TP 的并行策略设计也要做对应的调整。
这里做一个详细的展开讨论，首先看一下 MLA 的计算公式：
MLA 相比 MHA 的主要特点是：线性运算多了一个步骤：在 QKV 生成之前，先要做降维 down_project 线性运算，然后进行升维 up_project 线性运算。分头运算是在 up_project 阶段。
根据公式绘制运算流图，非矩阵吸收的 MLA（KV up_project 运算在压缩计算后）如下所示，这种形态常用在 prefill 阶段（Q 的输入为完整序列时）。
MLA 一般形态主要的计算步骤：
- 1、Q、K、V 的下采样线性计算
- 2、Q 上采样运算
- 3、KV 上采样运算
- 4、attention 计算
- 5、O 线性运算
分别来看一下这些计算步骤的 TP 策略：
- 下采样阶段还没有 heads，所以一般而言 Q、K、V 的下采样不进行 TP 切分。考虑到独立切分适应任何场景，也可参考前面提的矩阵切分场景来分割下采样。
- 如果按照 heads 进行切分的话，步骤 2、3 的 W 矩阵需要进行列切 (column split)，步骤 5 的 W 矩阵行切 (row split)；
- 步骤 4 的 attention 则进行 heads 维度切分 (head split)。
MLA 还有第二种形态，就是将 KV 的 up_project 运算移动到 Q 上面，即 K、V 输入没有了 up_project 计算（多 heads 通过广播实现）。MLA 吸收矩阵形态的逻辑如下所示：
MLA（吸收矩阵）常用在 decode 中，这种形态的 TP 切分策略：
- 步骤 1 不采用 TP 切分；
- 步骤 2 的 W 矩阵需要进行列切 (column split)；
- 步骤 5 的 W 矩阵行切 (row split)；
- 步骤 4 的 attention 则进行 heads 维度切分 (head split)；
- 步骤 3 的两次线性运算均进行 heads 维度切分 (head split)；
注：MLA 的 KV 矩阵吸收运算（W_UK -> W_UQ W_UV -> W_O）是重新建立一个线性映射矩阵（即，两次的 W_UK 不相等），MLA 的矩阵吸收的证明见参考文献 7。
4. 附 1：常见的集群通信方式
5. 参考文献
- DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model
- 集群通信 - hccl
- [Megatron2]: Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM
- [Megatron3]: Reducing Activation Recomputation in Large Transformer Models
- Attention Is All You Need
- arXiv:1911.02972
- Ring Attention with Blockwise Transformers for Near-Infinite Context
- [CONTEXT PARALLELISM] 对 ring attention 实现方式分析：arXiv:2411.01783
- deepseek 技术解读 (1)- 彻底理解 MLA（Multi-Head Latent Attentio…
- [merged attention] FLASHINFER: EFFICIENT AND CUSTOMIZABLE ATTENTION ENGINE FOR LLM INFERENCE SERVING
- [tree-attention] arXiv:2408.04093
更多推理框架分析知识：知乎专栏《LLM 推理优化必备知识》。欢迎点赞、收藏、关注！