Blackwell 上的 DeepGEMM FusedMoE（vLLM）

Blackwell 上的 DeepGEMM FusedMoE

• 基准版本：vLLM v0.30.0（本地 tag 检出）
• 代码路径前缀：vllm/model_executor/layers/fused_moe/（下文链接均指向 v0.30.0 源码，链接文字即函数 / 类 / 变量名）

TL;DR

今天我们来扒一下当前最火热的大模型推理框架之一 vLLM 在 Blackwell（SM100/SM103）上的 FusedMoE 栈，它（MoERunner + FusedMoEKernelModularImpl 三段式）通过 moe_backend=deep_gemm（或 oracle/auto 落到 DeepGEMM）支持三种配方：

量化方法	Experts 类	grouped GEMM	量化组合	硬件门槛
FP8 block 128×128	DeepGemmExperts（Triton 兜底由 TritonOrDeepGemmExperts 组合）	m_grouped_fp8_gemm_nt_contiguous	W8A8（静态 block 权重 + 动态 128 激活）	SM90 起（SM100 上激活尺度走 E8M0 路径）
MXFP8（1×32）	DeepGemmExperts（block_shape == [1,32] 分支，见 __init__）	同上 + recipe (1,32)	权重 MXFP8（E8M0/32） + 动态 MXFP8 激活	SM100
W4A8 MXFP4	DeepGemmFP4Experts	m_grouped_fp8_fp4_gemm_nt_contiguous	MXFP4 权重（E2M1 + E8M0/32） × FP8 激活（128）	SM100/120

每个专家的计算是固定的六步流水线：

1. 对齐 M 维度（compute_aligned_M_and_alignment）→
2. 重排序 permute（deepgemm_moe_permute / ep_scatter：Triton _fwd_kernel_ep_scatter_1/_fwd_kernel_ep_scatter_2，产出 m_indices/inv_perm）→
3. 计算 GEMM1（m_grouped_fp8_gemm_nt_contiguous）→
4. 激活运算，SwiGLU + 动态量化（DeepGemmExperts._act_mul_quant：silu_mul_quant_fp8_packed_triton / silu_mul_per_token_group_quant_fp8_colmajor）→
5. 计算 GEMM2（同 m_grouped_fp8_gemm_nt_contiguous）→
6. ep_gather 加权归约（deepgemm_unpermute_and_reduce / ep_gather：Triton _fwd_kernel_ep_gather）

SM100 与 Hopper 的差别集中在尺度格式（UE8M0 int32 packed vs fp32）与 FP4/MXFP8 两条 Blackwell 专属配方上；TP 与 EP（无 DP / DP+SP）的分岔逻辑两代一致。

本文讲的代码路径：（以 vLLM 0.30.0 为例）

• 模型侧：DeepSeek V4/Kimi K3 等非 mega MoE 层，经由 FusedMoEFactory → MoERunner → FusedMoEKernelModularImpl.apply；
• Kernel 侧：DeepGemmExperts / DeepGemmFP4Experts / BatchedDeepGemmExperts，以及 TritonOrDeepGemmExperts 的 fallback 组合；
• 后端选择：moe_backend=deep_gemm / batched_deep_gemm / auto（oracle）。

---

1. 全景

1.1 六步速览

