企业官网建设流程全解析

1. 从 Dense GEMM 说起：为什么 MoE 不能直接套用稠密矩阵乘

1.1 Dense GEMM 的计算本质与硬件友好性

Dense GEMM（通用矩阵乘法）是整个深度学习推理和训练的算力基石。它的计算模式非常规整：给定两个稠密矩阵 A（M×K）和 B（K×N），输出 C（M×N），每个输出元素都是 A 的一行与 B 的一列做完整内积。这种规整性带来的最大好处是计算密度极高、访存模式可预测、Tensor Core 利用率容易打满。

我在实际 profiling 一个典型的 Transformer FFN 层时，Dense GEMM 的 SM 利用率通常能稳定在 70% 以上，Tensor Core 的 MMA 指令流水几乎不会断。原因很简单：所有 token 都要经过同一个权重矩阵，batch 维度足够大，K 维度足够深，硬件可以一次性把数据喂饱。

但 MoE（Mixture of Experts）架构把这个前提彻底打破了。MoE 的核心思想是：每个 token 不再经过一个共享的 FFN，而是由 Router（门控网络）动态选择 Top-K 个专家（Expert）来处理。这意味着不同 token 走的是不同的权重矩阵，原本一个大的 Dense GEMM 被拆成了若干个大小不一的小 GEMM。

1.2 MoE 带来的计算模式剧变

假设一个 MoE 层有 8 个专家，每个专家是一个独立的 FFN，Router 对每个 token 选出 Top-2 专家。那么对于一批 4096 个 token 的输入，理想情况下每个专家会分到大约 1024 个 token（4096×2/8）。但实际情况远没有这么均匀——Router 的负载倾斜是常态，热门专家可能分到 2000+ token，冷门专家可能只有几百个。

这就导致两个直接后果：

第一，每个专家的 GEMM 的 M 维度（token 数）是动态变化的，而且专家之间的 M 差异可能达到 3-5 倍。你没法像 Dense GEMM 那样预先确定一个漂亮的 tile 配置。

第二，如果逐个专家串行调用 GEMM，GPU 的并行度会被严重浪费。一个专家只有几百个 token 时，M 维度太小，SM 占用率极低，kernel launch 开销反而成了瓶颈。

所以 MoE 的核心工程挑战就变成了：如何把多个大小不一的专家 GEMM 高效地合并成一次或少数几次 kernel 调用。这就是 GroupedGEMM 要解决的问题，也是从 Dense GEMM 走向 TP MoE 的第一道坎。

注意：很多人第一次接触 MoE 时会想“不就是多个 FFN 嘛，循环调用不就行了”。实测下来，在 8 专家 Top-2 的配置下，朴素循环调用的端到端延迟比 GroupedGEMM 方案高出 2.5 倍以上，专家数量越多差距越大。

2. GroupedGEMM：把碎片化的小矩阵乘合并起来

2.1 GroupedGEMM 的核心思想

GroupedGEMM 的本质是一次 kernel launch 处理多个独立的 GEMM 问题。它接收一组矩阵对 (A_i, B_i)，每个问题有自己的 M_i、K_i、N_i，然后在一个 kernel 内部把这些计算任务分配给不同的 thread block 并行执行。

这跟 batched GEMM 的区别在于：batched GEMM 要求所有问题的 shape 完全一致（比如都是 M×K×N），而 GroupedGEMM 允许每个 group 的 shape 不同。MoE 场景下，每个专家的 N（hidden dim）和 K（hidden dim）通常是一样的，但 M（token 数）不同，所以 GroupedGEMM 天然适配。

我用过一个比较直观的类比：Dense GEMM 像是让一个班的学生做同一套试卷，老师可以统一收卷；MoE 像是每个学生做不同的试卷，如果还按人头一个个收，效率极低。GroupedGEMM 就是“按试卷类型分组收”，一次收完所有同类型的。

2.2 GroupedGEMM 在 MoE 中的具体实现路径

