跳转至

MoE Marlin 计算流程示例

本文用一个小例子说明 Int4MoePlugin 中 MoE Marlin 路径的完整计算过程，包括：router topK、expert 重排、padding、Marlin Gate/Up GEMM、SwiGLU、Marlin Down GEMM、slot 聚合。

1. 示例参数

假设：

Text Only
1
2
3
4
5
6
7
numTokens = 3
 topK = 2
totalSlots = numTokens * topK = 6
numExperts = 3
moeBlockSize = 4 (代码中实现，decode是8， prefill是32)
hiddenSize = H
moeInterSize = I

每个 token 会选出 topK=2 个 expert，因此一共有 6 个 routed slot：

Text Only
1
2
3
token0: slot0 = token0,k0   slot1 = token0,k1
token1: slot2 = token1,k0   slot3 = token1,k1
token2: slot4 = token2,k0   slot5 = token2,k1

slot 和 token 的关系是：

Text Only
1
2
3
slot = tokenId * topK + k
tokenId = slot / topK
k = slot % topK

2. Router TopK 输出

假设 router softmax + topK 后得到如下路由结果：

Text Only
1
2
3
4
5
6
slot0 = token0,k0 -> expert1, weight = w0
slot1 = token0,k1 -> expert0, weight = w1
slot2 = token1,k0 -> expert1, weight = w2
slot3 = token1,k1 -> expert1, weight = w3
slot4 = token2,k0 -> expert0, weight = w4
slot5 = token2,k1 -> expert2, weight = w5

对应数组：

Text Only
1
2
topkIndices = [1, 0, 1, 1, 0, 2]
topkWeights = [w0, w1, w2, w3, w4, w5]

这里的 topkIndices[slot] 表示该 slot 被路由到哪个 expert，topkWeights[slot] 表示该 slot 的 routing 权重。

3. 统计每个 expert 的 slot 数

launchCountSlotsPerExpertKernel 会统计每个 expert 被分到多少 slot：

Text Only
1
2
3
expert0: slot1, slot4       -> count = 2
expert1: slot0, slot2, slot3 -> count = 3
expert2: slot5              -> count = 1

得到：

Text Only
1
slotsPerExpertWorkspace = [2, 3, 1]

4. 计算 padded count 和 offset

Marlin MoE GEMM 要求每个 expert 的 slot 数按照 moeBlockSize 对齐。

这里：

Text Only
1
moeBlockSize = 4

所以：

Text Only
1
2
3
expert0: count=2 -> paddedCount=4
expert1: count=3 -> paddedCount=4
expert2: count=1 -> paddedCount=4

得到：

Text Only
1
paddedCounts = [4, 4, 4]

然后对 paddedCounts 做 exclusive prefix sum，得到每个 expert 在 sortedTokenIds 中的起始 offset：

Text Only
1
2
3
paddedOffsets[0] = 0
paddedOffsets[1] = 4
paddedOffsets[2] = 8

因此：

Text Only
1
2
paddedOffsets = [0, 4, 8]
numTokensPostPadded = 12

最终 padded 后的 Marlin row 总数是 12，而真实 slot 数只有 6。多出来的是 padding row。

5. 构建按 expert 分桶的 slot list

launchBuildSlotListsKernel 会遍历所有 slot，根据 topkIndices[slot] 把 slot 放入对应 expert 的 list：

Text Only
1
2
3
expert0: [1, 4]
expert1: [0, 2, 3]
expert2: [5]

这一步输出到 slotsByExpertWorkspace，逻辑上是：

Text Only
1
slotsByExpertWorkspace[expertId][offset] = slot

实际内存是一维数组：

Text Only
1
slotsByExpertWorkspace[expertId * totalSlots + offset] = slot

每个 expert 预留 totalSlots 个位置，因为最坏情况下所有 slot 都可能路由到同一个 expert。

6. 构建 Marlin 使用的 sortedTokenIds / expertIds

launchBuildMarlinIndicesKernel 会根据 slotsByExpertWorkspace、paddedCounts、paddedOffsets 生成 Marlin GEMM 需要的索引。

真实 slot 写在前面，padding 位置写 sentinel：

Text Only
1
sentinel = totalSlots = 6

重排后：

Text Only
1
2
3
4
sortedTokenIds:
expert0: [1, 4, 6, 6]
expert1: [0, 2, 3, 6]
expert2: [5, 6, 6, 6]

展开成一维：

Text Only
1
sortedTokenIds = [1, 4, 6, 6,  0, 2, 3, 6,  5, 6, 6, 6]