步	位置（v0.30.0）	关键函数 / 类 / Triton kernel	输入 → 输出	重点 tensor（形状 / dtype）	说明
0	models/deepseek_v4/nvidia/model.py:1260-1353	DeepseekV4DecoderLayer.forward；TileLang mhc_pre_tilelang / mhc_fused_post_pre_tilelang、sp_all_gather / sp_reduce_scatter	attention 输出 → hidden [M,H] bf16	hidden_states [M,H] bf16	mhc norm/混合 + SP 的 all_gather/reduce_scatter
1	runner/moe_runner.py:681-791 / 862-937	MoERunner.forward / _forward_impl；shared experts maybe_forward_async	hidden、router_logits → _forward_impl	hidden_states [M,H] bf16；router_logits [M,E_global] fp32	gate → shared expert（可多流 overlap）→ SP context → dispatch
2	runner/moe_runner.py:626-631、router/fused_topk_bias_router.py:345	router.select_experts（DSV4 为 FusedTopKBiasRouter）	router_logits → topk_weights [M,topk] fp32、topk_ids [M,topk]	topk_weights [M,topk] fp32；topk_ids [M,topk] int32（HT/LL 为 int64）	select_experts，dtype 由 prepare/finalize 约束
3	modular_kernel.py:1452-1547	FusedMoEKernelModularImpl.apply（_prepare / _fused_experts / _finalize）	(hidden, topk) → 三段式 apply	output [M,H] bf16（empty_like）	output = empty_like；_prepare → _fused_experts → _finalize
4	prepare_finalize/no_dp_ep.py:14-98（或 deepep_ht.py / deepep_ll.py）	MoEPrepareAndFinalizeNoDPEPModular.prepare → _quantize_input → moe_kernel_quantize_input	hidden bf16 → a1q fp8 [M,K] + scale	a1q [M,K] fp8_e4m3；scale [M,K/128] fp32（FP8 block）/ [M,K/32] uint8（MXFP8 UE8M0，permute 时打包）	本地 block 量化（NoDPEP）或 dispatch（HT/LL）
5	experts/deep_gemm_moe.py:287-394（FP8）/ :535-636（FP4）	DeepGemmExperts.apply / DeepGemmFP4Experts.apply；deepgemm_moe_permute / ep_scatter	a1q → permuted 缓冲 + m_indices + inv_perm	a1q_perm [M_sum,K] fp8_e4m3；expert_ids [M_sum] int32（m_indices）；inv_perm [M,topk] int32	deepgemm_moe_permute（ep_scatter）
6	同上	m_grouped_fp8_gemm_nt_contiguous（FP4 为 m_grouped_fp8_fp4_gemm_nt_contiguous）；mk_alignment_scope 包裹	[M_sum,K] → [M_sum,N]	mm1_out [M_sum,N] bf16（N = 2×intermediate_size_per_partition）	grouped GEMM1（mk_alignment_scope 包裹）
7	同上 + fp8_utils.py:250/414/683	_act_mul_quant → silu_mul_quant_fp8_packed_triton（本地别名 fused_silu_mul_fp8_quant_packed）/ silu_mul_per_token_group_quant_fp8_colmajor / per_token_group_quant_fp8_packed_for_deepgemm	[M_sum,N] → [M_sum,N/2] fp8 + packed scale	a2q [M_sum,N/2] fp8_e4m3；scale int32（UE8M0 packed）或 fp32（col-major）	SwiGLU(gate×up) + clamp + 动态量化
8	同上	m_grouped_fp8_gemm_nt_contiguous（GEMM2）；deepgemm_unpermute_and_reduce / ep_gather	[M_sum,K] → [M,H] bf16	mm2_out [M_sum,K] bf16 → output [M,K] bf16（fp32 累加后 cast）	grouped GEMM2 + ep_gather（fp32 加权累加）

1.2 调用图

DeepseekV4DecoderLayer.forward()                      models/deepseek_v4/nvidia/model.py:1260
 ├── self.attn(...)                                   （DP attention 时另有同步处理）
 └── self.ffn(...) = DeepseekV4MoE                    model.py:1019 / 类 :774
      ├── self.gate(hidden_states) → router_logits    （DSV4 用 fp32 logits）
      └── self.experts(...) = MoERunner               layer.py:408（FusedMoEFactory 构造）
           └── MoERunner.forward()                    runner/moe_runner.py:681
                ├── self._forward_entry(...)          :731  = torch.ops.vllm.moe_forward(_shared)
                │    └── _forward_impl(...)           :862
                │         ├─ shared experts maybe_forward_async   :892
                │         ├─ gate → router_logits                 :900
                │         └─ with _sequence_parallel_context():   :907
                │              ├─ _maybe_dispatch(...)            :911 / :799
                │              ├─ _apply_quant_method(...)        :916 / :595
                │              │    ├─ router.select_experts(...) :626
                │              │    └─ routed_experts.forward_modular
                │              │         └─ FusedMoEKernelModularImpl.apply   modular_kernel.py:1452
                │              │              ├─ _prepare(...)         :1211  量化 + dispatch
                │              │              ├─ _fused_experts(...)   :1300  → DeepGemmExperts.apply
                │              │              │    ├─ compute_aligned_M_and_alignment  deep_gemm_moe.py:320
                │              │              │    ├─ deepgemm_moe_permute              :331
                │              │              │    ├─ with mk_alignment_scope(...)      :356
                │              │              │    │    ├─ grouped GEMM1                :358
                │              │              │    │    ├─ _act_mul_quant               :370
                │              │              │    │    └─ grouped GEMM2                :375
                │              │              │    └─ deepgemm_unpermute_and_reduce      :386
                │              │              └─ _finalize(...)        :1384
                │              └─ _maybe_combine(...)             :924 / :826
                └── shared + routed 相加 → 可选 all-reduce  :782 / :480

---

2. 分步拆解

2.1 Step 0：attention → hidden states

DSV4 的 decoder 层（DeepseekV4DecoderLayer.forward）里，attention 前后有 mhc（post/pre mix + norm 融合 tilelang kernel，mhc_pre_tilelang / mhc_fused_post_pre_tilelang）与 SP 的通信（sp_all_gather / sp_reduce_scatter）；输出 hidden_states [num_tokens, hidden] bf16 交给 self.ffn（DeepseekV4MoE）。这一步与 MoE 后端无关，略过。

2.2 Step 1：MoERunner——路由之前

MoERunner.forward：

