M1 讲解：MoE 路由与 moe_align_block_size

① 先看全貌：为什么要排序和补齐

MoE 把 MLP 换成 E 个专家，每个 token 只送到其中 top-k 个。直接写成循环：

for e in range(E):
    idx = tokens_routed_to(e)
    y[idx] += w[idx, e] * (x[idx] @ W_e)       # E 次小 GEMM，E 次 kernel launch

E = 64–256 时要启动几百个小 kernel，每个都喂不饱 GPU。高效的做法是一次 launch 算完所有专家：把所有 (token, slot) 对按专家排成一列，每个专家的那一段补齐到 BLOCK_M 的整数倍。这样每 BLOCK_M 行组成的块只属于一个专家，GEMM kernel 的每个 program 只要知道“我这块是哪个专家”，就能去读对应的权重。

② route：softmax → top-k → 重新归一化

drillsroute（18–22 行）

probs = torch.softmax(router_logits.float(), dim=-1)    # 每个 token 在 E 个专家上的概率
w, ids = torch.topk(probs, top_k, dim=-1)               # 取最大的 k 个
w = w / w.sum(-1, keepdim=True)                         # 让这 k 个权重加起来等于 1

一个 token、4 个专家，logits = [2.0, 1.0, 0.5, −1.0]，top-2：

topk_ids = [0, 1]，topk_weights = [0.731, 0.269]。最后的输出是 0.731 · expert_0(x) + 0.269 · expert_1(x)。

docstring 提到不同模型的路由方式不同：Mixtral、Qwen-MoE 是这里的“softmax + top-k + 重新归一化”；DeepSeek-V3 用 sigmoid 分数加分组限制（先选专家组，再在组内选专家）。M2 的 GEMM 部分对这些路由方式都一样，只要拿到 topk_ids 和 topk_weights。

③ moe_align_block_size：一个完整的例子

drillsmoe_align_block_size（25–44 行）

T = 4 个 token，K = 2，E = 4 个专家，block_m = 2：

topk_ids = [[0, 2],      # token 0 去专家 0 和 2
            [2, 1],      # token 1
            [0, 2],      # token 2
            [3, 2]]      # token 3

第 1 步：展平，给每个 (token, slot) 一个编号

flat = topk_ids.reshape(-1)          # [0, 2, 2, 1, 0, 2, 3, 2]

编号 f = t × K + k：f = 5 是 token 2 的第 1 个 slot（去专家 2）。后面所有的“token id”其实都是这个 f。

第 2 步：每个专家几个，补齐后多长，从哪开始

counts    = torch.bincount(flat, minlength=E)            # [2, 1, 4, 1]
padded    = (counts + block_m - 1) // block_m * block_m  # [2, 2, 4, 2]
pad_start = torch.cumsum(padded, 0) - padded             # [0, 2, 4, 8]
total     = padded.sum()                                 # 10

bincount 数每个值出现几次。cumsum − 自己 是排他前缀和：每个专家的段从哪里开始。和 F1 里从页数算 indptr 是同一个技巧。

第 3 步：稳定排序，算组内排名

order = torch.argsort(flat, stable=True)                 # [0, 4, 3, 1, 2, 5, 7, 6]
sorted_experts = flat[order]                             # [0, 0, 1, 2, 2, 2, 2, 3]
group_start = torch.cumsum(counts, 0) - counts           # [0, 2, 3, 7]（不补齐时的起点）
rank = arange(n) - group_start[sorted_experts]           # [0, 1, 0, 0, 1, 2, 3, 0]
dest = pad_start[sorted_experts] + rank                  # [0, 1, 2, 4, 5, 6, 7, 8]

• 为什么要稳定排序：同一个专家内部保持原来的 token 顺序（f = 1, 2, 5, 7，而不是任意顺序）。结果确定、可复现，和 vLLM 的输出也能逐元素对上。
• rank：排好之后，每个元素在自己专家的组里排第几。
• dest：在补齐后的布局里的位置 = 这个专家段的起点 + 组内排名。

第 4 步：写出 sorted_token_ids 和 expert_ids

sorted_token_ids = torch.full((total,), n)               # 先全部填 padding 值 n = T*K = 8
sorted_token_ids[dest] = order                           # 再把真实的 f 放到 dest 位置
expert_ids = repeat_interleave(arange(E), padded // block_m)

token id 是 f = t·K + k，所以 token = f // K。padding 的值是 T·K = 8，一个越界的编号，GEMM 读到它就知道这一行是空的（mask 掉）。专家 1 和专家 3 各只有 1 个 token，各浪费了 1 行。

padding 的代价

token 越多、block_m 越小，浪费越少；但 block_m 太小 GEMM 的效率又低。decode 时 T 很小，每个专家可能只分到一两个 token，padding 浪费会非常大。这正是 M2 里 decode 场景特别慢的原因。

④ 对照 vLLM

vLLM 的 Python 接口：vllmmoe_align_block_size（注意参数顺序是 topk_ids, block_size, num_experts，和这题不同）。它的 docstring 里还讲了 expert parallel 的处理：不在本卡上的专家，expert_ids 标成 −1，GEMM 直接跳过这些块。

CUDA 实现在 vllmmoe_align_block_size_kernel 和 vllmcount_and_sort_expert_tokens_kernel：

vLLM 的 CUDA 版组内顺序不确定，但 GEMM 的结果不受影响：每一行算的是同一个 (token, slot) 和同一个专家的乘积，只是写回的顺序不同。测试是按“每个专家的 token 集合”而不是按位置比较的话就不受影响。

⑤ 常见坑与自测

• 把 sorted_token_ids 当成 token 编号：它是 f = t·K + k，要 // K 才是 token。M2 里读 x 的时候用 f // K，读路由权重的时候用 f 本身。
• padding 值用 −1 或 0：0 是一个合法的 f，会被当成真实数据算进去；参考解和 vLLM 都用 T·K。
• 用 Python 循环按专家拼：E = 256 时每一步都多几百次 kernel launch，比 GEMM 本身还慢。
• counts 漏了没有 token 的专家：bincount 要带 minlength=E，否则最后几个专家的计数缺失，后面的 cumsum 全部错位。

源码与资料索引

本仓库：drillssolutions/m1_moe_align.py · drillsmine/m1_moe_align.py · drillstests/test_m1_moe_align.py · drillscontent/m1.md

vLLM：vllmmoe_align_block_size.py · vllmmoe_align_sum_kernels.cu · vllmfused_moe.py（M2）

论文：Mixtral of Experts · DeepSeek-V3 技术报告（sigmoid 路由与分组限制）