在 MoE 的前向计算中，GroupedGEMM 通常用在两个位置：

• Expert FFN 的 up projection：把每个专家分到的 token 与其对应的 up weight 做 GEMM
• Expert FFN 的 down projection：把中间激活与 down weight 做 GEMM

实现上，需要先做一步token 重排（permutation）：根据 Router 的输出，把所有 token 按照目标专家分组，同一专家的 token 在内存中连续排列。这个重排操作本身也有开销，但相比逐个专家调用的浪费，完全值得。

具体的数据流是这样的：

1. Router 计算每个 token 对每个专家的 gate 分数
2. 取 Top-K，得到每个 token 的目标专家列表
3. 统计每个专家的 token 数量，计算前缀和（prefix sum），确定每个专家在重排后 buffer 中的偏移
4. 按偏移把 token 散射（scatter）到对应位置
5. 调用 GroupedGEMM，传入每个专家的 M_i 和权重指针
6. 计算完成后，再把结果按原顺序收集（gather）回来，乘以 gate 权重并求和

这套流程里，第 3 步的前缀和计算和第 4 步的 scatter 是容易被忽视的性能热点。我见过不少实现把这两步用 Python 循环写，结果重排开销比 GEMM 本身还大。

2.3 GroupedGEMM 的参数选择与调优

GroupedGEMM 的性能高度依赖于 tile 配置。由于每个专家的 M_i 不同，固定的 tile size 很难对所有专家都最优。常见的做法是：

• 对 M_i 较大的专家，用较大的 tile（如 128×128）
• 对 M_i 较小的专家，用较小的 tile（如 64×64）或直接走 SIMT 路径
• 设置一个阈值，M_i 低于阈值的专家合并处理或走 fallback kernel

下面是一个典型的 GroupedGEMM 调用参数示例（以 CUDA 生态的常见接口风格为例）：

# 假设有 E 个专家，每个专家的 token 数存在 expert_sizes 中 # 权重指针数组和输出指针数组 grouped_gemm( A_ptrs, # 每个专家输入 token 的起始指针 B_ptrs, # 每个专家权重的起始指针 C_ptrs, # 每个专家输出的起始指针 expert_sizes, # 每个专家的 M_i N, # 输出维度（所有专家相同） K, # 输入维度（所有专家相同） workspace, # 临时空间 stream )

实操心得：expert_sizes 这个数组一定要放在 GPU 上，不要每步都从 CPU 传。我踩过的坑是每次 forward 都做一次 H2D 拷贝，在小 batch 场景下这个拷贝延迟占了总延迟的 15%。

3. TP MoE：当专家并行遇上张量并行

3.1 为什么 MoE 需要张量并行

GroupedGEMM 解决了单卡上多专家的高效计算问题，但当模型规模继续增大，单个专家的参数量就可能超过单卡显存。比如一个 8 专家的 MoE 层，每个专家是一个 4096×14336 的 FFN，光权重就占 4096×14336×2×2 bytes ≈ 235MB（FP16），8 个专家加起来接近 2GB，再加上 attention 和其他层，单卡很快就不够用。

这时候就需要TP（Tensor Parallelism，张量并行）：把每个专家的权重矩阵按列或按行切分到多张卡上，每张卡只持有部分权重，计算时通过 AllGather / ReduceScatter 等集合通信操作来同步中间结果。

TP MoE 的核心矛盾在于：MoE 本身是“专家并行”（不同 token 走不同专家），TP 是“权重切分”（同一个专家切到多卡）。两者叠加后，通信模式变得非常复杂。

3.2 TP MoE 的通信模式拆解

以一个典型的 TP=4、EP（Expert Parallel）=2 的配置为例，总共 8 张卡。每个专家被切到 4 张卡上（TP group），同时有 2 个专家分别放在不同的 EP group 里。

前向计算时，数据流大致是：