1. _forward_entry（custom op）→ _forward_impl；entry 由 _select_forward 在 __init__ 里选好，最后落到 torch.ops.vllm.moe_forward 或者 torch.ops.vllm.moe_forward_shared，取决于有没有 shared experts。

   self._forward_entry = self._select_forward()
   
   def _select_forward(self) -> Callable:
       if current_platform.is_tpu():
           ...
       if current_platform.is_cpu():
           ...
   
       return (
           torch.ops.vllm.moe_forward
           if self._shared_experts is None
           else torch.ops.vllm.moe_forward_shared
       )
   
   def forward(...):
       result = self._forward_entry(
           hidden_states,
           router_logits,
           shared_experts_input,
           input_ids,
           self._encode_layer_name(),
           self.moe_config.hidden_dim_unpadded
           if self._quant_method.has_unpadded_output
           else 0,
       )

   torch.ops.vllm.moe_forward_shared 会调用到 DeepseekV4MoE.experts 层实现中（详细过程见附录 B），而 experts 的类其实就是 MoERunner（经过 FusedMoEFactory 构造得到）。

2. shared experts 若支持多流，先 maybe_forward_async 异步发射，与后续 dispatch 重叠；

3. gate 输出 router_logits（FP8 模型里 DSV4 用 fp32 logits，见 self.gate）；

4. with self._sequence_parallel_context(): 内做 dispatch → quant method → combine（_sequence_parallel_context）。

_maybe_dispatch / _maybe_combine 只在 naive 模式（do_naive_dispatch_combine：dp>1 或 SP 且不用 all2all kernels）才走 get_ep_group().dispatch_router_logits/combine 的 AG+RS；用了 DeepEP HT/LL 时 dispatch/combine 在 prepare/finalize 内部完成。

2.3 Step 2：路由与 topk dtype

模块化路径下先显式选专家（_apply_quant_method）：

topk_weights, topk_ids = self.router.select_experts(
    hidden_states=hidden_states,
    router_logits=router_logits,
    topk_indices_dtype=self._quant_method.topk_indices_dtype,
    input_ids=input_ids,
)

topk_indices_dtype 由 prepare/finalize 反向决定（FusedMoEMethodBase.topk_indices_dtype → moe_kernel.prepare_finalize.topk_indices_dtype()）：

prepare/finalize	topk dtype	说明
NoDPEP（TP / 无 DP 的 EP）	None（保持 int32 等原 dtype）	纯本地，无约束（MoEPrepareAndFinalizeNoDPEPModular.topk_indices_dtype）
DeepEP HT	torch.int64	DeepEP dispatch 要求 int64（DeepEPHTPrepareAndFinalize.topk_indices_dtype）
DeepEP LL	torch.int64	同上（DeepEPLLPrepareAndFinalize.topk_indices_dtype）

DSV4 的路由器是 FusedTopKBiasRouter（sqrtsoftplus + e_score_correction_bias + hash 查表 + vision bias_vl），输出 topk_weights fp32、topk_ids（select_experts 里按 dtype 需求转换）。DeepGEMM 的 permute 内核要求 topk_ids 必须是有符号类型（topk_ids.dtype.is_signed），用 -1 表示无效槽位。

2.4 Step 3：FusedMoEKernel.apply 三段式

FusedMoEKernelModularImpl.apply：

output = torch.empty_like(hidden_states)                       # :1494
a1q, a1q_scale, expert_tokens_meta, topk_ids, topk_weights = self._prepare(...)   # :1500
fused_out = self._fused_experts(in_dtype=hidden_states.dtype, a1q=a1q, ...)        # :1516
return self._finalize(output, fused_out, hidden_states, topk_weights, topk_ids, ...)  # :1536

_prepare 调 self.prepare_finalize.prepare（defer_input_quant=self.fused_experts.expects_unquantized_inputs）。DeepGEMM 的 experts 不消费未量化输入，所以 expects_unquantized_inputs 为 False（默认），量化在 prepare 阶段完成。

2.5 Step 4：prepare——量化与 dispatch

以最常见的 NoDPEP（MoEPrepareAndFinalizeNoDPEPModular）为例：

def _quantize_input(a1, quant_config, defer_input_quant=False):     # :14-39
    input_sf = quant_config.a1_scale
    a1q, a1q_scale = moe_kernel_quantize_input(
        a1, input_sf,
        quant_dtype=quant_config.quant_dtype,
        per_act_token_quant=quant_config.per_act_token_quant,
        block_shape=quant_config.block_shape,
        is_scale_swizzled=quant_config.is_scale_swizzled,
        mx_alignment=quant_config.mx_alignment,
    )

moe_kernel_quantize_input 按 quant_dtype 分派：

quant_dtype	实现	输出
fp8（DSV4 fp8 权重路径）	_fp8_quantize → block 时 per_token_group_quant_fp8	fp8 激活 + fp32 分组尺度（128 组）
mxfp8（Kimi K3 等）	_mxfp8_e4m3_quantize	fp8 激活 + uint8 UE8M0 分组尺度（32 组，permute 时打包）
nvfp4 / int8 / …	对应实现（_nvfp4_quantize、_int8_quantize）	—

EP 场景：

• DeepEP HT（prepare_async）：quant_config.is_block_quantized and not defer_input_quant 时先量化再 dispatch（quant_config.is_block_quantized 分支），否则 dispatch bf16、落地后量化（_receiver）。dispatch 返回的 expert_topk_ids 带 rank 偏移，接收侧把它还原回全局 id 以便 expert_map 映射，并构造 ExpertTokensMetadata。这种“每专家 token 数已知”的元数据是 kernel 侧省掉一次计数 kernel 的关键。
• DeepEP LL（prepare_async）：low_latency_dispatch 按 SUPPORTED_HIDDEN_SIZES 走 batched 格式，激活为 [E_local, max_tokens, K]（activation_format 为 BatchedExperts），专家侧换成 BatchedDeepGemmExperts（masked GEMM）。

