DeepSeek MoE架构与SGLang框架完全指南 - 大模型专家系统入门到精通！

本文详细解析了DeepSeek MoE架构的结构与推理流程，介绍了SGLang框架中MoE的实现方法。文章阐述了包含256个路由专家和1个共享专家的计算流程，包括Gate、Topk和Expert computation三个关键步骤。同时探讨了SGLang框架中的多种优化技术，如SBO单批次重叠、token重排和Triton Fused MoE优化策略，帮助读者深入理解MoE架构的工作原理和性能优化

python零基础入门小白

241人浏览 · 2025-11-17 09:54:52

python零基础入门小白 · 2025-11-17 09:54:52 发布

一、什么是 MoE？

Mixture of Experts（MoE）是一种神经网络架构设计，通过组合多个”专家”（子网络）来处理输入数据。

其核心思想是：不是让整个模型处理每个输入，而是为每个输入动态选择最合适的专家子集。这种设计使模型能够在保持计算效率的同时，显著扩大参数规模。

本文简略梳理 DeepSeek MoE 结构的推理流程以及主流 SGLang（v0.5.5）框架是如何进行推理，希望能帮到大家~如有错误，欢迎各位大佬指正。

二、MoE 的结构

MoE deepseek

以 DeepSeek MoE layer 为例，其包含 256 个路由专家和 1 个共享专家，每个 Token 选择 Top-K = 8 个路由专家再加权求和，共享专家独立于路由网络。

三、计算流程

参考 SGLang 的 sglang/srt/models/deepseek_v2.py/ 中 class DeepseekV2MoE 实现。

总体流程如下所示：

SGLang - deepseekv2.py

Gate

gate 是一个 linear 线性层，接收输入 hidden_states 并且为每个专家计算一个原始得分 logits。

输入：

hidden_states shape = [num_tokens, hidden_dim]

输出：

router_logits shape = [num_tokens, n_experts]

# native代码为：
logits = F.linear(hidden_states, self.weight, None)

Topk

接收 router_logits 并决定每个 token 实际去哪个专家，并计算相对应的权重。

输入：

hidden_states shape = [num_tokens, n_experts]

输出：

topk_weights shape = [num_tokens, K], topk_ids shape = [k]

native top-k 代码的计算过程：

def fused_topk_torch_native(
hidden_states: torch.Tensor,
gating_output: torch.Tensor,
topk: int,
renormalize: bool,
correction_bias: torch.Tensor = None, # 偏置项，用于调整专家选择的公平性
):
if correction_bias is not None:
n_router_experts = gating_output.shape[-1]
scores = gating_output.softmax(dim=-1)
# 添加bias
scores_for_choice = scores_view(-1, n_router_experts) + correction_bias.unsqueeze(0)
# 计算top_k
topk_ids = torch.topk(scores_for_choice, k=topk, dim=-1, sorted=False)[1]
# 取出top-k 对应的权重
# scores.shape = [M, E]
# topk_ids.shape = [M, k], 专家编号
# 比如
#    scores = [[0.05, 0.10, 0.75, 0.10]]   # softmax 后
#    topk_ids = [[2, 0]]                   # 选专家 2 和 0
#    topk_weights = scores.gather(1, topk_ids)
#    结果 [[0.75, 0.05]]
topk_weights = scores.gather(1, topk_ids)
else:
assert (
hidden_states.shape[0] == gating_output.shape[0]
), f"Number of tokens mismatch, {hidden_states.shape=} vs {gating_output.shape=}"
M, _ = hidden_states.shape
topk_weights = torch.empty(
M, topk, dtype=torch.float32, device=hidden_states.device
)
topk_ids = torch.empty(M, topk, dtype=torch.int32, device=hidden_states.device)
topk_weights = F.softmax(gating_output.float(), dim=-1)
topk_weights, topk_ids = torch.topk(topk_weights, topk, dim=-1)
if renormalize:
topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
return topk_weights, topk_ids

DeepSeek V2/V3/R1 系列模型使用 grouped_topk 机制，每个 token 只能选择固定数量的专家组，然后在每个专家组内选择 top-k 个专家。

具体的步骤如下所示：

按 expert group 分组，选出 top-k groups，这里代码中一般为 2 个 group
在选中的 groups 内部再选 top-k experts
支持融合共享专家
返回每个 token 的 top-k 专家的 ID 和对应的 topk_weights 代码实现还是和之前一样，只不过将 scores_for_choice 这个变量进行了 group 维度切分，即