1. AllGather：每张卡上的 token 需要收集到所有 TP rank，因为每个 rank 只持有部分权重，需要完整输入才能计算部分输出
2. GroupedGEMM：在本地做部分权重的 GEMM
3. ReduceScatter：把各 rank 的部分结果求和并散射回各自的 token

这里的关键是：AllGather 和 ReduceScatter 的通信量跟 token 数成正比。token 越多，通信开销越大。而在 MoE 中，token 经过 Router 后会被分发到不同专家，如果专家分布在不同 EP group，还需要额外的 All-to-All 通信来把 token 送到目标专家所在的卡。

我实测过一个 TP=8、EP=1 的配置，在 4096 token 的 batch 下，AllGather + ReduceScatter 的通信时间占了 MoE 层总时间的 40% 以上。这个比例在 batch 更大时会下降，但在推理场景（batch 通常较小）下非常致命。

3.3 TP MoE 的权重布局与切分策略

TP MoE 的权重切分有两种主流策略：

策略一：按列切分（Column Parallel）

每个专家的 up projection 权重按输出维度切分到 TP 张卡上。每张卡计算部分输出，然后通过 AllGather 拼成完整结果。这种策略下，每张卡需要完整的输入 token，所以输入需要广播。

策略二：按行切分（Row Parallel）

每个专家的 down projection 权重按输入维度切分。每张卡计算部分和，然后通过 AllReduce 或 ReduceScatter 求和。这种策略下，输入也需要按维度切分。

实际实现中，通常 up projection 用列并行，down projection 用行并行，这样一次前向只需要一次 AllGather（up 之前）和一次 ReduceScatter（down 之后），通信次数最少。

权重布局上，我建议把同一专家的所有 TP shard 放在连续的显存地址上，这样 GroupedGEMM 的指针数组可以按专家顺序排列，减少指针跳转带来的 cache miss。

4. 从 Dense GEMM 到 TP MoE 的完整实操路径

4.1 环境准备与基础依赖

要复现一套 TP MoE 的前向计算，你需要：

• CUDA 环境：建议 12.0 以上，GroupedGEMM 相关的 cutlass 接口在新版本里更稳定
• cutlass：GroupedGEMM 的高效实现基本都基于 cutlass，建议用 3.x 版本
• 通信库：NCCL 2.18 以上，AllGather / ReduceScatter 的性能在新版本有显著优化
• PyTorch：2.1 以上，方便做 baseline 对比和调试

如果你只是想先理解原理，可以先用 PyTorch 的torch.matmul循环调用来做 baseline，确认数值正确后再替换成 GroupedGEMM。

4.2 第一步：实现 Router 与 Token 重排

Router 通常是一个简单的线性层 + Softmax + TopK。关键是要把 TopK 的结果转成 token 到专家的映射。

import torch import torch.nn.functional as F def route_tokens(hidden_states, gate_weight, top_k): # hidden_states: [num_tokens, hidden_dim] # gate_weight: [num_experts, hidden_dim] logits = F.linear(hidden_states, gate_weight) # [num_tokens, num_experts] scores = F.softmax(logits, dim=-1) topk_scores, topk_indices = torch.topk(scores, top_k, dim=-1) # 归一化 gate 权重 topk_scores = topk_scores / topk_scores.sum(dim=-1, keepdim=True) return topk_scores, topk_indices

重排的核心是计算每个专家的 token 偏移：

def compute_expert_offsets(topk_indices, num_experts): # 统计每个专家被选中的次数 expert_counts = torch.bincount(topk_indices.flatten(), minlength=num_experts) # 前缀和得到每个专家的起始偏移 offsets = torch.cumsum(expert_counts, dim=0) - expert_counts return offsets, expert_counts

注意：torch.bincount在 GPU 上的性能不如手写的 atomicAdd kernel，如果 token 数很大（>100k），建议自己写一个简单的 histogram kernel。

4.3 第二步：GroupedGEMM 的接入与调试

接入 GroupedGEMM 时，最容易出错的地方是指针数组的构造。每个专家的输入指针 = base_ptr + offset_i × hidden_dim，权重指针 = weight_base + i × expert_weight_size。