2.6 Step 5：experts.apply 五步流水线（以 FP8 DeepGemmExperts 为准）

DeepGemmExperts.apply，对照 FP4 版 DeepGemmFP4Experts.apply。

2.6.1 对齐 M_sum

M_sum, _ = compute_aligned_M_and_alignment(
    M=topk_ids.size(0), num_topk=topk_ids.size(1),
    local_num_experts=local_num_experts,
    alignment=get_mk_alignment_for_contiguous_layout()[0],
    expert_tokens_meta=expert_tokens_meta)

compute_aligned_M_and_alignment 两条路径：

• 有 CPU 侧每专家 token 数（HT/LL dispatch 提供）→ 每专家 round_up(align) 后求和，精确（expert_num_tokens_cpu）；
• 否则（TP/无 meta）→ 上界估算 M*topk + min(M*topk, E_local)*(align-1) 再 round_up；并且在 SM100/SM120 上会用 get_theoretical_mk_alignment_for_contiguous_layout 收缩实际对齐粒度（per_call_align），返回 (M_sum, align_used)。

2.6.2 permute：ep_scatter → 连续布局 + m_indices + inv_perm

a1q_perm = _resize_cache(workspace13.view(dtype=torch.float8_e4m3fn), (M_sum, K))
a1q, a1q_scale, expert_ids, inv_perm, align_used = deepgemm_moe_permute(
    aq=a1q, aq_scale=a1q_scale, topk_ids=topk_ids,
    local_num_experts=local_num_experts, expert_map=expert_map,
    expert_tokens_meta=expert_tokens_meta, aq_out=a1q_perm,
    block_size=self.block_shape[1] if self.mxfp8 else None,   # MXFP8: 32；FP8: 默认 128
)

deepgemm_moe_permute 做四件事：

1. 分配三类输出：aq_out [M_sum,H]（激活）、aq_scale_out（尺度）、expert_ids [M_sum] int32（即 DeepGEMM 的 m_indices）、inv_perm [M,topk] int32；
2. 尺度格式分叉（pack_ue8m0）：pack_ue8m0 = aq_scale.dtype == torch.uint8（MXFP8 的 UE8M0 尺度）→ empty_strided((M_sum, ceil(sf_k/4)), (1, round_up(M_sum,4)), int32)，即 4 个 8-bit 指数打包进 int32、并按 TMA 对齐；否则（fp32 尺度）→ zeros((M_sum, sf_k), fp32)。zeros（而非 empty）保证 padding 行的尺度是 0；
3. m_indices 初始化为 -1（expert_ids）：未被 scatter 写入的行 = padding，DeepGEMM 调度器会跳过负值块；
4. 执行 ep_scatter。

ep_scatter 是两个 Triton kernel：

• _fwd_kernel_ep_scatter_1（grid=num_experts）：对每专家 token 数按 ALIGN_M 取整后前缀和 → expert_start_loc[e] 给出每个专家段在 [M_sum] 中的起点；同时把该段前 n_e 个 m_indices 填成专家 id（padding 槽保持 -1）；
• _fwd_kernel_ep_scatter_2（grid=min(T,8192)）：逐 token 加载 x/scale，atomic_add 得到目标行，写 aq_out[dest]、inv_perm[token,k] = dest；expert_id < 0（无效）直接跳过（expert_id >= 0）；PACK_UE8M0 时把 4 个 uint8 尺度合并成一个 int32（PACK_UE8M0）。

permute 之后：

• a1q [M_sum, K]：每个专家一段、段内按 token 顺序，每段 M 按 align_used 对齐；
• expert_ids [M_sum] int32：行 → 专家 id（padding = -1）；
• inv_perm [M, topk]：反查表，finalize 时用。

2.6.3 grouped GEMM1

with mk_alignment_scope(align_used):                      # :356
    mm1_out = _resize_cache(workspace2, (M_sum, N))
    m_grouped_fp8_gemm_nt_contiguous(
        (a1q, a1q_scale), (w1, self.w1_scale), mm1_out, expert_ids, **gemm_kwargs)

• mk_alignment_scope 临时把 DeepGEMM 的 contiguous-layout BLOCK_M 上限设为 align_used。必要性见 2.6.1 / 2.8：DeepGEMM 调度器按 m_block_idx * BLOCK_M 去读 m_indices 取专家 id，若实际对齐粒度比 BLOCK_M 小，就会读到别的专家的 id → 计算错乱甚至 IMA（set_mk_alignment_for_contiguous_layout 的 docstring 写得很明确）；
• gemm_kwargs 对 MXFP8 传 recipe_a/b=(1,32)（gemm_kwargs）；
• wrapper（m_grouped_fp8_gemm_nt_contiguous）统一带 disable_ue8m0_cast=not is_deep_gemm_e8m0_used()。