其中 6 == totalSlots，表示 padding row，不对应真实 slot。

同时每 moeBlockSize=4 个 row 是一个 Marlin block，每个 block 对应一个 expert：

Text Only
1
2
3
rows 0..3   -> expert0
rows 4..7   -> expert1
rows 8..11  -> expert2

所以：

Text Only
1
expertIds = [0, 1, 2]

topkWeightsFlat 也会按 sortedTokenIds 同步重排：

Text Only
1
topkWeightsFlat = [w1, w4, 0, 0,  w0, w2, w3, 0,  w5, 0, 0, 0]

padding row 的 weight 为 0。

7. 第一次 Marlin GEMM：Gate/Up Projection

第一次调用：

Text Only
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
moeAwqW4A16MarlinGemm(
    hiddenStates2D,
    gateUpOutputTensor,
    gateUpQWeightsTensor,
    gateUpScalesTensor,
    sortedTokenIdsTensor,
    expertIdsTensor,
    numTokensPostPaddedTensor,
    topkWeightsFlatTensor,
    marlinWorkspaceTensor,
    moeBlockSize,
    topK,
    mulTopkWeights = false)

它计算的是：

Text Only
1
gate_up[slot] = hidden[token] @ W_gate_up[expert]

其中：

Text Only
1
2
3
hiddenStates2D shape      = [numTokens, hiddenSize]
gateUpOutputTensor shape  = [totalSlots, 2 * moeInterSize]
W_gate_up[expert] shape   = [hiddenSize, 2 * moeInterSize]

Marlin 处理过程可以理解为：

Text Only
1
2
3
4
5
6
7
8
for each paddedRow in [0, numTokensPostPadded):
    slot = sortedTokenIds[paddedRow]
    if slot == totalSlots:
        continue  // padding row

    token = slot / topK
    expert = expertIds[paddedRow / moeBlockSize]
    gateUpOutput[slot] = hidden[token] @ W_gate_up[expert]

套入本例：

Text Only
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
padded row 0: slot1 -> token0 -> expert0
gateUpOutput[1] = hidden[token0] @ W_gate_up[expert0]

padded row 1: slot4 -> token2 -> expert0
gateUpOutput[4] = hidden[token2] @ W_gate_up[expert0]

padded row 4: slot0 -> token0 -> expert1
gateUpOutput[0] = hidden[token0] @ W_gate_up[expert1]

padded row 5: slot2 -> token1 -> expert1
gateUpOutput[2] = hidden[token1] @ W_gate_up[expert1]

padded row 6: slot3 -> token1 -> expert1
gateUpOutput[3] = hidden[token1] @ W_gate_up[expert1]

padded row 8: slot5 -> token2 -> expert2
gateUpOutput[5] = hidden[token2] @ W_gate_up[expert2]

padding row 不产生有效 slot output。

最终 gateUpOutputTensor 仍然按原始 slot id 排列：

Text Only
1
slot0, slot1, slot2, slot3, slot4, slot5

这点很重要：Marlin 内部按 expert 分组计算，但输出写回原始 slot layout，方便后续算子继续按 slot = tokenId * topK + k 使用。

8. SwiGLU 激活

第一次 GEMM 的输出每行包含 gate 和 up 两部分：

Text Only
1
2
gateUpOutput[slot] = [gate, up]
shape = [totalSlots, 2 * moeInterSize]

swiGluActivation 计算：

Text Only
1
activation[slot, i] = silu(gate[slot, i]) * up[slot, i]

其中：

Text Only
1
silu(x) = x / (1 + exp(-x))

输出：

Text Only
1
activationOutputTensor shape = [totalSlots, moeInterSize]

继续本例：

Text Only
1
2
3
4
activation[0] = silu(gateUpOutput[0].gate) * gateUpOutput[0].up
activation[1] = silu(gateUpOutput[1].gate) * gateUpOutput[1].up
...
activation[5] = silu(gateUpOutput[5].gate) * gateUpOutput[5].up

9. 第二次 Marlin GEMM：Down Projection

第二次调用计算 down projection：

Text Only
1
downOutput[slot] = W_down[expert](activation[slot])

在代码中第二次调用传入：

Text Only
1
2
topK = 1
mulTopkWeights = true

所以实际输出是：

Text Only
1
downOutput[slot] = topkWeight[slot] * (activation[slot] @ W_down[expert])

为什么这里才乘 routing weight？

MoE FFN 的数学形式是：

Text Only
1
2
expert_out = W_down(silu(W_gate x) * W_up x)
final = sum(topk_weight * expert_out)