def bias_grouped_topk_impl(
hidden_states: torch.Tensor,
gating_output: torch.Tensor,
correction_bias: torch.Tensor,
topk: int,
...):
# 1. 用sigmoid对每个专家独立激活
scores = gating_output.sigmoid()
num_token = scores.shape[0]
num_experts = scores.shape[1]
scores_for_choice = scores.view(num_token, -1) + correction_bias.unsqueeze(0)
# 2. group
# 对每个group内部取top-2 experts的score并求和 -> 得到该group的"总强度"
# 这是一种 group importance scoring策略(避免只看max)
group_scores = (
scores_for_choice.view(num_token, num_expert_group, -1)
.topk(2, dim=-1)[0]
.sum(dim=-1)
)  # [n, n_group]
# 3. 选择top-k groups
group_idx = torch.topk(group_scores, k=topk_group, dim=-1, sorted=False)[1] # [n, top_k_group]
# 4. 构建mask， 只在选中的groups中选择expert
group_mask = torch.zeros_like(group_scores)  # [n, G]
group_mask.scatter_(1, group_idx, 1)  # 标记选中的 groups
score_mask = (
group_mask.unsqueeze(-1)
.expand(num_token, num_expert_group, scores.shape[-1] // num_expert_group)
.reshape(num_token, -1))  # [n, e]
# 5. 在masked区域内做final top-k
#把未选中group的scores设为
-inf， 确保不会被选择
#从有效区域内选择topk
tmp_scores = scores_for_choice.masked_fill(~score_mask.bool(), float("-inf"))
_, topk_ids = torch.topk(tmp_scores, k=topk, dim=-1, sorted=(True if num_fused_shared_experts > 0 else False))
topk_weights = scores.gather(1, topk_ids)
...

Expert computation

moe 的实现有很多种，目前 SGlang 实现了如下三种：

DeepEPMoE
FusedMoE
FlashInferFP4MoE

在 EP 并行的模式下，其主要的流程都为：

（1）dispatch

将输入 token 分发到指定的 rank 上的专家，具体实现即为 all2all 通信，其基本的流程都如下所示：

init：设备数据 buffer 的初始化，预留缓冲区
token 计算：计算向其他 rank 的 expert 发送的 token 数量，以及要接收到的 token 数量，并且检查 token 是否在 rank 中
更新缓冲区：将 num_token_per_rank，num_token_per_expert 即每个 token 放到发送缓冲区里面
执行 All2ALL：数据发送到目标 rank，并且接受其他设备的数据后将其分配到本地专家

（2）run_moe_core

SGLang 结合不同的优化方法和硬件类别，实现了多个 MoeRunnerBackend ，目前 FP8MoEMethod 支持 cutlass、deep_gemm、TritonMoe 的实现。

（3）combine

将专家的输出组合成最终结果。

（4）shared expert forward

执行共享专家的计算，deepseek 为一个，这里理论上可以与通信的过程 overlap，因为 shared expert 的结果和通信 dispatch 的数据没有关联。

计算之后，与路由专家的结果相加，如果做了 TP 并行的话，最后还会执行一个 all-reduce 通信操作来得到最终的 MoE layer 的输出结果。

四、优化点

SBO（single batch overlap）

在使用 DeepEP 的情况下，可以使用 SBO 对 shared expert 推理和 dispatch 的通信过程进行 overlap，相关 Issue：# Single Batch Overlap for MoE Models

https://github.com/sgl-project/sglang/pull/9660

SBO

目前 sglang 在 sglang/python/sglang/srt/single_batch_overlap.py 中实现了该方法，其流程为：

分发 tokens 到 routed experts。
执行 routed MoE 计算（up + down GEMM）。
同时（overlap）：

在部分 SM 上执行 shared experts（通过限制 GEMM 使用的 SM 数）
在剩余 SM 上准备/执行 combine 操作

通过 CUDA Event 确保 combine 在 routed MoE 完成后才开始。
最终输出融合结果通过手动的控制 down_gemm、combine 在不同的 stream 上使用的 SM 数量来实现 overlap , 相关代码如下：

def execute_sbo(
forward_shared_experts: Callable[[], Any],  # 共享专家前向函数
experts: FusedMoE,
hidden_states: torch.Tensor,
topk_output: TopKOutput,
alt_stream: Optional[torch.cuda.Stream] = None,
disable_sbo: bool = False,
):
# Step1 Dispatch, 分发token到experts
dispatch_output = experts.dispatcher.dispatch(...) # 将hidden_states 按照top-k分发的各rank上的expert buffer
# Step2 计算overlap参数
# 1. 分配使用多少个 SM 分别给通信和计算
# 2. 创建cuda event和stream
# 3. 是否启用overlap
combine_overlap_args, down_gemm_overlap_args, meta_overlap_args = _compute_overlap_args(...)
# Step3 执行routed MoE Forward
# 这里如果启用down_gemm_overlap_args,使用自定义的cute dsl kernel(flashinfer_cutedsl_moe_masked）进行推理，实现GEMM 和 combine的overlap
combine_input = experts.run_moe_core(
dispatch_output,
down_gemm_overlap_args=down_gemm_overlap_args
)
# Step4 异步启动shared experts
if (not disable_sbo) and SboFlags.enable_combine_shared_two_stream_overlap():
with deep_gemm_wrapper.configure_deep_gemm_num_sms(compute_num_sms):
forward_shared_experts()  # ← 在受限 SM 上异步执行！
# Step 5 执行combine
# 在alt_stream执行
hidden_states = experts.dispatcher.combine( combine_input=combine_input, overlap_args=combine_overlap_args, # 包含 alt_stream, event, signal 等 )

token 重排

每个专家处理的 token 不是连续的，现在是 topk_ids 的shape = （num_tokens，num_experts）。

即每个专家要处理的 token 在内存中不是连续排布的在 MOE 中：

每个 token 可能被分配给多个 expert
每个 expert 只处理分配给它的 token
如果这些 token 在原始输入 x 中是分散的，比如 token（0，5，100 都去 expert 3）

那么：

GPU 无法高效的批处理
内存访问不连续，带宽利用率低
难以启动一个大的、高效的 kernel 所以可以将要发给 expert 0 的 tokens 放一起，expert 1 的放一起，就能提高性能

解决方法：

对 hidden_states 重排，假设现在有如下 topk_ids：

# 每个token 两个experts
topk_ids = torch.tensor([
[3, 1],   # token 0 → expert 3 和 1
[0, 2],   # token 1 → expert 0 和 2
[1, 3],   # token 2 → expert 1 和 3
[0, 1],   # token 3 → expert 0 和 1
[2, 0],   # token 4 → expert 2 和 0
[3, 2],   # token 5 → expert 3 和 2
])

将（6，2）的 shape 拉成一维

flat_idx = topk_ids.view(-1)
# => [3, 1, 0, 2, 1, 3, 0, 1, 2, 0, 3, 2]
# ↑  ↑  ↑  ↑  ↑  ↑  ↑  ↑  ↑  ↑  ↑  ↑
#t0e0
t0e1 t1e0 t1e1 t2e0 t2e1 ...

argsort() -> 得到排序索引，即按照专家大小来排序

# 这里返回的是：如果对flat_ids排序，每个元素原来的位置在哪
idxs = flat_idx.argsort()
# 按 value 排序后（从小到大）：
# value:  0  0  0  1  1  1  2  2  2  3  3  3
# index:  2  6  9  1  4  7  3  8 11  0  5 10
# 故idxs = [2, 6, 9, 1, 4, 7, 3, 8, 11, 0, 5, 10]

映射回原始 token ID

token_indices = idxs // 2
# idxs = [2, 6, 9, 1, 4, 7, 3, 8,11, 0, 5,10]
# //2  = [1, 3, 4, 0, 2, 3, 1, 4, 5, 0, 2, 5]

token_indices 表示：排序后的每个“token-expert对”，对应的原始 token 是哪个

获取 sorted_tokens

sorted_token = x[token_indices]

即：

[
x[1],  # token 1 → expert 0
x[3],  # token 3 → expert 0
x[4],  # token 4 → expert 0
x[0],  # token 0 → expert 1
x[2],  # token 2 → expert 1
x[3],  # token 3 → expert 1
x[1],  # token 1 → expert 2
x[4],  # token 4 → expert 2
x[5],  # token 5 → expert 2
x[0],  # token 0 → expert 3
x[2],  # token 2 → expert 3
x[5],  # token 5 → expert 3
]

获取每个专家被多少个 tokens 选中

# cnts =
# token 0: [3,1] → 第0行，第3列和第1列设为1 → [0,1,0,1,0]
# token 1: [0,2] → 第1行，第0列和第2列设为1 → [1,0,1,0,0]
# token 2: [1,3] → 第2行，第1列和第3列设为1 → [0,1,0,1,0]
# token 3: [0,1] → 第3行，第0列和第1列设为1 → [1,1,0,0,0]
#所以
cnts =
#[[0, 1, 0, 1, 0],
# [1, 0, 1, 0, 0],
# [0, 1, 0, 1, 0],
# [1, 1, 0, 0, 0]]
tokens_per_experts = cnts.sum(dim=0)
#[2,3,1,2,0]

循环遍历每个专家，进行推理

# tokens_per_experts.shape = (5, 1) = (num_experts, 1)
for i, num_tokens in enumerate(tokens_per_experts):
end_idx = start_idx + num_tokens
if num_tokens == 0:
continue
tokens_for_this_experts = sorted_tokens[start_idx : end_idx]
# 用expert i 的权重计算
layer_w13_weight = layer.w13_weight[i]
layer_w2_weight = layer.w2_weight[i]
gate_up = F.linear(tokens_for_this_expert, layer_w13_weight)
gate_up = act(gate_up)
expert_out = F.linear(gate_up, layer_w2_weight)
outputs.append(expert_out)
start_idx = end_idx

combine - 聚合所有的 outputs 结果

outs = torch.cat(outputs, dim=0) if len(outputs) else \                                             sorted_tokens.new_empty(0)
# 创建空tensor， 并用idxs 还原到原始的x排序
new_x = torch.empty_like(outs)
new_x[idxs] = outs

乘以路由权重

new_x.view(*topk_ids.shape, -1) \   # (N, K, D_out)
.type(topk_weights.dtype) \   # weights (N,K)
.mul_(topk_weights.unsqueeze(dim=-1)) \
.sum(dim=1) \     # 对k维度求和 -> 得到每个token的最终输出
.type(new_x.dtype) # 恢复dtype

Triton Fused MoE 优化策略

融合 gate_proj 与 up_proj，两者输入都是相同 hidden_states，可以在 ffn_dim 上将权重 cat 到一起，然后 chunk

class SiluAndMul(nn.Module):
def __init__(self):
super().__init__()
@torch.compile
def forward(self, x:torch.Tensor) -> torch.Tensor:
x, y = x.chunk(2,-1)
return F.silu(x) * y

融合不同专家的 gate_proj/up_proj/down_proj，如将八个 down_proj 计算放在一起
将 softmax、topk 等小算子融合

SGLang kernel

moe_fused_gate.cu 代码走读：

https://github.com/sgl-project/sglang/blob/main/sgl-kernel/csrc/moe/moe_fused_gate.cu

这里 kernel 的输入为经过 router 的 linear 层计算得到的结果。

详细的代码可以参考 bbuf 大佬的解读，有非常完整的说明：图解 DeepSeek V3 biased_grouped_topk cuda 融合算子 fused_moe_gate kernel

https://zhuanlan.zhihu.com/p/1895178845830771205

DeepGemm kernel

核心代码链接 Todo：代码解读：

https://github.com/deepseek-ai/DeepGEMM/blob/03d0be3d2d03b6eed3c99d683c0620949a13a826/deep_gemm/include/deep_gemm/fp8_gemm.cuh#L40

五、如何系统的学习大模型 AI ？

由于新岗位的生产效率，要优于被取代岗位的生产效率，所以实际上整个社会的生产效率是提升的。

但是具体到个人，只能说是：

“最先掌握AI的人，将会比较晚掌握AI的人有竞争优势”。

这句话，放在计算机、互联网、移动互联网的开局时期，都是一样的道理。

我在一线互联网企业工作十余年里，指导过不少同行后辈。帮助很多人得到了学习和成长。

我意识到有很多经验和知识值得分享给大家，也可以通过我们的能力和经验解答大家在人工智能学习中的很多困惑，所以在工作繁忙的情况下还是坚持各种整理和分享。但苦于知识传播途径有限，很多互联网行业朋友无法获得正确的资料得到学习提升，故此将并将重要的AI大模型资料包括AI大模型入门学习思维导图、精品AI大模型学习书籍手册、视频教程、实战学习等录播视频免费分享出来。

一直在更新，更多的大模型学习和面试资料已经上传带到CSDN的官方了，有需要的朋友可以扫描下方二维码免费领取【保证100%免费】👇👇

在这里插入图片描述