2.6.4 SwiGLU + clamp + 动态量化

_act_mul_quant 按尺度格式三分支：

条件	kernel	产物
scale_fmt == UE8M0（SM100/120 默认）	silu_mul_quant_fp8_packed_triton（模块内别名 fused_silu_mul_fp8_quant_packed）	gate×silu(up) + clamp + 量化 + packed int32 UE8M0（一条龙）
其他且是门控激活	silu_mul_per_token_group_quant_fp8_colmajor（use_ue8m0 = (fmt == FLOAT32_CEIL_UE8M0)）	同上但尺度为 fp32/列主序
非门控兜底	先 self.activation，再 per_token_group_quant_fp8	fp32 尺度

配方标量（gemm1_alpha / gemm1_beta）让同一个 kernel 覆盖 silu（α=1, β=0）与 swigluoai；gemm1_clamp_limit 来自模型 swiglu_limit。

2.6.5 grouped GEMM2 与 ep_gather

mm2_out = _resize_cache(workspace2, (M_sum, K))
m_grouped_fp8_gemm_nt_contiguous((a2q, a2q_scale), (w2, self.w2_scale), mm2_out, expert_ids, ...)
if apply_router_weight_on_input: topk_weights = torch.ones_like(topk_weights)
deepgemm_unpermute_and_reduce(a=mm2_out, topk_ids=topk_ids, topk_weights=topk_weights,
                              inv_perm=inv_perm, expert_map=expert_map, output=output)

ep_gather（kernel _fwd_kernel_ep_gather）：grid (cdiv(H,BLOCK_D), min(T,1024))，BLOCK_D = gcd(H,1024)；对每个 token 的 topk 槽位，按 inv_perm 取回专家输出行，fp32 累加 weight * value，最后 cast 回 bf16 写 output [M,H]。注意 topk_ids == -1（EP 时该 token 不在本 rank）直接跳过。

2.7 Step 6：finalize

场景	finalize 行为	关键代码
NoDPEP（TP / 无 DP 的 EP）	TopKWeightAndReduceDelegate → 换 TopKWeightAndReduceContiguous，但 DeepGemm 返回的是 TopKWeightAndReduceNoOP，output.copy_ 直通（ep_gather 已完成加权归约）	MoEPrepareAndFinalizeNoDPEPModular.finalize；finalize_weight_and_reduce_impl
DeepEP HT	对 NoOP：combine 时 topk_weights=None（加权已在 ep_gather 做），combine 后 output.copy_；output_is_reduced=True	_finalize
DeepEP LL	要求 TopKWeightAndReduceDelegate（LL 的 combine 内部做加权），low_latency_combine	_finalize

之后回到 MoERunner.forward 尾部：shared_output + fused_output 把 shared expert 输出相加，再按条件做最终 all-reduce（_maybe_reduce_final_output）：not SP and not skip_final_all_reduce and (tp>1 or ep>1) and not output_is_reduced。

2.8 Step 7：工程细节

（a）workspace 复用。 workspace13（激活/量化输出）与 workspace2（GEMM 输出）按最大需求分配、跨层共享；kernel 内用 _resize_cache 把同一块存储在不同阶段视图化成 (M_sum,K)、(M_sum,N) 等形状。workspace_shapes 里 workspace1=(M_sum, max(act_out,K))、workspace2=(M_sum, max(N,K))。cudagraph 捕获时形状固定，因此 M_sum 必须覆盖最坏情况。

（b）TritonOrDeepGemmExperts 的动态选择。 _select_experts_impl：

if is_deep_gemm_e8m0_used() or _valid_deep_gemm(hidden_states, w1, w2):
    return self.experts           # DeepGemm
return self.fallback_experts      # Triton

_valid_deep_gemm 的条件：deep_gemm 可用、M ≥ align 且 N/K 是 align 倍数、N > 512、权重是 fp8、三张量连续。workspace 则悲观按 DeepGemm 尺寸分配（workspace_shapes），因为 DeepGemm 的 workspace 严格更大。

（c）权重预处理链（FP8）。 加载后 Fp8MoEMethod._setup_kernel → convert_to_fp8_moe_kernel_format → prepare_fp8_moe_layer_for_deepgemm：

1. 若 checkpoint 尺度非 E8M0 且当前启用 E8M0 → requant_weight_ue8m0_inplace（把权重 scale 重新量化为 UE8M0）；
2. deepgemm_post_process_fp8_weight_block → _upcast_e8m0_to_fp32 → deepgemm_post_process_weight_scale_block → transform_sf_into_required_layout(sf, mn, k, recipe=(1, block_n, block_k), num_groups, is_sfa=False)，得到内核要求的 MN-major、TMA 对齐 packed 布局。