如果在 gate/up GEMM 后就乘 routing weight，会变成：

Text Only
1
silu(weight * gate) * (weight * up)

这会改变 SwiGLU 的非线性结果，所以第一次 Gate/Up GEMM 必须：

Text Only
1
mulTopkWeights = false

第二次 Down GEMM 才能：

Text Only
1
mulTopkWeights = true

本例中：

Text Only
1
2
3
4
5
6
downOutput[0] = w0 * (activation[0] @ W_down[expert1])
downOutput[1] = w1 * (activation[1] @ W_down[expert0])
downOutput[2] = w2 * (activation[2] @ W_down[expert1])
downOutput[3] = w3 * (activation[3] @ W_down[expert1])
downOutput[4] = w4 * (activation[4] @ W_down[expert0])
downOutput[5] = w5 * (activation[5] @ W_down[expert2])

输出 shape：

Text Only
1
downOutputTensor shape = [totalSlots, hiddenSize]

10. 聚合 slot 输出回 token

最后调用 launchAggregateSlotOutputsKernel，把同一个 token 的 topK 个 slot 加起来：

Text Only
1
output[token, dim] = sum_{k=0}^{topK-1} downOutput[token * topK + k, dim]

本例：

Text Only
1
2
3
output[token0] = downOutput[slot0] + downOutput[slot1]
output[token1] = downOutput[slot2] + downOutput[slot3]
output[token2] = downOutput[slot4] + downOutput[slot5]

代入权重和 expert：

Text Only
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
output[token0] =
    w0 * (activation[0] @ W_down[expert1]) +
    w1 * (activation[1] @ W_down[expert0])

output[token1] =
    w2 * (activation[2] @ W_down[expert1]) +
    w3 * (activation[3] @ W_down[expert1])

output[token2] =
    w4 * (activation[4] @ W_down[expert0]) +
    w5 * (activation[5] @ W_down[expert2])

聚合后输出 shape：

Text Only
1
output shape = [numTokens, hiddenSize]

在 plugin 中这个 shape 对应原始 hidden states 的 [B, S, hiddenSize]。

11. 和普通 GEMM 的区别

普通 dense GEMM 是：

Text Only
1
C = A @ B

MoE Marlin GEMM 是：

Text Only
1
2
3
4
for each routed slot:
    token = slot / topK
    expert = expertIds[paddedRow / moeBlockSize]
    C[slot] = A[token] @ B[expert]

它多了三层机制：

  1. 按 expert 重排：把同一个 expert 的 slot 放在一起，提高 expert 权重访问局部性。
  2. 按 moeBlockSize padding：让 Marlin 每个 block 都处理规整数量的 routed slots。
  3. slot id 映射：内部按 expert 分组算，但输出写回原始 slot layout，方便后续 SwiGLU、Down GEMM 和最终 aggregate。
  4. INT4 packed weight + scale：Marlin 内部对 AWQ INT4 权重边反量化边做 W4A16 GEMM。

12. 完整流程总览

Text Only
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
hiddenStates [numTokens, hiddenSize]
    |
    | router logits -> topK
    v
topkIndices/topkWeights [numTokens, topK]
    |
    | count slots per expert
    v
slotsPerExpertWorkspace
    |
    | pad counts + exclusive scan offsets
    v
paddedCounts / paddedOffsets / numTokensPostPadded
    |
    | build slot lists by expert
    v
slotsByExpertWorkspace
    |
    | build Marlin indices
    v
sortedTokenIds / expertIds / topkWeightsFlat
    |
    | Marlin W4A16 Gate/Up GEMM, mulTopkWeights=false
    v
gateUpOutput [totalSlots, 2 * moeInterSize]
    |
    | SwiGLU: silu(gate) * up
    v
activationOutput [totalSlots, moeInterSize]
    |
    | Marlin W4A16 Down GEMM, mulTopkWeights=true
    v
downOutput [totalSlots, hiddenSize]
    |
    | aggregate slots by token
    v
output [numTokens, hiddenSize]

13. 关键结论

  • topkIndices 决定每个 slot 去哪个 expert。
  • sortedTokenIds 是按 expert 分组、padding 后的 slot id 列表。
  • expertIds 是每个 Marlin block 对应的 expert。
  • totalSlots 作为 sentinel，表示 padding row。
  • Gate/Up GEMM 不乘 routing weight，因为后面有 SwiGLU 非线性。
  • Down GEMM 融合乘 routing weight，最终 aggregate 只需要求和。
  • Marlin 内部按 expert-contiguous layout 计算，但输出保持原始 slot layout。

📚 相关文章推荐

评论