调试时建议先用小规模数据验证数值正确性：设置 2 个专家、Top-1、batch=4，手动计算每个专家的输出，跟 GroupedGEMM 的结果对比。确认无误后再放大规模。

一个常见的数值陷阱是：gate 权重的归一化方式。有些实现用 Softmax 后的原始分数，有些用 TopK 后重新归一化的分数。这两种方式在 Top-1 时等价，但在 Top-2 及以上时结果不同。一定要跟训练时的配置保持一致。

4.4 第三步：TP 通信的插入与重叠

TP MoE 的性能关键在于通信与计算的重叠。理想情况下，AllGather 可以在上一个 micro-batch 的 GEMM 计算时异步进行。

实现上，可以用 CUDA stream 来做通信和计算的重叠：

# 伪代码示意 comm_stream = torch.cuda.Stream() compute_stream = torch.cuda.current_stream() with torch.cuda.stream(comm_stream): allgather_handle = dist.all_gather_into_tensor_async(...) # 在 compute_stream 上做其他不依赖通信结果的计算 ... # 等待通信完成 allgather_handle.wait()

但要注意，GroupedGEMM 的输入依赖 AllGather 的结果，所以重叠的空间有限。真正有效的重叠策略是把不同专家的通信和计算交错开：专家 0 的 GEMM 在计算时，专家 1 的 AllGather 已经在进行。

4.5 第四步：端到端性能验证

验证时建议分三个层次：

如果单层性能提升不明显，优先检查 token 重排的开销和 GroupedGEMM 的 tile 配置。如果端到端通信占比过高，考虑增大 batch 或调整 TP/EP 比例。

5. 常见问题与排查技巧实录

5.1 专家负载不均导致的长尾延迟

这是 MoE 最经典的问题。Router 训练不充分时，80% 的 token 可能集中到 2-3 个专家，其余专家几乎空闲。GroupedGEMM 虽然能并行处理，但最慢的那个专家决定了整体延迟。

排查方法：打印每个专家的 token 数分布，计算 max/mean 比值。如果超过 3，说明负载严重不均。

解决思路有三个：一是加负载均衡损失（aux loss）重新训练 Router；二是推理时用 capacity factor 限制每个专家的最大 token 数，超出的 token 走 fallback 路径；三是用专家并行的方式把热门专家切到多卡上。

5.2 AllGather 通信量过大

TP MoE 中 AllGather 的通信量 = token 数 × hidden dim × 2 bytes（FP16）。在 TP=8、hidden=4096、token=8192 时，单次 AllGather 就是 512MB 的通信量。如果每层都做，累积起来非常可观。

优化手段：一是用 FP8 或 INT8 做通信压缩；二是把多个层的 AllGather 合并成一次；三是用 NVLink 全互联的拓扑，避免跨节点通信。

5.3 GroupedGEMM 的 workspace 溢出

GroupedGEMM 通常需要一块 workspace 来存放中间结果。如果 workspace 分配太小，会出现静默的错误结果（不报错但数值不对）。建议 workspace 大小按max_expert_tokens × N × sizeof(dtype) × 2来分配，留 2 倍余量。

5.4 常见问题速查表

最后分享一个我踩过的坑：GroupedGEMM 的指针数组在 CPU 上构造后传到 GPU，如果用了 pinned memory 但忘了同步，会出现随机性的错误结果。这个 bug 极难复现，因为只在特定时序下触发。后来我改成直接在 GPU 上用 kernel 构造指针数组，问题彻底消失。

这套从 Dense GEMM 到 TP MoE 的路径，核心就是把“碎片化计算”和“跨卡通信”这两个问题分别用 GroupedGEMM 和 TP 通信原语解决，再用重叠策略把通信开销藏起来。实际落地时，数值正确性和性能调优往往要反复迭代好几轮，但只要把 token 重排、指针构造、通信重叠这三个关键环节做扎实，整体方案就能稳定跑起来。