（d）FP4 权重尺度。 _pack_deepgemm_mxfp4_scales：对 w13/w2 各调一次 deepgemm_post_process_weight_scale_block(..., quant_block_shape=(1,32))。权重本身保持 (E, N, K//2) 的 packed uint8/FP4，在 apply 里 w.view(torch.int8) 传给 GEMM。

（e）oracle 选择。

• FP8（_get_priority_backends）：默认优先级 AITER → FLASHINFER_TRTLLM → FLASHINFER_CUTLASS → DEEPGEMM → VLLM_CUTLASS → TRITON → …；SM100 + DeepEP v2 contiguous 布局时 TRTLLM 提权到最前（_move_to_front）。
• MXFP4（_get_priority_backends）：CUDA 上 FLASHINFER_TRTLLM_MXFP8 → DEEPGEMM_MXFP4（激活 key 由 _backend_activation_key 给出，= kFp8Dynamic128Sym）→ MARLIN → BATCHED_MARLIN；选中后 hidden/intermediate 由 mxfp4_round_up_hidden_size_and_intermediate_size round up 到 128（Mxfp4MoeBackend.DEEPGEMM_MXFP4）。
• 显式指定：--moe-backend deep_gemm（FP8→DEEPGEMM；FP4→DEEPGEMM_MXFP4），--moe-backend batched_deep_gemm（LL 的 batched 形式）。

（f）E8M0 开关与白名单。 DeepGemmQuantScaleFMT：VLLM_USE_DEEP_GEMM_E8M0=1（默认）且 DeepGEMM 支持时，SM100/120 → UE8M0，Hopper → FLOAT32_CEIL_UE8M0；否则 FLOAT32。已知 Blackwell 上 E8M0 掉精度的模型白名单（Qwen3.5 系列）会被 should_auto_disable_deep_gemm 在 oracle 层自动关掉 DeepGEMM。

---

3. TP 与 EP 四种形态

并行决策（FusedMoEParallelConfig.make）：use_ep = dp*pcp*tp > 1 and enable_expert_parallel；EP 时 tp_size=1、ep_size=flattened，非 EP 时 ep_size=1、expert_map=None。

形态	prepare/finalize	权重切分	尾部归约
TP（ep=1）	NoDPEP：本地量化，num_dispatchers=1	每个专家中间维 I//tp（intermediate_size_per_partition）	tp>1 → all-reduce（_maybe_reduce_final_output）
EP，无 DP/SP	NoDPEP + expert_map（全局 id → 本地 id）	专家维切分：n_local = n_physical / ep_size	ep>1 → all-reduce
EP + DP/SP（all2all）	DeepEP HT：先量化后 dispatch、ExpertTokensMetadata、combine 不加权	同上（可叠加 EPLB 冗余副本）	output_is_reduced=True，跳过
EP + DeepEP LL	batched 格式 + BatchedDeepGemmExperts + LL combine（内部加权）	同上	同上

对 kernel 而言，TP 与 EP 的差别只有两点：expert_map 是否存在（影响 ep_scatter 里 id→本地 id 的映射）与每专家 token 数来源（CPU 列表 vs 估算）。这也是 DeepGEMM 的 contiguous-layout 设计能同时服务两种并行的原因。

---

4. 与 Hopper 的差异（一页对照）

维度	SM90（Hopper）	SM100+（Blackwell）
FP8 block 配方	✅ DeepGemmExperts	✅ 同左
激活/权重尺度格式	FLOAT32_CEIL_UE8M0（fp32 存储）	UE8M0（4×uint8 → int32 packed，TMA 对齐 stride）
SwiGLU+量化 kernel	silu_mul_per_token_group_quant_fp8_colmajor	silu_mul_quant_fp8_packed_triton
MXFP8（1×32）	❌	✅ 复用 fp8_fp4 GEMM 的 (1,32) recipe（family 100）
W4A8 MXFP4	❌（family 限制）	✅ DeepGemmFP4Experts（family 100/120）
对齐 M_sum 估算	固定 128 对齐	可用 get_theoretical_mk_alignment_for_contiguous_layout 收缩每调用 BLOCK_M

整体流水线（permute → GBMM → act/quant → GBMM → ep_gather）与 TP/EP 的 prepare/finalize 拓扑两代完全一致，差异都收敛在"尺度格式"与"配方可用性"两处。

---

5. 小结

1. SM100+ 上 DeepGEMM 在通用 FusedMoE 栈里有三条配方：FP8 block（主力）、MXFP8 1×32（Blackwell 专属）、W4A8 MXFP4（Blackwell 专属）；它们共享同一套 permute/GEMM/gather 骨架。
2. 一个专家的全部计算可概括为：ep_scatter 重排 → grouped GEMM1 → 门控激活+量化 → grouped GEMM2 → ep_gather 加权归约；通信只在 EP 时由 prepare/finalize（HT/LL）介入。
3. SM100 的工程关键点是 E8M0 packed 尺度（量化 kernel、scale buffer stride、DG wrapper 的 disable_ue8m0_cast）与 mk_alignment_scope（把 DG 调度器的 BLOCK_M 上限锁到本次调用的实际对齐，避免 cudagraph 下 m_indices 调度错位）。
4. 选型：DSV4-Flash 的 fp4 checkpoint 在非 mega 路径下落 DeepGemmFP4Experts；fp8 checkpoint 与 Kimi K3 的 MXFP8 落 DeepGemmExperts；LL 场景落 BatchedDeepGemmExperts。

下一篇：MegaMoE——deep_gemm_mega_moe 与 FlashInfer moe_ep mega 的对称内存单 kernel 实现。

---

附录 A：符号速查

主题	符号（v0.30.0）
MoERunner（gate / quant method / reduce）	MoERunner（类）、forward、_forward_impl、_apply_quant_method、_sequence_parallel_context、_maybe_reduce_final_output
FusedMoEKernel 三段式	FusedMoEKernelModularImpl.apply、_prepare、_fused_experts、_finalize、expects_unquantized_inputs
NoDPEP prepare/finalize	MoEPrepareAndFinalizeNoDPEPModular、_quantize_input
DeepEP HT / LL	DeepEPHTPrepareAndFinalize、DeepEPLLPrepareAndFinalize
DeepGemmExperts / DeepGemmFP4Experts	DeepGemmExperts、DeepGemmFP4Experts
Triton 兜底组合	TritonOrDeepGemmExperts、FallbackExperts
permute / ep_scatter / ep_gather	compute_aligned_M_and_alignment、deepgemm_moe_permute、ep_scatter、ep_gather、deepgemm_unpermute_and_reduce
TopK 加权归约三件套	TopKWeightAndReduceDelegate、TopKWeightAndReduceNoOP、TopKWeightAndReduceContiguous
输入量化分派	moe_kernel_quantize_input、_fp8_quantize、_mxfp8_e4m3_quantize、_nvfp4_quantize、_int8_quantize
scale FMT / alignment / GEMM 包装	DeepGemmQuantScaleFMT、is_deep_gemm_e8m0_used、get_theoretical_mk_alignment_for_contiguous_layout、mk_alignment_scope、m_grouped_fp8_gemm_nt_contiguous、m_grouped_fp8_fp4_gemm_nt_contiguous
激活量化 kernel	silu_mul_quant_fp8_packed_triton、silu_mul_per_token_group_quant_fp8_colmajor、per_token_group_quant_fp8、per_token_group_quant_fp8_packed_for_deepgemm
权重尺度预处理	requant_weight_ue8m0_inplace、_upcast_e8m0_to_fp32、deepgemm_post_process_weight_scale_block、deepgemm_post_process_fp8_weight_block、prepare_fp8_moe_layer_for_deepgemm（均在 fp8_utils.py）
FP8 / FP4 oracle	_get_priority_backends、convert_to_fp8_moe_kernel_format、_pack_deepgemm_mxfp4_scales、_get_priority_backends、select_mxfp4_moe_backend、mxfp4_round_up_hidden_size_and_intermediate_size

---

附录 B：direct_register_custom_op 速查

vLLM 内注册 custom op 的统一入口（direct_register_custom_op），vLLM 几乎所有 custom op（moe_forward*、sequence_parallel_chunk_impl 等）都走它。它的定位很清楚：绕过 torch.library.custom_op 的复杂 dispatch 逻辑（后者有明显 overhead），直接把算子 define + impl 到指定后端。

用法示例

from vllm.utils.torch_utils import direct_register_custom_op

direct_register_custom_op(
    op_name="moe_forward_shared",        # 注册为 torch.ops.vllm.moe_forward_shared
    op_func=_moe_forward_shared,         # 实现：查表转发到 layer._forward_impl
    fake_impl=_moe_forward_shared_fake,  # 编译期 meta 形状推断
    tags=(torch.Tag.needs_fixed_stride_order,),
)

# 之后即可调用：
torch.ops.vllm.moe_forward_shared(hidden_states, router_logits, ...)

参数意义

参数	默认值	意义
op_name	—	算子名，注册后完整路径为 torch.ops.<lib>.<op_name>（vLLM 默认 torch.ops.vllm.*）。
op_func	—	实现函数。它的类型标注就是 op 的 schema：infer_schema() 从函数签名反推参数/返回类型，标注不全或返回 union 类型都无法注册（_moe_forward / _moe_forward_shared 拆成两个 op 正是这个限制逼出来的）。
mutates_args	None → 内部补 []	声明哪些输入参数会被原地修改（按参数名，如 ["hidden_states"]）。它决定 functionalization 如何对待这些 buffer：声明了才允许原地写；不声明却实际原地改，会在 torch.compile 下引发正确性问题。
fake_impl	None	meta 设备上的 fake 实现，负责编译期形状/类型推断。custom op 对 Dynamo 是黑盒，输出形状的唯一来源就是它；不传则没有 fake kernel。
target_lib	vllm_lib	注册到哪个 Library，默认 Library("vllm", "FRAGMENT")（FRAGMENT 允许跨文件逐片追加 schema，全 vLLM 共用一个库）。op 的生命周期与 Library 对象绑定。
dispatch_key	current_platform.dispatch_key	实现绑定的后端 key：CUDA/ROCm 为 "CUDA"、XPU 为 "XPU"、CPU 为 "CPU"（见 vllm/platforms/*.py）。运行时按设备类型走对应实现。
tags	()	torch.library 的 tag 元组。vLLM 最常用 torch.Tag.needs_fixed_stride_order：要求输出保持固定 stride 顺序、禁止 Inductor 做布局改写，保证 compile / CUDA Graph 下与 eager 的内存布局一致。

关键理解点

1. 实现只有三步：infer_schema(op_func, mutates_args) 生成 schema 字符串 → my_lib.define(op_name + schema_str, tags=tags) 注册 → my_lib.impl(op_name, op_func, dispatch_key=...) 绑定实现 → 若有 fake 则 my_lib._register_fake(...)。等价于官方 custom_op 的最简子集。
2. schema 完全来自签名标注：参数要标注成 torch.Tensor / int / str 等可映射类型；返回类型不支持 union——这就是同一个 MoE 前向要拆成 moe_forward 与 moe_forward_shared 两个 op 的根因。
3. FRAGMENT 库语义：Library("vllm", "FRAGMENT") 创建后，各文件可以各自往里追加算子（区别于 DEF 模式的一次性定义）；同名算子重复注册会直接报错，因此 op_name 是全局唯一约定。
4. fake impl 的设计约束：为保持 fake 是"输入的纯函数"、让编译期子图去重（dedup）生效，vLLM 常把影响输出形状的运行时信息显式作为 op 参数传入（如 hidden_dim_unpadded），而不是在 fake 里回查 layer 状态。
5. 与 CUDA Graph 的配合：op 边界确立后，整个 MoE 在 cudagraph capture 时就是一次调用；needs_fixed_stride_order tag 保证输出 tensor 布局在重放中稳定（cudagraph 复用内存，输出 stride 漂移会破坏下游 kernel 的地址假设）。
6. 实例对照：_moe_forward（带 mutates_args=["hidden_states"]）与 _moe_forward_shared（不声明）的注册见 direct_register_custom_op——前者向编译器声明会原地写输入，后者按纯函数对待。

用例说明

direct_register_custom_op(
    op_name="moe_forward",
    op_func=_moe_forward,
    mutates_args=["hidden_states"],
    fake_impl=_moe_forward_fake,
    tags=(torch.Tag.needs_fixed_stride_order,),
)


direct_register_custom_op(
    op_name="moe_forward_shared",
    op_func=_moe_forward_shared,
    fake_impl=_moe_forward_shared_fake,
    tags=(torch.Tag.needs_fixed_stride_order,),
)

关键理解点：

1. _moe_forward_shared 本体只是个"查表 + 转发"：它通过 layer_name 找回 layer 再调 _forward_impl。也就是说从编译器的视角，路由、dispatch、量化、grouped GEMM 全部被藏在这一个 op 里面，torch Dynamo 不会去 trace 内部。
2. 为什么传 layer_name 而不是 layer 本身：nn.Module 引用不能进编译图。__init__ 里通过 register_layer_for_moe_forward_op 把 layer 塞进 static_forward_context，运行时 get_layer_from_name 再取回来；一个值得品味的细节是，_encode_layer_name 函数在非 LayerName 模式下会统一传常量 "from_forward_context"，而在 get_layer_from_name 遇到这个常量时，会自动按 moe_layer_index 自增顺序取层——这样所有 MoE 层共享同一个编译图，避免每层因常量不同而图分化。
   1. layer 是怎么注册进去的？

   DeepseekV4DecoderLayer.__init__                   # model.py:1132
   └─ DeepseekV4MoE(prefix=f"{prefix}.ffn")          → prefix = "model.layers.{i}.ffn"
       └─ FusedMoEFactory(prefix=f"{prefix}.experts")   # model.py:996-1017（非 mega 分支）
               ├─ layer_name = prefix                        # layer.py:217 → "model.layers.{i}.ffn.experts"
               ├─ routed_experts = RoutedExperts(layer_name, ...)   # layer.py:374
               └─ runner = MoERunner(layer_name=..., ...)           # layer.py:405-428
                   │   ← DSV4 没传 runner_cls/routed_experts_cls，所以就是默认的 MoERunner + RoutedExperts
                   └─ register_layer_for_moe_forward_op(vllm_config, self)   # moe_runner.py:299
                       ├─ compilation_config.static_forward_context[layer_name] = self   # :58-69
                       └─ static_all_moe_layers.append(layer_name)

   get_layer_from_name 查表时，no_compile_layers["model.layers.{i}.ffn.experts"] 拿到的就是这一层的 MoERunner 对象本身。
3. fake impl：编译期在 meta 上给出两个输出的形状。hidden_dim_unpadded > 0 分支对应 TRT-LLM MXFP4 路径（真实 kernel 写出的宽度比 hidden_states.shape[-1] 窄）；shared_out 在 latent MoE 时取 shared_experts_input 的形状。
4. torch.Tag.needs_fixed_stride_order：要求该 op 的输出保持固定 stride 顺序，不允许 Inductor 做 flexible-layout 之类的布局改写，保证 CUDA Graph 下输出内存布局与 eager 一致。这是 vLLM custom op 的模式化做法（sequence_parallel_chunk_impl 等也带此 tag）。
5. 上下对比 moe_forward 和 moe_forward_shared 注册可以发现：mutates_args=["hidden_states"]，会声明该值存在原地写输入的情况。